path: root/src/main.c

/*
Copyright (C) 2017 Paweł Redman

This program is free software; you can redistribute it and/or
modify it under the terms of the GNU General Public License
as published by the Free Software Foundation; either version 3
of the License, or (at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program; if not, write to the Free Software Foundation,
Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301  USA
*/

#include "shared.h"
#include <time.h>
#include <signal.h>
#include <sys/select.h>
#include <fcntl.h>

int server_sockfd;

static struct timespec time_ref;

uint64_t get_time(void)
{
	static struct timespec now;

	clock_gettime(CLOCK_MONOTONIC, &now);
	return (uint64_t)(now.tv_sec - time_ref.tv_sec) * TIME_SECOND +
	       now.tv_nsec + 1;
}

static int parse_req_addr(char *begin, char *end, uint32_t *out)
{
	char *p;
	size_t i = 0, parts[4] = {0, 0, 0, 0};

	for (p = begin; p < end; p++) {
		if (*p >= '0' && *p <= '9') {
			parts[i] *= 10;
			parts[i] += *p - '0';

			if (parts[i] > 255)
				return 1;
		} else if (*p == '.') {
			i++;

			if (i > 3)
				return 1;
		} else
			break;
	}

	if (i != 3)
		return 1;

	*out = (parts[0] << 24) | (parts[1] << 16) | (parts[2] << 8) | parts[3];
	return 0;
}

static void handle_request(struct sockaddr_in *ret_addr, uint32_t query)
{
	db_entry *entry;
	job_t *reply_job;
	uint64_t now = get_time();

	DEBUG("query %08X\n", query);

	entry = db_find(query);
	if (!entry) {
		eprintf("handle_request: out of memory\n");
		return;
	}

	reply_job = job_create(JOB_REPLY, entry);
	reply_job->ipv4 = query;
	memcpy(&reply_job->ret_addr, ret_addr, sizeof(struct sockaddr_in));

	if (entry->revdns.state == DB_VALID
	    && entry->revdns.exp_time <= now) {
		entry->revdns.state = DB_INVALID;
		DEBUG("revDNS for %08X expired\n", query);
	}

	if (entry->whois.state == DB_VALID
	    && entry->whois.exp_time <= now) {
		entry->whois.state = DB_INVALID;
		DEBUG("WHOIS for %08X expired\n", query);
	}

	if (entry->revdns.state == DB_VALID &&
	    entry->whois.state == DB_VALID) {
		job_enqueue(reply_job);
		goto out;
	}

	if (entry->revdns.state == DB_INVALID) {
		job_t *job;

		DEBUG("revDNS cache miss for %08X\n", query);
		job = job_create(JOB_REVDNS, entry);
		job->ipv4 = query;
		job_enqueue(job);
		entry->revdns.state = DB_IN_PROGRESS;
	}

	if (entry->whois.state == DB_INVALID) {
		job_t *job;

		DEBUG("WHOIS cache miss for %08X\n", query);
		job = job_create(JOB_WHOIS, entry);
		job->ipv4 = query;
		job_enqueue(job);
		entry->whois.state = DB_IN_PROGRESS;
	}

	eli_append(&entry->waiting_jobs, reply_job, waiting_list);

out:
	pthread_mutex_unlock(&entry->mutex);
}

static volatile sig_atomic_t signals_terminate = false;
static volatile sig_atomic_t signals_reload = false;

// NOTE: The code assumes that signal_handler will only be called in the main
//       thread, which (as far as I know) is always the case on Linux.
static void signal_handler(int signum)
{
	if (signum == SIGINT || signum == SIGTERM)
		signals_terminate = true;

	if (signum == SIGUSR1)
		signals_reload = true;
}

static int handle_signals(void)
{
	if (signals_terminate) {
		DEBUG("received a terminating signal\n");
		return 1;
	}

	if (signals_reload) {
		signals_reload = false;
		eprintf("reloading the lists...\n");

		// Don't initiate an exit if the lists failed to reload.
		// It's better to continue operating with empty lists instead
		// having to be restarted (with an empty revDNS/WHOIS cache).
		if (lists_reload("schachts.list"))
			eprintf("error: couldn't reload the lists\n");
	}

	return 0;
}

static struct {
	int port;
	char *lists_file;
} config = {
	.port = 1337,
	.lists_file = "schachts.list"
};

