/*
** CBLM v3
** $Id: packet.c,v 1.14 2006/06/10 05:29:33 sella Exp $
** Copyright (c) 2004 James M. Sella. All Rights Reserved.
** Released under the GPL Version 2 License.
** http://www.digitalgenesis.com
** 
** Cheap Bastard Latency Monitor v3
*/

#include "packet.h"
#include "cblm.h"
#include "probe.h"
#include "prefs.h"
#include "data.h"
#include "db.h"
#include "util.h"
#include "data_hash.h"

/* Global variables */
extern struct args_t args;
extern struct prefs_t prefs;
extern struct probe_t *probes;
extern int running, restart;
extern pthread_mutex_t probe_mutex, packet_list_mutex, db_mutex, hash_mutex;

extern MYSQL mysql;

struct packet_list_t *packets = NULL;
struct packet_list_t **pl_end = &packets;
long p_count = 0;

void transmit_thread() {
	struct probe_t *curp;
	struct tos_t *curt;
	struct protoent *protocol;
	long remaining, pending;
	struct timeval now;
	long start, end, slices_past;
	double slice_packets, should_be_done;
	struct timespec timer;
	long slices;
	unsigned int set = 0;
	time_t t;
	int s = 0, no_probes = 0, optval;

	if (args.debug == 0 && fileno(stdin) != -1) fclose(stdin);
	if (args.debug == 0 && fileno(stdout) != -1) fclose(stdout);
	if (args.debug == 0 && fileno(stderr) != -1) fclose(stderr);

	/* Setup transmit socket */
	protocol = (struct protoent*) getprotobyname("udp");
	if (!protocol) {
		syslog(LOG_CRIT, "Failed to get protocol. Verify that /etc/services is readable from current root directory.");

		running = 0;
		goto exit;
	}

	if ((s = socket(AF_INET, SOCK_DGRAM, protocol->p_proto)) < 0) {
		syslog(LOG_CRIT, "Failed to open socket: %s", strerror(errno));
		goto exit;
	}

#ifdef SO_SNDBUF
	optval = SNDBUF_SIZE;
	if (setsockopt(s, SOL_SOCKET, SO_SNDBUF, (void*) &optval, sizeof(optval)) < 0) {
		syslog(LOG_CRIT, "Failed to setsockopt SO_SNDBUF to %d: %s", optval, strerror(errno));
	}
#endif

#ifdef IP_TTL
	optval = prefs.ttl;
	if (setsockopt(s, IPPROTO_IP, IP_TTL, (void*) &optval, sizeof(optval)) < 0) {
		syslog(LOG_CRIT, "Failed to setsockopt IP_TTL to %d: %s", optval, strerror(errno));
	}
#endif

#ifdef IP_DONTFRAGMENT
	if (prefs.fragment == 1) {
		optval = 0;
		if (setsockopt(s, IPPROTO_IP, IP_DONTFRAGMENT, (void*) &optval, sizeof(optval)) < 0) {
			syslog(LOG_CRIT, "Failed to setsockopt IP_DONTFRAGMENT to %d: %s", optval, strerror(errno));
		}
	}
#endif

	/* Slice sleep time. */
	timer.tv_sec = 0;
	timer.tv_nsec = 1000000 * 1000 / SLICES_PER_SEC;

	while (running) {
		/* Update probe list. */
		read_probes();

		/* If there are no probes, wait and try again. */
		if (probes == NULL) {
			no_probes = 1;
			syslog(LOG_INFO, "No remote probe(s) available. Retry in %lu seconds.", NOPROBES_WAIT);

			t = time(NULL) + NOPROBES_WAIT;
			while (running && t >= time(NULL)) {
				usleep(250000); /* 1/4 sec */
			};

			continue;
		}

		if (no_probes == 1) {
			no_probes = 0;
			syslog(LOG_INFO, "Remote probe(s) found. Transmit cycle beginning.");
		}

		pthread_mutex_lock(&probe_mutex);

		/* Initialize data for transmit cycle */
		gettimeofday(&now, NULL);
		start = tvtousec(now);
		end = start + CYCLE_USEC;
		slice_packets = (double) prefs.rate / CYCLE_SECONDS / SLICES_PER_SEC;
		remaining = prefs.rate;
		slices = 0;

		if (args.debug >= DEBUG_INTERNAL) {
			syslog(LOG_INFO, "Transmit cycle (%lu sec) sending %0.5f pkts/slice over %lu slices.", CYCLE_SECONDS, slice_packets, SLICES_PER_CYCLE);
		}

		if (args.debug >= DEBUG_ERROR && (probes == NULL || prefs.tos == NULL)) {
			syslog(DEBUG_ERROR, "No probes and/or ToS found. Nothing will be sent. probes == %08x, prefs.tos == %08x", probes, prefs.tos);
		}

		/* Transmit cycle loop. */
		while (remaining > 0 && running) {
			/* How many should we transmit this cycle? */
			slices_past = (CYCLE_USEC - (end - tvtousec(now))) / (1000000 / SLICES_PER_SEC);
			should_be_done = (slices_past * slice_packets);
			pending = should_be_done - (prefs.rate - remaining);

			if (pending > 0) {
				/* Loop for probes. */
				for (curp = probes; curp != NULL; curp = curp->next) {
					/* Loop though ToS. */
					for (curt = prefs.tos; curt != NULL; curt = curt->next) {
						transmit_packet(s, &(curp->ip), pending, curt, set);
					}
				}

				/* Subtract what we transmitted. */
				remaining -= pending;
			}

			/* Update now to see if the cycle is complete. */
			gettimeofday(&now, NULL);

			/* Sleep */
			nanosleep(&timer, NULL);

			slices++;
		}

		if (args.debug >= DEBUG_INTERNAL && running) {
			syslog(LOG_INFO, "Transmit cycle complete in %0.4f sec using %lu of %lu slices.", (tvtousec(now) - start) / 1000000.0, slices, SLICES_PER_CYCLE);
		}

		pthread_mutex_unlock(&probe_mutex);

		/* Increament to the next set. */
		set++;

		/* Update preferences. */
		if (running && prefs.last != last_prefs()) {
			if (args.debug >= DEBUG_NOTICE) {
				syslog(LOG_INFO, "Preferences have been updated. Restarting.\n");
			}
			running = 0;
			restart = 1;
		}
	}

exit:
	if (s > 0) {
		close(s);
	}

	pthread_exit(NULL);
}

