#include <sys/cdefs.h>

#include <sys/param.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <sys/socketvar.h>
#include <sys/sysctl.h>

#include <netinet/in.h>
#include <netinet/in_pcb.h>
#define TCPSTATES
#include <netinet/tcp_fsm.h>
#include <netinet/tcp_var.h>

#include <err.h>
#include <netdb.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

#define	TCPDROP_FOREIGN		0
#define	TCPDROP_LOCAL		1

struct host_service {
	char hs_host[NI_MAXHOST];
	char hs_service[NI_MAXSERV];
};

static bool tcpdrop_list_commands = false;

static void sockinfo(const struct sockaddr *, struct host_service *);
static bool tcpdrop(const struct sockaddr *, const struct sockaddr *);
static bool tcpdropbyname(const char *, const char *, const char *,
    const char *);
static bool tcpdropconn(const struct in_conninfo *);

static void
sockinfo(const struct sockaddr *sa, struct host_service *hs)
{
	static const int flags = NI_NUMERICHOST | NI_NUMERICSERV;
	int rv;

	rv = getnameinfo(sa, sa->sa_len, hs->hs_host, sizeof hs->hs_host,
	    hs->hs_service, sizeof hs->hs_service, flags);
	if (rv == -1)
		err(1, "getnameinfo");
}

static bool
tcpdrop(const struct sockaddr *lsa, const struct sockaddr *fsa)
{
	struct host_service local, foreign;
	struct sockaddr_storage addrs[2];
	int rv;

	memcpy(&addrs[TCPDROP_FOREIGN], fsa, fsa->sa_len);
	memcpy(&addrs[TCPDROP_LOCAL], lsa, lsa->sa_len);

	sockinfo(lsa, &local);
	sockinfo(fsa, &foreign);

	if (tcpdrop_list_commands) {
		//printf("tcpdrop %s %s %s %s\n", local.hs_host, local.hs_service,
		//    foreign.hs_host, foreign.hs_service);
		return (true);
	}

	rv = sysctlbyname("net.inet.tcp.drop", NULL, NULL, &addrs,
	    sizeof addrs);
	if (rv == -1) {
		//warn("%s %s %s %s", local.hs_host, local.hs_service,
		//    foreign.hs_host, foreign.hs_service);
		return (false);
	}
	//printf("%s %s %s %s: dropped\n", local.hs_host, local.hs_service,
	//    foreign.hs_host, foreign.hs_service);
	return (true);
}

static bool
tcpdropbyname(const char *lhost, const char *lport, const char *fhost,
    const char *fport)
{
	static const struct addrinfo hints = {
		/*
		 * Look for streams in all domains.
		 */
		.ai_family = AF_UNSPEC,
		.ai_socktype = SOCK_STREAM,
	};
	struct addrinfo *ail, *local, *aif, *foreign;
	int error;
	bool ok, infamily;

	error = getaddrinfo(lhost, lport, &hints, &local);
	if (error != 0)
		errx(1, "getaddrinfo: %s port %s: %s", lhost, lport,
		    gai_strerror(error));

	error = getaddrinfo(fhost, fport, &hints, &foreign);
	if (error != 0) {
		freeaddrinfo(local); /* XXX gratuitous */
		errx(1, "getaddrinfo: %s port %s: %s", fhost, fport,
		    gai_strerror(error));
	}

	ok = true;
	infamily = false;

	/*
	 * Try every combination of local and foreign address pairs.
	 */
	for (ail = local; ail != NULL; ail = ail->ai_next) {
		for (aif = foreign; aif != NULL; aif = aif->ai_next) {
			if (ail->ai_family != aif->ai_family)
				continue;
			infamily = true;
			if (!tcpdrop(ail->ai_addr, aif->ai_addr))
				ok = false;
		}
	}

	if (!infamily) {
		//warnx("%s %s %s %s: different address families", lhost, lport,
		//    fhost, fport);
		ok = false;
	}

	freeaddrinfo(local);
	freeaddrinfo(foreign);

	return (ok);
}

static bool
tcpdropconn(const struct in_conninfo *inc)
{
	struct sockaddr *local, *foreign;
	struct sockaddr_in6 sin6[2];
	struct sockaddr_in sin4[2];

	if (inc->inc_isipv6) {
		memset(sin6, 0, sizeof sin6);

		sin6[TCPDROP_LOCAL].sin6_len = sizeof sin6[TCPDROP_LOCAL];
		sin6[TCPDROP_LOCAL].sin6_family = AF_INET6;
		sin6[TCPDROP_LOCAL].sin6_port = inc->inc_lport;
		memcpy(&sin6[TCPDROP_LOCAL].sin6_addr, &inc->inc6_laddr,
		    sizeof inc->inc6_laddr);
		local = (struct sockaddr *)&sin6[TCPDROP_LOCAL];

		sin6[TCPDROP_FOREIGN].sin6_len = sizeof sin6[TCPDROP_FOREIGN];
		sin6[TCPDROP_FOREIGN].sin6_family = AF_INET6;
		sin6[TCPDROP_FOREIGN].sin6_port = inc->inc_fport;
		memcpy(&sin6[TCPDROP_FOREIGN].sin6_addr, &inc->inc6_faddr,
		    sizeof inc->inc6_faddr);
		foreign = (struct sockaddr *)&sin6[TCPDROP_FOREIGN];
	} else {
		memset(&sin4[TCPDROP_LOCAL], 0, sizeof sin4[TCPDROP_LOCAL]);

		sin4[TCPDROP_LOCAL].sin_len = sizeof sin4[TCPDROP_LOCAL];
		sin4[TCPDROP_LOCAL].sin_family = AF_INET;
		sin4[TCPDROP_LOCAL].sin_port = inc->inc_lport;
		memcpy(&sin4[TCPDROP_LOCAL].sin_addr, &inc->inc_laddr,
		    sizeof inc->inc_laddr);
		local = (struct sockaddr *)&sin4[TCPDROP_LOCAL];

		sin4[TCPDROP_FOREIGN].sin_len = sizeof sin4[TCPDROP_FOREIGN];
		sin4[TCPDROP_FOREIGN].sin_family = AF_INET;
		sin4[TCPDROP_FOREIGN].sin_port = inc->inc_fport;
		memcpy(&sin4[TCPDROP_FOREIGN].sin_addr, &inc->inc_faddr,
		    sizeof inc->inc_faddr);
		foreign = (struct sockaddr *)&sin4[TCPDROP_FOREIGN];
	}

	return (tcpdrop(local, foreign));
}

