112-hackathon/src/dns_lkm/nf_conntrack_sdns.c

/* SDNS extension for IP connection tracking
 * (C) 2021 by Beau Kujath <beau@breakpointingbad.com>
 */

#define pr_fmt(fmt) KBUILD_MODNAME ": " fmt

#include <linux/module.h>
#include <linux/moduleparam.h>
#include <linux/skbuff.h>
#include <linux/in.h>
#include <linux/ip.h>
#include <linux/tcp.h>
#include <linux/udp.h>
#include <linux/netfilter.h>
#include <linux/slab.h>

#include <net/netfilter/nf_conntrack.h>
#include <net/netfilter/nf_conntrack_expect.h>
#include <net/netfilter/nf_conntrack_helper.h>
#include <linux/netfilter/nf_conntrack_irc.h>

#define MAX_PORTS 8
static unsigned short ports[MAX_PORTS];
static unsigned int ports_c;


static char *sdns_buffer;


//static DEFINE_SPINLOCK(nf_sdns_lock);



struct dnsinfo {
        __be16 txid;
};



struct nf_ct_sdns_master {
		unsigned int txid;
        /* Valid seq positions for cmd matching after newline */
        //u_int32_t seq_aft_nl[IP_CT_DIR_MAX][NUM_SEQ_TO_REMEMBER];
        /* 0 means seq_match_aft_nl not set */
        //u_int16_t seq_aft_nl_num[IP_CT_DIR_MAX];
        /* pickup sequence tracking, useful for conntrackd */
        //u_int16_t flags[IP_CT_DIR_MAX];
};


unsigned int (*nf_nat_sdns_hook)(struct sk_buff *skb,
				enum ip_conntrack_info ctinfo,
				unsigned int protoff,
				unsigned int matchoff,
				unsigned int matchlen,
				struct nf_conntrack_expect *exp) __read_mostly;
EXPORT_SYMBOL_GPL(nf_nat_sdns_hook);

#define HELPER_NAME "sdns"



MODULE_AUTHOR("Beau Kujath <beau@breakpointingbad.com>");
MODULE_DESCRIPTION("sdns connection tracking helper");
MODULE_LICENSE("GPL");
MODULE_ALIAS("ip_conntrack_sdns");
MODULE_ALIAS_NFCT_HELPER(HELPER_NAME);

module_param_array(ports, ushort, &ports_c, 0400);
MODULE_PARM_DESC(ports, "port numbers of dns servers");





static int help(struct sk_buff *skb, unsigned int protoff,
		struct nf_conn *ct, enum ip_conntrack_info ctinfo)
{


	struct udphdr _udph;
	struct dnsinfo _dnsi;
	const struct dnsinfo *dnsi;
	unsigned int dataoff, datalimit;

	struct nf_ct_sdns_master *ct_sdns_info = nfct_help_data(ct);
	struct udphdr *uh = skb_header_pointer(skb, protoff, sizeof(_udph), &_udph);

	printk(KERN_INFO "\ntriggering conntrack sdns helper function\n");

	
	
	//int dport = ntohs(uh->dest);
	//printk(KERN_INFO "original protoff: (%u), len: (%u)\n", protoff, sizeof(_udph));
	//printk(KERN_INFO "dst port for packet: %u\n", dport);


	if (uh == NULL)
		return NF_ACCEPT;


	dataoff = protoff + sizeof(struct udphdr);
	datalimit = dataoff + sizeof(struct dnsinfo);

	// could probably delete this if check
	if (dataoff >= skb->len) {
		pr_debug("sdns: dataoff(%u) >= skblen(%u)\n", dataoff,
			 skb->len);
		return NF_ACCEPT;
	}


	if (datalimit >= skb->len) {
		pr_debug("sdns: datalimit(%u) >= skblen(%u)\n", datalimit,
			 skb->len);
		return NF_ACCEPT;
	}

	dnsi = skb_header_pointer(skb, dataoff,
                                 sizeof(_dnsi), &_dnsi);

	if (dnsi == NULL)
		return NF_ACCEPT;
 

	//printk(KERN_INFO "ctinfo val: %d\n", ctinfo);

	if (ctinfo == IP_CT_NEW) {
		// keep track of txid if it is new entry
		int txid = ntohs(dnsi->txid);
		ct_sdns_info->txid = txid;
		printk(KERN_INFO "tracking txid for new ct entry: %d\n", ct_sdns_info->txid);

	} else if (ctinfo == IP_CT_IS_REPLY) {

		// check the txid of reply on reply entry
		int txid = ntohs(dnsi->txid);
		printk(KERN_INFO "dns txid of reply: %d, dns txid stored in ct struct: %d \n", txid, ct_sdns_info->txid);

		if (txid == ct_sdns_info->txid) {
			printk(KERN_INFO "yes, txid on reply matched so accept!");
			return NF_ACCEPT;
		} else {
			printk(KERN_INFO "WARNING, txid on reply did not match so removing ct entry!");
			nf_ct_delete(ct, 0, 0);
			return NF_DROP;
		}
	}


	return NF_ACCEPT;
}

static struct nf_conntrack_helper sdns[MAX_PORTS] __read_mostly;
static struct nf_conntrack_expect_policy sdns_exp_policy;

static int __init nf_conntrack_sdns_init(void)
{
	int i, ret;

	printk(KERN_INFO "initializing conntrack sdns helper\n");

	sdns_buffer = kmalloc(65536, GFP_KERNEL);
	if (!sdns_buffer)
		return -ENOMEM;

	/* If no port given, default to standard sdns port */
	if (ports_c == 0) {
		printk(KERN_INFO "no ports given so assigning default port\n");
		ports[ports_c++] = IRC_PORT;
	}

	for (i = 0; i < ports_c; i++) {
		nf_ct_helper_init(&sdns[i], AF_INET, IPPROTO_UDP, HELPER_NAME,
				  IRC_PORT, ports[i], i, &sdns_exp_policy,
				  0, help, NULL, THIS_MODULE);
	}

	ret = nf_conntrack_helpers_register(&sdns[0], ports_c);
	if (ret) {
		pr_err("failed to register helpers\n");
		kfree(sdns_buffer);
		return ret;
	}

	return 0;
}

static void __exit nf_conntrack_sdns_fini(void)
{
	nf_conntrack_helpers_unregister(sdns, ports_c);
	kfree(sdns_buffer);
}

module_init(nf_conntrack_sdns_init);
module_exit(nf_conntrack_sdns_fini);