void receive_thread() {
	int s, n, c, fd, optval;
	pid_t pid;
	fd_set fdset, fds;
	struct timeval timeout;

	if (args.debug == 0 && fileno(stdin) != -1) fclose(stdin);
	if (args.debug == 0 && fileno(stdout) != -1) fclose(stdout);
	if (args.debug == 0 && fileno(stderr) != -1) fclose(stderr);

	/* Bind our socket */
	if ((s = bind_port(&fdset)) == -1) {
		running = 0;

		goto exit;
	}
	n = s + 1;

	while (running) {
		/* Copy our fdset, select screws with it. */
		FD_ZERO(&fds);
		bcopy(&fdset, &fds, sizeof(fd_set));

		/* Use a 3/4 sec timeout. */
		timeout.tv_sec = 0;
		timeout.tv_usec = 750000;

		/* Wait for a connection. */
		c = select(n, &fds, 0, 0, &timeout);
		if (c == -1) {
			syslog(LOG_CRIT, "select() %s", strerror(errno));
			return;
		} else if (c > 0) {
			/* Read incoming packets */
			receive_packet(s);
		}
	}

	/* Close ports */
	if (s > 0) {
		close(s);
	}

exit:
	pthread_exit(NULL);
}

void data_thread() {
	struct timespec timer;
	struct packet_t *packet;
	struct packet_list_t *pl_ptr;
	time_t tm = 0;

	if (args.debug == 0 && fileno(stdin) != -1) fclose(stdin);
	if (args.debug == 0 && fileno(stdout) != -1) fclose(stdout);
	if (args.debug == 0 && fileno(stderr) != -1) fclose(stderr);

	/* Initialize our hash */
	pthread_mutex_lock(&hash_mutex);
	hash_init();
	pthread_mutex_unlock(&hash_mutex);

	/* Slice sleep time. */
	timer.tv_sec = 0;
	timer.tv_nsec = 1000000 * 1000 / SLICES_PER_SEC;

	while (running) {
		if (time(NULL) > tm) {
			tm = time(NULL) + (prefs.maxage / 3);
			register_probe();
		}

		while (packets) {
			/* Retrieve packet from the packet_list. */
			pthread_mutex_lock(&packet_list_mutex);

			/* Grab reference to current packet. */
			pl_ptr = packets;
			packet = &(pl_ptr->packet);

			/* Move to the next packet. */
			packets = packets->next;

			if (packets == NULL) {
				pl_end = &packets;
			}

			/* Decrease packet count by 1 */
			p_count--;

			pthread_mutex_unlock(&packet_list_mutex);

			/* Process the packet. */
			process_packet(packet);

			/* Free the node. */
			free(pl_ptr);
			pl_ptr = NULL;
		}

		/* Sleep */
		if (p_count < 1) {
			nanosleep(&timer, NULL);
		}
	}

	/* Free our hash */
	pthread_mutex_lock(&hash_mutex);
	hash_free();
	pthread_mutex_unlock(&hash_mutex);

	pthread_exit(NULL);
}

