wstationd/src/wstationd.c

#include <stdio.h>
#include <unistd.h>
#include <time.h>
#include <errno.h>
#include <string.h>
#include <stdlib.h>
#include <limits.h>
#include <fcntl.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <libgen.h>
#include <getopt.h>
#include "wstationd.h"

#define BUFFERSIZE 4096


static char *exec_name;


static void print_help()
{
	fprintf(stdout, "usage: %s [-p port] [-b addr] directory\n", exec_name);
	fprintf(stdout, "  -p, --port <port>     Port to listen on (default: 10800)\n");
	fprintf(stdout, "  -b, --bind <address>  Address to bind to (default: 0.0.0.0)\n");
	fprintf(stdout, "  -h, --help            Print this message and exit\n");
	fprintf(stdout, "  -v, --version         Print version and exit\n");
}


int main(int argc, char *argv[])
{
	const struct option longopts[] = {
		{ "port", required_argument, NULL, 'p' },
		{ "bind", required_argument, NULL, 'b' },
		{ "version", no_argument, NULL, 'v' },
		{ "help", no_argument, NULL, 'h' },
		{ NULL, 0, NULL, 0 }
	};

	fd_set readfds, masterfds;

	int port = 10800; // Default port
	uint32_t address = INADDR_ANY; // Default listen address
	int connflag, sockflag = 1;
	int sockfd = 0, connfd = 0;
	struct sockaddr_in sock, client;
	socklen_t client_sz = sizeof(client);
	char connbuff[BUFFERSIZE], timebuff[15], fnbuff[35];
	ssize_t connbuff_sz;

	time_t timeraw;
	struct tm* timeptr;

	// Allocate space for the output fds
	int writefds[FD_SETSIZE];
	memset(&writefds, 0, sizeof(int) * FD_SETSIZE);

	exec_name = basename(argv[0]);

	for(int ch; (ch = getopt_long(argc, argv, "b:p:vh0", longopts, NULL)) != -1;){
		switch(ch){
			case 'p':
				port = (int)strtol(optarg, NULL, 10);
				if(port < 1 || port > 65535){
					fprintf(stderr, "%s: invalid port number\n", exec_name);
					print_help();
					goto shutdown_error;
				}
			break;

			case 'b':
				if(inet_pton(AF_INET, optarg, &address) == -1){
					fprintf(stderr, "%s: invalid bind address\n", exec_name);
					print_help();
					goto shutdown_error;
				}
			break;

			case 'v':
				fprintf(stdout, "wstationd v%s (%s %s)\n", VERSION,__DATE__, __TIME__);
			goto shutdown_clean;

			case '?':
			case 'h':
				print_help();
			goto shutdown_clean;
		}
	}

	// Accept final argument, the directory
	if(argc == (optind+1)){
		// Fail if we can't chdir to that directory
		if(chdir(argv[argc-1]) == -1){
			fprintf(stderr, "%s: cannot change directory - %s\n",
				exec_name, strerror(errno)
			);
			goto shutdown_error;
		}
	} else {
		print_help();
		goto shutdown_error;
	}
	
	// Establish the socket
	sockfd = socket(AF_INET, SOCK_STREAM, 0);
	if(sockfd == -1){
		fprintf(stderr, "%s: cannot open socket - %s\n",
			exec_name, strerror(errno)
		);
		goto shutdown_error;
	}

	// This sets REUSE on the socket so it's easily reallocated if
	// this program dies
	setsockopt(sockfd, SOL_SOCKET, SO_REUSEADDR, &sockflag, sizeof(sockflag));

	memset(&sock, 0, sizeof(sock));
	sock.sin_family = AF_INET;
	sock.sin_addr.s_addr = address;
	sock.sin_port = htons(port);

	// Bind our socket to the specified address and port
	if(bind(sockfd, (struct sockaddr*)&sock, sizeof(sock)) < 0){
		fprintf(stderr, "%s: cannot bind - %s\n",
			exec_name, strerror(errno)
		);
		goto shutdown_error;
	}

	// Listen on port
	if(listen(sockfd, 10) != 0){
		fprintf(stderr, "%s: cannot listen - %s\n",
			exec_name, strerror(errno)
		);
		goto shutdown_error;
	}

	FD_ZERO(&masterfds);
	FD_SET(sockfd, &masterfds);

	while(1){
		readfds = masterfds;
		if(select(FD_SETSIZE, &readfds, NULL, NULL, NULL) == -1){
			fprintf(stderr, "%s: socket select error - %s\n",
				exec_name, strerror(errno)
			);
			goto shutdown_error;
		}

		for(int i = 0; i < FD_SETSIZE; i++){
			if(FD_ISSET(i, &readfds)){
				// Deal with new connections
				if(i == sockfd){
					connfd = accept(sockfd, (struct sockaddr*)&client, &client_sz);
					if(connfd == -1){
						fprintf(stderr, "%s: accept failed - %s\n",
							exec_name, strerror(errno)
						);
						continue;
					}

					// Get current flags before setting nonblock
					if((connflag = fcntl(connfd, F_GETFL)) == -1){
						fprintf(stderr, "%s: cannot get connection flags - %s\n",
							exec_name, strerror(errno)
						);
						close(connfd);
						continue;
					}

					// Set nonblocking on socket
					if(fcntl(connfd, F_SETFL, connflag | O_NONBLOCK) == -1){
						fprintf(stderr, "%s: cannot set connection flags - %s\n",
							exec_name, strerror(errno)
						);
						close(connfd);
						continue;
					}

					// Collect the time of the connection, process it into 
					// a short string
					time(&timeraw);
					timeptr = localtime(&timeraw);
					strftime(&timebuff[0], 15, "%Y%m%d%H%M%S", timeptr);
					timeptr = NULL;

					// Take the short date/time connection string and add the client address
					// to create a filename
					snprintf(&fnbuff[0], 34, "%s-%s.dat",
						timebuff, inet_ntoa(client.sin_addr)
					);

					// Open a file to dump our data to
					if((writefds[connfd] = open(fnbuff, O_WRONLY | O_CREAT, 0644)) == -1){
						writefds[connfd] = 0;
						fprintf(stderr, "%s: cannot open - %s\n",
							exec_name, strerror(errno)
						);

						close(connfd);
					} else {
						FD_SET(connfd, &masterfds);
					}

				// Read from existing connections
				} else {
					// Do the read
					connbuff_sz = recv(i, connbuff, BUFFERSIZE, 0);
					if(connbuff_sz > 0){
						// If the read fails, throw an error, and set the
						// buffer size to 0, which forces a disconnect.
						if(write(writefds[i], connbuff, connbuff_sz) == -1){
							fprintf(stderr, "%s: file write error - %s\n",
								exec_name, strerror(errno)
							);
							connbuff_sz = 0;
						}
					}

					// If the buffer is empty (which with select(), will only happen
					// on disconnect, disconnect and cleanup.
					if(connbuff_sz == 0){
						close(writefds[i]);
						writefds[i] = 0;
						FD_CLR(i, &masterfds);
						close(i);
					} else if(connbuff_sz == -1){
						fprintf(stderr, "%s: recv error - %s\n",
							exec_name, strerror(errno)
						);
						FD_CLR(i, &masterfds);
						close(i);
					}
				}
			}
		}
	}

	shutdown_clean:
	if(sockfd > 0) close(sockfd);
	return 0;

	shutdown_error:
	if(sockfd > 0) close(sockfd);
	return 1;
}