path: root/src/main.c

#ifdef __linux__
#define _POSIX_C_SOURCE 2
#define _DEFAULT_SOURCE
#endif

#include <unistd.h>
#include <stdlib.h>
#include <stdint.h>
#include <stdio.h>
#include <syslog.h>
#include <string.h>
#include <sys/socket.h>
#include <arpa/inet.h>
#include <netinet/in.h>
#include <netdb.h>
#include <sys/types.h>
#include <sys/time.h>
#include <pwd.h>
#include <grp.h>
#include <errno.h>
#include <signal.h>
//TODO: use threads.h instead
#include <pthread.h>

// illumos-only headers
#ifdef __illumos__
#include <priv.h>
#endif

#include "packet.h"
#include "netascii.h"
#include "handlers.h"

int nbtpd_stop = 0;

void stop_handler(int i) {
	syslog(LOG_ERR, "caught shutdown signal, cleaning up and awaiting threads");
	if (i != 0) {
		nbtpd_stop = i;
	} else {
		nbtpd_stop = 1;
	}
	return;
}

void usage(char *name) {
	//TODO: print to stderr; cleanup to be compliant
	fprintf(stderr,
		"usage: %s -h [-d [-u USER] [-g GROUP]] [-a INET_ADDRESS] [-p PORT] [-l LEVEL]\n",
		name
	);
	fprintf(stderr, "\td: daemonize.\n");
	fprintf(stderr,
		"\tu: username to run as (default: nobody).  Must be specified after -d.\n"
	);
	fprintf(stderr,
		"\tg: group to run as (default: nobody).  Must be specified after -d.\n"
	);
	fprintf(stderr, "\ta: address to bind to (default: 127.0.0.1)\n");
	fprintf(stderr, "\tp: port to bind to (default: 69)\n");
	fprintf(stderr, "\tl: log level; one of:\n");
	fprintf(stderr, "\t\temerg\n\t\talert\n\t\tcrit\n\t\terr\n\t\twarn\n");
	fprintf(stderr, "\t\tnotice\n\t\tinfo (default)\n\t\tdebug\n");
}

int set_loglevel(char *level) {
	if (strncmp("emerg", level, 6) == 0) {
		setlogmask(LOG_UPTO(LOG_EMERG));
		return 0;
	}
	if (strncmp("alert", level, 6) == 0) {
		setlogmask(LOG_UPTO(LOG_ALERT));
		return 0;
	}
	if (strncmp("crit", level, 5) == 0) {
		setlogmask(LOG_UPTO(LOG_CRIT));
		return 0;
	}
	if (strncmp("err", level, 4) == 0) {
		setlogmask(LOG_UPTO(LOG_ERR));
		return 0;
	}
	if (strncmp("warn", level, 5) == 0) {
		setlogmask(LOG_UPTO(LOG_WARNING));
		return 0;
	}
	if (strncmp("notice", level, 7) == 0) {
		setlogmask(LOG_UPTO(LOG_NOTICE));
		return 0;
	}
	if (strncmp("info", level, 5) == 0) {
		setlogmask(LOG_UPTO(LOG_INFO));
		return 0;
	}
	if (strncmp("debug", level, 6) == 0) {
		setlogmask(LOG_UPTO(LOG_DEBUG));
		return 0;
	}
	return -1;
}