int transmit_packet(int s, struct in_addr *dst, unsigned int pending, struct tos_t *tos, unsigned int set) {
	int i, c, ret_val = 1, pattern_len;
	struct sockaddr_in addr;
	char buf[MAX_SIZE];
	struct packet_t* packet = (struct packet_t*) buf;
	size_t packet_len = (prefs.size < sizeof(struct packet_t)) ? sizeof(struct packet_t) : prefs.size;
	struct in_addr in;
	static unsigned int seq = 0, last_set = -1;
	fd_set fds;
	struct timeval timeout;
	unsigned int opt;

	if (args.debug >= DEBUG_DUMP) {
		syslog(LOG_DEBUG, "Transmitting %lu pkts to %s with ToS %lu (dscp: %u, prec: %u) and PRI %lu (size %lu).", pending, inet_ntoa(*dst), tos->tos, tos->tos >> 2, tos->tos >> 5, tos->pri, packet_len);
	}

	/* Reset the seq, if needed. */
	if (set != last_set) {
		seq = 0;
		last_set = set;
	}

	/* Initialize our addr. */
	bzero((char*) &addr, sizeof(addr));
	addr.sin_family = AF_INET;
	addr.sin_addr = *dst;
	addr.sin_port = prefs.port;

#ifdef IP_TOS
	opt = tos->tos & 0x0000000F;
	if (setsockopt(s, IPPROTO_IP, IP_TOS, (void*) &opt, sizeof(opt)) < 0) {
		syslog(LOG_CRIT, "Failed to set ToS value %u (dscp: %u, prec: %u): %s", opt, opt >> 2, opt >> 5, strerror(errno));
	}
#endif

#ifdef SO_PRIORITY
	opt = tos->pri & 0x00000007;
	if (setsockopt(s, SOL_SOCKET, SO_PRIORITY, (void*) &opt, sizeof(opt)) < 0) {
		syslog(LOG_CRIT, "Failed to set PRI value %u: %s", opt, strerror(errno));
	}
#endif

	/* Set generic portions of our packet. */
	memset(packet, 0, packet_len);

	packet->src = args.source;
	packet->dst = *dst;
	memcpy(packet->key, prefs.key, sizeof(packet->key));
	packet->tx_tos = tos->tos;
	packet->tx_pri = tos->pri;
	packet->set = set;
	packet->tot = htons(prefs.rate);

	pattern_len = strlen(prefs.pattern);
	for (i = 0; i < packet_len - sizeof(struct packet_t); i += pattern_len) {
		if (pattern_len + i < packet_len - sizeof(struct packet_t)) {
			memcpy(packet->pattern + i, prefs.pattern, pattern_len);
		} else {
			memcpy(packet->pattern + i, prefs.pattern, packet_len - sizeof(struct packet_t) - i + 1);
		}
	}

	/* Transmit our packets. */
	for (i = 0; i < pending; i++) {
		/* Copy our fdset, select screws with it. */
		FD_ZERO(&fds);
		FD_SET(s, &fds);

		/* Use a 3/4 sec timeout. */
		timeout.tv_sec = 0;
		timeout.tv_usec = 750000;

		/* Wait for a connection. */
		c = select(s + 1, 0, &fds, 0, &timeout);
		if (c == -1) {
			syslog(LOG_CRIT, "select() %s", strerror(errno));
			return;
		} else if (c > 0) {
			/* Set dynamic portions of packet. */
			gettimeofday(&packet->tx, NULL);
			packet->tx = htontv(packet->tx);
			packet->seq = htons(seq++);

			if (sendto(s, (void*) packet, packet_len, 0, (struct sockaddr*) &addr, sizeof(struct sockaddr_in)) == -1) {
				syslog(LOG_CRIT, "sendto() %s", strerror(errno));
				goto exit;
			}
		}
	}

	/* Success */
	ret_val = 0;

exit:
	return ret_val;
}

