digFTP/src/main.cpp

#include <iostream>
#include <string>
#include <unistd.h>
#include <cstring>
#include <future>
#include <thread>
#include <chrono>
#include <stdio.h>
#include <stdlib.h>
#include <sys/ioctl.h>
#include <sys/poll.h>
#include <sys/socket.h>
#include <sys/time.h>
#include <netinet/in.h>
#include <errno.h>
#include <fcntl.h>

#include "server.h"
#include "client.cpp"
#include "util.h"

using namespace std::chrono_literals;

// I stole the majority* of this multi-client sock code from IBM
// of all places. Had to still modify it for my own uses but
// laziness wins again!

struct pollfd fds[MAXCLIENTS];
struct clientfd {
	Client* client;
	std::thread thread;
	bool close = false;
} fdc[MAXCLIENTS];

bool run = true, compress_array = false;

void runClient(struct clientfd* cfd) {
	char inbuf[BUFFERSIZE];
	printf("[d] C(%i) Initialized\n", cfd->client->control_sock);
	// Loop as long as it is a valid file descriptor.
	while (fcntl(cfd->client->control_sock, F_GETFD) != -1) {
		printf("[d] C(%i) Attempting read...\n", cfd->client->control_sock);
		int rc = recv(cfd->client->control_sock, inbuf, sizeof(inbuf), 0);
		if (rc < 0) {
			if (errno != EWOULDBLOCK) {
				perror("recv() failed");
				break;
			}
			continue;
		}
		
		if (rc == 0 || cfd->client == nullptr) {
			printf("[d] C(%i) closed\n", cfd->client->control_sock);
			break;
		}
		
		std::string lin(inbuf);
		int len = lin.find("\r\n", 0);
		int cmdend = lin.find(" ", 0);
		if (cmdend >= len || cmdend == std::string::npos) cmdend = len;
		std::string cmd = toUpper(lin.substr(0, cmdend));
		std::string args = "";
		if (len > cmdend) args = lin.substr(cmdend+1, len-cmdend-1);
			
		printf("[d] C(%i) >> '%s' '%s'\n", cfd->client->control_sock, cmd.c_str(), args.c_str());
		
		if (cfd->client->receive(cmd, args) < 0) break;
		inbuf[0] = '\0';
	}
	printf("[d] C(%i) Marking for deletion...\n", cfd->client->control_sock);
	cfd->close = true;
	cfd->thread.detach();
}

int main(int argc , char *argv[]) {
	int opt = 1,
		master_socket = -1,
		newsock = -1,
		nfds = 1,
		current_size = 0;
	
	char inbuf[BUFFERSIZE];
	struct sockaddr_in ctrl_address;
	
	if ((master_socket = socket(AF_INET , SOCK_STREAM , 0)) < 0) {
		perror("socket() failed");
		exit(-1);
	}
	
	if (setsockopt(master_socket, SOL_SOCKET, SO_REUSEADDR, (char *)&opt, sizeof(opt)) < 0) {
		perror("setsockopt() failed");
		close(master_socket);
		exit(-1);
	}
	
	if (ioctl(master_socket, FIONBIO, (char *)&opt) < 0) {
		perror("ioctl() failed");
		close(master_socket);
		exit(-1);
	}
	
	ctrl_address.sin_family = AF_INET;
	ctrl_address.sin_addr.s_addr = INADDR_ANY;
	ctrl_address.sin_port = htons(CONTROL_PORT);
	
	if (bind(master_socket, (struct sockaddr *)&ctrl_address, sizeof(ctrl_address))<0 < 0) {
		perror("bind() failed");
		close(master_socket);
		exit(-1);
	}
	
	if (listen(master_socket, 3) < 0) {
		perror("listen() failed");
		close(master_socket);
		exit(-1);
	}
	
	memset(fds, 0 , sizeof(fds));
	memset(fdc, 0 , sizeof(fdc));
	
	fds[0].fd = master_socket;
	fds[0].events = POLLIN;
	
	printf("[i] Server started.\n");
	while (run) {
		int pc = poll(fds, nfds, -1);
		
		if (pc < 0) {
			perror("poll() failed");
			break;
		}
		
		if (pc == 0) {
			printf("poll() timed out\n");
			break;
		}
		
		current_size = nfds;
		for (int i = 0; i < current_size; i++) {
			inbuf[0] = '\0';
			
			if(fds[i].revents == 0)
				continue;
			
			if(fds[i].revents != POLLIN) {
				printf("[!] C(%i) Error! revents = %d\n", fds[i].fd, fds[i].revents);
				goto conn_close;
			}
			
			if (fds[i].fd == master_socket) {
				do {
					newsock = accept(master_socket, NULL, NULL);
					if (newsock < 0) {
						if (errno != EWOULDBLOCK) {
							perror("accept() failed");
							run = false;
						}
						break;
					}
					
					printf("[d] C(%i) Accepted client%i\n", newsock);
					// If we assign thead if it is still attached,
					// we suffer greatly, and by we I mean me.
					if (fdc[nfds].thread.joinable()) {
						printf("[!] C(%i) Thread still joinable! Detaching...\n", newsock);
						// Pray the thread will end itself.
						fdc[nfds].thread.detach();
					}
					fds[nfds].fd = newsock;
					fds[nfds].events = POLLIN;
					fdc[nfds].client = new Client(newsock);
					fdc[nfds].thread = std::thread(runClient, &fdc[nfds]);
					nfds++;
				} while (newsock != -1);
			} else {
				/*
				int rc = recv(fds[i].fd, inbuf, sizeof(inbuf), 0);
				if (rc < 0) {
					if (errno != EWOULDBLOCK) {
						perror("recv() failed");
						fdc[i].close = true;
					}
					continue;
				}
				
				if (rc == 0 || fdc[i].client == nullptr) {
					printf("[d] (%i) closed\n", fds[i].fd);
					fdc[i].close = true;
				}
				
				std::string lin(inbuf);
				int len = lin.find("\r\n", 0);
				int cmdend = lin.find(" ", 0);
				if (cmdend >= len || cmdend == std::string::npos) cmdend = len;
				std::string cmd = toUpper(lin.substr(0, cmdend));
				std::string args = "";
				if (len > cmdend) args = lin.substr(cmdend+1, len-cmdend-1);
				
				printf("[d] (%i) >> '%s' '%s'\n", fds[i].fd, cmd.c_str(), args.c_str());
				
				if (fdc[i].client->receive(cmd, args) < 0) fdc[i].close = true;
				inbuf[0] = '\0';
				*/
				
				if (fdc[i].close) {
					conn_close:
					printf("[d] C(%i) Deleting client...\n", fds[i].fd);
					close(fds[i].fd);
					fds[i].fd = -1;
					if (fdc[i].thread.joinable()) fdc[i].thread.detach();
					fdc[i].client = nullptr;
					fdc[i].close = false;
					compress_array = true;
				}
			}
		}
	
		if (compress_array) {
			compress_array = false;
			printf("[d] Compressing...\n");
			for (int i = 0; i < nfds; i++) {
				if (fds[i].fd == -1) {
					for(int j = i; j < nfds; j++) {
						printf("[d] Compressing: id %i to fd %i\n", j, fds[j+1].fd);
						fds[j].fd = fds[j+1].fd;
					}
					i--;
					nfds--;
				}
			}
			printf("[d] Compressing complete!\n");
		}
	}
	return 0;
}