int main(int argc, char **argv) {
	int retme = 0;
	int daemonize = 0, loglevset = 0;
	char addr[17], user[32], group[32];
	memset(addr, '\0', sizeof(addr));
	memset(user, '\0', sizeof(user));
	memset(group, '\0', sizeof(group));
	strcpy(addr, "127.0.0.1");
	strcpy(user, "nobody");
	strcpy(group, "nobody");
	int port = 69;
	int ch = 0;
	while ((ch = getopt(argc, argv, "da:p:u:g:l:h")) != -1) {
		switch (ch) {
			case 'a':
				strncpy(addr, optarg, sizeof(addr) - 1);
				break;
			case 'p':
				port = atoi(optarg);
				if ((port <= 0) || (port >= 65536)) {
					fprintf(stderr, "invalid port specified.\n");
					return -1;
				}
				break;
			case 'd':
				daemonize = 1;
				break;
			case 'g':
				if (daemonize) {
					strncpy(group, optarg, sizeof(group) - 1);
				} else {
					fprintf(stderr, "-g requires -d\n");
					return -1;
				}
				break;
			case 'u':
				if (daemonize) {
					strncpy(user, optarg, sizeof(user) - 1);
				} else {
					fprintf(stderr, "-u requires -d\n");
					return -1;
				}
				break;
			case 'l':
				loglevset = 1;
				if (set_loglevel(optarg)) {
					fprintf(stderr, "invalid option for -l\n");
					usage(argv[0]);
					return -1;
				}
				break;
			 case 'h': case '?':
				usage(argv[0]);
				return -1;
		}
	}
	if (!loglevset) {
		setlogmask(LOG_UPTO(LOG_INFO));
	}
#ifdef __illumos__
	openlog(argv[0], LOG_PID | LOG_NDELAY, LOG_FTP);
#else
	openlog(argv[0], LOG_PID | LOG_PERROR | LOG_NDELAY, LOG_FTP);
#endif
	syslog(LOG_INFO, "starting up...");
	signal(SIGTERM, &stop_handler);
	signal(SIGINT, &stop_handler);
	char *buf = NULL;
	int s = socket(AF_INET, SOCK_DGRAM, 0);
	if (s <= 0) {
		syslog(LOG_ERR, "unable to bind socket!");
		return -1;
	}
	if (daemonize) {
		if (daemon(1, 0)) {
			syslog(LOG_ERR, "failed to daemonize: %s", strerror(errno));
			retme = -1;
			goto cleanup;
		} else {
			syslog(LOG_INFO, "daemonized");
		}
	}
	struct timeval timeout = { 30, 0 };
	if (setsockopt(s, SOL_SOCKET, SO_RCVTIMEO, &timeout, sizeof(timeout)) < 0) {
		syslog(LOG_ERR, "unable to set socket timeout: %s", strerror(errno));
		retme = -1;
		goto cleanup;
	}
	struct sockaddr_in saddr;
	saddr.sin_family = AF_INET;
	saddr.sin_port = htons(port);
	saddr.sin_addr.s_addr = inet_addr(addr);
	if (bind(s, (struct sockaddr*)&saddr, sizeof(saddr)) < 0) {
		syslog(LOG_ERR, "socket bind failed: %s", strerror(errno));
		retme = -1;
		goto cleanup;
	}
	syslog(LOG_DEBUG, "socket bind success");
	if (daemonize) {
		//TODO: use getpwnam_r() and getgrnam_r()
		struct group *g = getgrnam((const char *)&group);
		if (setgid((*g).gr_gid) == -1) {
			syslog(
				LOG_ERR,
				"failed to drop group privileges: %s",
				strerror(errno)
			);
			retme = -1;
			goto cleanup;
		}
		struct passwd *u = getpwnam((const char *)&user);
		if (setuid((*u).pw_uid) == -1) {
			syslog(
				LOG_ERR,
				"failed to drop user privileges: %s",
				strerror(errno)
			);
			retme = -1;
			goto cleanup;
		}
	}
#ifdef __illumos__
	//TODO: illumos priv.h privdrop
#endif
	buf = malloc(1024);
	if (buf == NULL) {
		syslog(LOG_CRIT, "unable to allocate memory: %s", strerror(errno));
		retme = -1;
		goto cleanup;
	}
	while (!nbtpd_stop) {
		struct sockaddr_in caddr;
		unsigned int clen = sizeof(caddr);
		memset(buf, '\0', 1024);
		ssize_t recvd_bytes = 0;
		if ((recvd_bytes = recvfrom(s, buf, 1024, 0,
				(struct sockaddr *)&caddr, &clen)) < 0) {
			continue;
		}
		if (recvd_bytes < 4) {
			syslog(
				LOG_ERR,
				"ignoring invalid packet from %s.",
				inet_ntoa(caddr.sin_addr)
			);
			continue;
		}
		syslog(
			LOG_INFO,
			"got connection from %s:%d",
			inet_ntoa(caddr.sin_addr),
			ntohs(caddr.sin_port)
		);
		nbd_nbtpd_args *args = malloc(sizeof(nbd_nbtpd_args));
		if (args == NULL) {
			syslog(
				LOG_CRIT,
				"unable to allocate memory: %s",
				strerror(errno)
			);
			continue;
		}
		nbd_tftp_opcode oc = ((uint16_t)buf[0] << 8) + buf[1];
		args->client = caddr;
		args->err = 0;
		memset(args->path, '\0', NBD_NBTPD_ARGS_PATH_MAX);
		memset(args->mode, '\0', NBD_NBTPD_ARGS_MODE_MAX);
		int i;
		for(i = 2; buf[i] != 0 && i < recvd_bytes && (i - 2) < (NBD_NBTPD_ARGS_PATH_MAX - 1); i++) {
			args->path[i - 2] = buf[i];
		}
		int s = i + 1;
		if (s >= recvd_bytes) {
			oc = 0;
		} else {
			for(i = s; buf[i] != 0 && i < recvd_bytes && (i - s) < (NBD_NBTPD_ARGS_MODE_MAX - 1); i++) {
				args->mode[i - s] = buf[i];
			}
		}
		//TODO: use std threads
		pthread_t _thread;
		void *(*func)(void *);
		switch (oc) {
			case RRQ:
				func = &read_req_resp;
				break;
			case WRQ:
				func = &write_req_resp;
				break;
			default:
				(*args).err = 4;
				func = &nbd_nbtpd_resp_error;
		}
		//TODO: use std threads
		int e;
		syslog(LOG_ERR, "spawning thread for %s", inet_ntoa(caddr.sin_addr));
		if ((e = pthread_create(&_thread, 0, func, (void *)args)) != 0) {
			syslog(
				LOG_CRIT,
				"unable to spawn thread: %s",
				strerror(e)
			);
			free(args);
		}
	}
cleanup:
	free(buf);
	close(s);
	pthread_exit(NULL);
	return retme;
}