int receive_packet(int s) {
	int ret_val = 1, packet_len;
	unsigned int tos = 0, pri = 0;
	char buf[MAX_SIZE];
	struct sockaddr_in addr;
	socklen_t len = sizeof(struct sockaddr);
	struct packet_t *packet = (struct packet_t*) buf;
	struct timeval tv, tv2;
	socklen_t tos_len, pri_len;

	while ((packet_len = recvfrom(s, buf, MAX_SIZE, 0, (struct sockaddr*) &addr, &len)) > 0) {
		/* Capture receive time */
		gettimeofday(&tv, NULL);

		/* Discard packets that are invalid. */
		if (packet_len < sizeof(struct packet_t)) {
			/* Discard obviously incorrect packets. */
			if (args.debug >= DEBUG_NOTICE) {
				syslog(LOG_NOTICE, "Recieved packet was too small (%lu < %lu). Discarding.", packet_len, sizeof(struct packet_t));
			}

			return ret_val;
		} else if (strncmp(packet->key, prefs.key, sizeof(prefs.key))) {
			/* Discard packets without correct key (might not even be from CBLM) */
			packet->key[17] = '\0'; /* packet->key is not null terminated. */

			if (args.debug >= DEBUG_NOTICE) {
				syslog(LOG_NOTICE, "Keys do not match (%s != %s). Discarding.", packet->key, prefs.key);
			}

			return ret_val;
		} else if (packet->dst.s_addr == 0 || packet->src.s_addr == 0) {
			if (args.debug >= DEBUG_NOTICE) {
				syslog(LOG_NOTICE, "Source (%lu) or destionation (%lu) are not set. Discarding.\n", packet->src.s_addr, packet->dst.s_addr);
			}

			return ret_val;
		}

		if (args.debug >= DEBUG_DUMP) {
			syslog(LOG_DEBUG, "Received packet (len: %lu)", packet_len);
			display_packet(packet, packet_len);
		}

#ifdef IP_TOS
		if (getsockopt(s, IPPROTO_IP, IP_TOS, (void *) &tos, &tos_len) < 0) {
			syslog(LOG_CRIT, "Failed to get ToS value %lu (dscp: %lu, prec: %lu): %s", tos, tos >> 2, tos >> 5, strerror(errno));
		}
#endif

#ifdef SO_PRIORITY
		if (getsockopt(s, SOL_SOCKET, SO_PRIORITY, (void *) &pri, &pri_len) < 0) {
			syslog(LOG_CRIT, "Failed to get ToS value %lu (dscp: %lu, prec: %lu): %s", tos, tos >> 2, tos >> 5, strerror(errno));
		}
#endif

		if (packet->ox.tv_sec == 0) {
			/* This is a packet sourced from a remote probe. */
			packet->ox = htontv(tv);
			packet->ox_tos = tos;
			packet->ox_pri = pri;

			if (args.debug >= DEBUG_DUMP) {
				tv = ntohtv(packet->ox);
				tv2 = ntohtv(packet->tx);
				subtract_timeval(&tv, &tv2);
				syslog(LOG_DEBUG, "OTT: %0.2f ms", tvtousec(tv) / 1000.0);
			}

			reflect_packet(s, packet, packet_len, (struct sockaddr_in*) &addr);
		} else {
			/* This is a packet we originated returning to us. */
			packet->rx = htontv(tv);
			packet->rx_tos = tos;
			packet->rx_pri = pri;

			if (args.debug >= DEBUG_DUMP) {
				tv = ntohtv(packet->rx);
				tv2 = ntohtv(packet->tx);
				subtract_timeval(&tv, &tv2);
				syslog(DEBUG_DUMP, "RTT: %0.2f ms", tvtousec(tv) / 1000.0);
			}

			/* Add this packet to the packet_list. */
			pthread_mutex_lock(&packet_list_mutex);

			*pl_end = malloc(sizeof(struct packet_list_t));
			memcpy(&(*pl_end)->packet, packet, sizeof(struct packet_t));

			(*pl_end)->next = NULL;
			pl_end = &(*pl_end)->next;
			p_count++;

			pthread_mutex_unlock(&packet_list_mutex);
		}
	}

	if (packet_len == -1) {
		if (errno != EAGAIN && errno == EINTR) {
			syslog(LOG_CRIT, "recvfrom() %d: %s", errno, strerror(errno));
		}

		/* Failure */
		ret_val = 0;
	}

	return ret_val;
}