static int parse_argv(int argc, char **argv)
{
	int opt;

	while ((opt = getopt(argc, argv, "vhp:l:")) != -1) {
		switch (opt) {
		case 'v':
			puts(PROGRAM_NAME " " PROGRAM_VERSION);
			exit(0);

		case 'h':
			puts("usage: " PROGRAM_NAME " [-h] [-v] [-p PORT] [-l LISTS-FILE]");
			exit(0);

		case 'p':
			config.port = atoi(optarg);
			break;

		case 'l':
			config.lists_file = optarg;
			break;

		case '?':
			return 1;
		}
	}

	return 0;
}

#define NUM_WORKERS 8 // FIXME: shouldn't be hardcoded

int main(int argc, char **argv)
{
	int error = 0 /* the exit code */, flags;
	struct sigaction sigact = {.sa_handler = signal_handler};
	struct sockaddr_in sockaddr;
	pthread_t workers[NUM_WORKERS];
	size_t i;

	clock_gettime(CLOCK_MONOTONIC_RAW, &time_ref);

	if (parse_argv(argc, argv))
		return 1;

	sigaction(SIGINT, &sigact, NULL);
	sigaction(SIGTERM, &sigact, NULL);
	sigaction(SIGUSR1, &sigact, NULL);

	if (lists_load(config.lists_file, 0)) {
		eprintf("fatal error: couldn't load the lists\n");
		goto error_lists;
	}

	server_sockfd = socket(AF_INET, SOCK_DGRAM, 0);
	if (server_sockfd == -1) {
		perror("socket");
		error = 1;
		goto error_socket;
	}

	// FIXME: can fcntl actually fail (in practice)?
	flags = fcntl(server_sockfd, F_GETFL, 0);
	fcntl(server_sockfd, F_SETFL, flags | O_NONBLOCK);

	sockaddr.sin_family = AF_INET;
	sockaddr.sin_addr.s_addr = INADDR_ANY;
	sockaddr.sin_port = htons(config.port);

	if (bind(server_sockfd, (void*)&sockaddr, sizeof(sockaddr)) == -1) {
		perror("bind");
		eprintf("couldn't bind to UDP port %hi\n", ntohs(sockaddr.sin_port));
		error = 1;
		goto error_bind;
	}

	for (i = 0; i < NUM_WORKERS; i++)
		pthread_create(workers + i, NULL, worker_main, NULL);

	while (1) {
		fd_set readfds;
		int rv;
		char buffer[256];
		struct sockaddr_in addr;
		socklen_t addrlen = sizeof(addr);
		ssize_t size;
		uint32_t query;

		FD_ZERO(&readfds);
		FD_SET(server_sockfd, &readfds);

		// In case the process catches a signal before the main loop
		// or outside the select call.
		if (signals_terminate)
			break;

		// The socket is non-blocking, so wait for incoming data with
		// select (not recvfrom). Furthermore, select returns if it's
		// interrupted by a signal, which makes breaking the main loop
		// easy to implement.
		rv = select(server_sockfd + 1, &readfds, NULL, NULL, NULL);
		if (rv == -1 && errno == EINTR) {
		interrupted_select:
			if (handle_signals())
				break;
			continue;
		} else if (rv == -1 && errno != EINTR) {
			perror("select");
			break;
		}

		size = recvfrom(server_sockfd, buffer, sizeof(buffer), 0,
		                (void*)&addr, &addrlen);
		if (size == -1) {
			// Sometimes select can wake up even if there's no
			// data to read.
			if (errno == EWOULDBLOCK ||
			    errno == EAGAIN)
				continue;

			// Just in case.
			if (errno == EINTR)
				goto interrupted_select;

			perror("recvfrom");
			break;
		}

		if ((ntohl(addr.sin_addr.s_addr) & LOCALHOST_MASK)
		    != LOCALHOST_NETWORK)
			continue;

		if (size < REQUEST_HEADER_LEN)
			continue;

		if (memcmp(buffer, REQUEST_HEADER, REQUEST_HEADER_LEN))
			continue;

		if (parse_req_addr(buffer + REQUEST_HEADER_LEN, buffer + size,
		                   &query))
			continue;

		handle_request(&addr, query);
	}

	job_quit();
	for (i = 0; i < NUM_WORKERS; i++)
		pthread_join(workers[i], NULL);

	db_destroy();

error_bind:
	shutdown(server_sockfd, SHUT_RDWR);
	close(server_sockfd);
error_socket:
	lists_destroy();
error_lists:
	return error;
}