42_ft_nmap/src/thread.c

#include <netinet/in.h>
#include <pthread.h>
#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

#include "interface.h"
#include "parsing.h"
#include "scan.h"
#include "thread.h"

bool g_start = false;
pthread_mutex_t g_start_mtx;
pthread_mutex_t g_getservent;

void *routine(void *p_data)
{
	while (1) {
		pthread_mutex_lock(&g_start_mtx);
		bool start = g_start;
		pthread_mutex_unlock(&g_start_mtx);
		if (start)
			break;
		usleep(100);
	}

	struct thread *thread_data = (struct thread *)p_data;
	struct scan scan_data = {
	    .dest_addr = thread_data->dest_addr,
	    .host = &thread_data->host,
	    .port_start = thread_data->port_start,
	    .type = thread_data->type,
	};

	for (uint16_t port = thread_data->port_start;
	     port <= thread_data->port_end; port++) {
		scan_data.port = port;
		scan_data.response = &thread_data->responses[port - 1];
		if (scan(&scan_data)) {
			free(p_data);
			return NULL;
		}
	}

	return NULL;
}

static struct thread *init_threads_data(const struct option_lst *options,
                                        char *ip_addr, const struct host *host,
                                        struct response *responses,
                                        uint8_t nb_threads)
{
	struct thread *threads = malloc(sizeof(struct thread) * nb_threads);
	if (!threads) {
		dprintf(2, "ft_nmap: allocation of threads failed\n");
		return NULL;
	}
	for (uint8_t i = 0; i < nb_threads; i++) {
		memcpy(&threads[i].host, host, sizeof(struct host));
		const char *ports = get_option_arg(options, FL_PORTS);
		if (parse_ports(ports, &threads[i].port_start,
		                &threads[i].port_end) < 0)
			goto error;
		threads[i].type = parse_type(get_option_arg(options, FL_SCAN));
		if (threads[i].type < 0)
			goto error;
		threads[i].dest_addr = ip_addr;
		threads[i].responses = responses;
	}

	return threads;
error:
	free(threads);
	return NULL;
}

int create_threads(const struct option_lst *options, char *ip_addr,
                   struct response *responses)
{
	uint16_t port_start = 0, port_end = 0;
	const char *ports = get_option_arg(options, FL_PORTS);
	if (parse_ports(ports, &port_start, &port_end) < 0)
		return -1;
	struct host host;
	if (get_interface_name(&host) < 0)
		return -1;

	const char *arg = get_option_arg(options, FL_SPEEDUP);
	// Launche single thread routine if it's a 1 port scan or if no speedup
	// option was passed
	if (!arg || !port_end) {
		struct thread *thread_data =
		    init_threads_data(options, ip_addr, &host, responses, 1);
		thread_data->port_start = port_start;
		thread_data->port_end = port_end;
		g_start = true;
		routine(thread_data);
		free(thread_data);
		return 0;
	}

	uint8_t nb_threads = atoi(arg);
	struct thread *threads_data =
	    init_threads_data(options, ip_addr, &host, responses, nb_threads);
	if (!threads_data)
		return -1;
	pthread_t *threads = malloc(sizeof(pthread_t) * nb_threads);
	if (!threads) {
		free(threads_data);
		return -1;
	}

	pthread_mutex_init(&g_start_mtx, NULL);
	pthread_mutex_init(&g_getservent, NULL);

	const uint16_t ports_per_thread =
	    (port_end - port_start + 1) / nb_threads;
	uint16_t remaining_ports = (port_end - port_start + 1) % nb_threads;
	for (uint8_t i = 0; i < nb_threads; i++) {
		threads_data[i].port_start = port_start + i * ports_per_thread;
		threads_data[i].port_end = (port_start - 1) +
		                           (i + 1) * ports_per_thread +
		                           (remaining_ports ? 1 : 0);
		if (remaining_ports) {
			remaining_ports--;
			port_start++;
		}
		if (pthread_create(&threads[i], NULL, routine,
		                   &threads_data[i])) {
			dprintf(2, "ft_nmap: error during pthread_create()\n");
			return -1;
		}
	}
	pthread_mutex_lock(&g_start_mtx);
	g_start = true;
	pthread_mutex_unlock(&g_start_mtx);

	for (uint8_t i = 0; i < nb_threads; i++) {
		if (pthread_join(threads[i], NULL)) {
			dprintf(2, "ft_nmap: error during pthread_join()\n");
			return -1;
		}
	}
	free(threads_data);
	free(threads);

	return 0;
}