int reflect_packet(int s, struct packet_t *packet, size_t packet_len, struct sockaddr_in *addr) {
	int c, ret_val = 1, optval;
	unsigned int tos, pri;
	fd_set fds;
	struct timeval timeout;

	tos = packet->tx_tos & 0x0000000F;
	pri = packet->tx_pri & 0x00000007;

	if (args.debug >= DEBUG_DUMP) {
		syslog(DEBUG_DUMP, "Reflecting packet to %s with ToS %lu (size %lu).", inet_ntoa((struct in_addr) addr->sin_addr), tos, packet_len);
	}

	/* Fixup the port. */
	addr->sin_port = prefs.port;

#ifdef IP_TOS
	if (setsockopt(s, IPPROTO_IP, IP_TOS, (void*) &tos, sizeof(tos)) < 0) {
		syslog(LOG_CRIT, "Failed to set ToS value %u (dscp: %u, prec: %u): %s", tos, tos >> 2, tos >> 5, strerror(errno));
	}
#endif

#ifdef SO_PRIORITY
	if (setsockopt(s, SOL_SOCKET, SO_PRIORITY, (void*) &pri, sizeof(pri)) < 0) {
		syslog(LOG_CRIT, "Failed to set PRI value %u: %s", pri, strerror(errno));
	}
#endif

	FD_ZERO(&fds);
	FD_SET(s, &fds);

	/* Use a 3/4 sec timeout. */
	timeout.tv_sec = 0;
	timeout.tv_usec = 750000;

#if 0
	if (sendto(s, (void*) packet, packet_len, 0, (struct sockaddr*) addr, sizeof(struct sockaddr_in)) == -1) {
		syslog(LOG_CRIT, "sendto() error: %s", strerror(errno));
		goto exit;
	}
#else
	/* Wait for socket write. */
	c = select(s + 1, 0, &fds, 0, &timeout);
	if (c == -1) {
		syslog(LOG_CRIT, "select() %s", strerror(errno));
		return;
	} else if (c > 0) {
		if (sendto(s, (void*) packet, packet_len, 0, (struct sockaddr*) addr, sizeof(struct sockaddr_in)) == -1) {
			syslog(LOG_CRIT, "sendto() error: %s", strerror(errno));
			goto exit;
		}
	}
#endif

	/* Success */
	ret_val = 0;

exit:
	return ret_val;
}

int bind_port(fd_set *fdset) {
	int s = 0, port, optval, len, flags;
	struct protoent *protocol;
	struct sockaddr_in socketaddr;

	protocol = (struct protoent*) getprotobyname("udp");
	if (!protocol) {
		syslog(LOG_CRIT, "Failed to get protocol. Verify that /etc/services is readable from current root directory.");

		return -1;
	}

	if (args.debug >= DEBUG_INTERNAL) {
		syslog(LOG_INFO, "Binding %s:%d", inet_ntoa(args.interface), ntohs(prefs.port));
	}

	/* Zero and setup our sockaddr_in structure. */
	bzero((char*) &socketaddr, sizeof(socketaddr));
	socketaddr.sin_family = AF_INET;
	socketaddr.sin_addr = args.interface;
	socketaddr.sin_port = prefs.port;

	if ((s = socket(AF_INET, SOCK_DGRAM, protocol->p_proto)) < 0) {
		syslog(LOG_CRIT, "Failed to open socket: %s", strerror(errno));
		return -1;
	}

#ifdef SO_RCVBUF
	/* Increase the receive buffer. */
	optval = RCVBUF_SIZE;
	if (setsockopt(s, SOL_SOCKET, SO_RCVBUF, (void*) &optval, sizeof(optval)) < 0) {
		syslog(LOG_CRIT, "Failed to setsockopt SO_RCVBUF to %d: %s", optval, strerror(errno));
	}
#endif

#ifdef SO_SNDBUF
	/* Increase the send buffer. */
	optval = SNDBUF_SIZE;
	if (setsockopt(s, SOL_SOCKET, SO_SNDBUF, (void*) &optval, sizeof(optval)) < 0) {
		syslog(LOG_CRIT, "Failed to setsockopt SO_SNDBUF to %d: %s", optval, strerror(errno));
	}
#endif

#ifdef IP_TTL
	optval = prefs.ttl;
	if (setsockopt(s, IPPROTO_IP, IP_TTL, (void*) &optval, sizeof(optval)) < 0) {
		syslog(LOG_CRIT, "Failed to setsockopt IP_TTL to %d: %s", optval, strerror(errno));
	}
#endif

#ifdef IP_DONTFRAGMENT
	if (prefs.fragment == 1) {
		optval = 0;
		if (setsockopt(s, IPPROTO_IP, IP_DONTFRAGMENT, (void*) &optval, sizeof(optval)) < 0) {
			syslog(LOG_CRIT, "Failed to setsockopt IP_DONTFRAGMENT to %d: %s", optval, strerror(errno));
		}
	}
#endif

#ifdef SO_NO_CHECK /* Linux only feature */
	optval = 1; /* 1 == disable checking of checksum */
	if (setsockopt(s, SOL_SOCKET, SO_NO_CHECK, (void*) &optval, sizeof(optval)) < 0) {
		syslog(LOG_CRIT, "Failed to setsockopt SO_NO_CHECK to %d: %s", optval, strerror(errno));
	}
#endif

	/* Set non-blocking */
	flags = fcntl(s, F_GETFL);
	if (fcntl(s, F_SETFL, flags | O_NONBLOCK) == -1) {
		syslog(LOG_CRIT, "Failed to fcntl O_NONBLOCK: %s", strerror(errno));
	}

	if (bind(s, (struct sockaddr*) &socketaddr, sizeof(socketaddr)) < 0) {
		syslog(LOG_CRIT, "Failed to bind to %s:%d: %s", inet_ntoa(args.interface), ntohs(prefs.port), strerror(errno));
		return -1;
	}

	/* Add new socket to our FDSET */
	FD_ZERO(fdset);
	FD_SET(s, fdset);

	if (args.debug >= DEBUG_INTERNAL) {
		syslog(LOG_INFO, "Bound socket '%d'.", s);
	}

	return s;
}

/*
** Local Variables:
** c-basic-offset: 3
** tab-width: 3
** End:
** vim: noet ts=3 sw=3
*/
