#include <linux/module.h>
#include <linux/kernel.h>
#include <linux/netfilter.h>
#include <linux/netfilter_ipv4.h>
#include <linux/ip.h>
#include <linux/tcp.h>
#include <linux/udp.h>
#include <linux/skbuff.h>
#include <net/netfilter/nf_conntrack.h>

MODULE_LICENSE("GPL");
MODULE_AUTHOR("Test Inc.");
MODULE_DESCRIPTION("Netfilter post routing hook");
MODULE_VERSION("1.0");

// Maximum number of skbs to keep in the list
#define MAX_SKB_LIST_SIZE 2000

static atomic_t wipe_flag = ATOMIC_INIT(0);

// Structure to hold skb copy in RCU list
struct skb_list_entry {
    struct list_head list;
    struct rcu_head rcu;
    struct sk_buff *skb_copy;
};

// RCU protected list head and control variables
static LIST_HEAD(skb_list);
static DEFINE_SPINLOCK(skb_list_lock);
static atomic_t skb_list_size = ATOMIC_INIT(0);
static atomic_t skb_copies_added = ATOMIC_INIT(0);

// RCU callback for freeing skb list entries
static void free_skb_entry_rcu(struct rcu_head *head)
{
    struct skb_list_entry *entry = container_of(head, struct skb_list_entry, rcu);

    if (entry->skb_copy)
        kfree_skb(entry->skb_copy);
    kfree(entry);
    atomic_dec(&skb_list_size);
}

// Function to add skb copy to RCU list
static int add_skb_to_list(struct sk_buff *skb)
{
    struct skb_list_entry *entry;
    struct skb_list_entry *oldest_entry;

    // Check if we need to remove old entries
    if (atomic_read(&skb_list_size) >= MAX_SKB_LIST_SIZE) {
        spin_lock_bh(&skb_list_lock);
        if (!list_empty(&skb_list)) {
            oldest_entry = list_first_entry(&skb_list, struct skb_list_entry, list);
            list_del_rcu(&oldest_entry->list);
            spin_unlock_bh(&skb_list_lock);
            call_rcu(&oldest_entry->rcu, free_skb_entry_rcu);
        } else {
            spin_unlock_bh(&skb_list_lock);
        }
    }

    // Allocate new entry
    entry = kmalloc(sizeof(struct skb_list_entry), GFP_ATOMIC);
    if (!entry) {
        printk(KERN_ERR "POST_ROUTING: Failed to allocate memory for skb entry\n");
        return -ENOMEM;
    }

    // Create skb copy
    entry->skb_copy = skb_copy(skb, GFP_ATOMIC);

    if (!entry->skb_copy) {
        kfree(entry);
        printk(KERN_ERR "POST_ROUTING: Failed to copy skb\n");
        return -ENOMEM;
    }

    // Initialize entry
    INIT_LIST_HEAD(&entry->list);

    // Add to RCU list
    spin_lock_bh(&skb_list_lock);
    list_add_tail_rcu(&entry->list, &skb_list);
    spin_unlock_bh(&skb_list_lock);

    atomic_inc(&skb_list_size);

    return 0;
}

// Function to cleanup all entries in the list
static void cleanup_skb_list(void)
{
    struct skb_list_entry *entry, *tmp;

    spin_lock_bh(&skb_list_lock);
    list_for_each_entry_safe(entry, tmp, &skb_list, list) {
        list_del_rcu(&entry->list);
        call_rcu(&entry->rcu, free_skb_entry_rcu);
    }
    spin_unlock_bh(&skb_list_lock);
}

/*
 * Wipe conntrack off of skb and delete it
 */
static inline void wipe_conntrack(struct sk_buff *skb)
{
    enum ip_conntrack_info ctinfo;
    struct nf_conn *ct = nf_ct_get(skb, &ctinfo);

    if (ct) {
        printk("nf_ct_is_confirmed(): %d, ct: %p\n",
                nf_ct_is_confirmed(ct), ct);
        if (nf_ct_is_confirmed(ct))
            nf_ct_delete(ct, 0, 0);
        nf_reset_ct(skb);
    }
}

static void wipe_conntrack_from_skb_list(void)
{
    struct skb_list_entry *entry;

    int old = atomic_cmpxchg(&wipe_flag, 0, 1);

    if (old != 0)
        return;

    rcu_read_lock();
    list_for_each_entry_rcu(entry, &skb_list, list) {
        if (entry->skb_copy)
            wipe_conntrack(entry->skb_copy);
    }
    rcu_read_unlock();

    atomic_set(&wipe_flag, 0);
}

// Hook function prototype
static unsigned int my_hook(void *priv, struct sk_buff *skb, const struct nf_hook_state *state);

// Netfilter hook operations structure
static struct nf_hook_ops post_routing_ops = {
    .hook = my_hook,
    .hooknum = NF_INET_POST_ROUTING,
    .pf = NFPROTO_IPV4,
    .priority = NF_IP_PRI_FIRST,
};

#define WIPE_THRESHOLD 500

// Hook function implementation
static unsigned int my_hook(void *priv, struct sk_buff *skb, const struct nf_hook_state *state)
{
    // Add skb copy to RCU list
    if (add_skb_to_list(skb) != 0)
        printk(KERN_ERR "POST_ROUTING: Failed to add skb to list\n");

    if (atomic_inc_return(&skb_copies_added) % WIPE_THRESHOLD == 0) {
        printk("%d skb copies added, we'll wipe their conntracks\n", WIPE_THRESHOLD);
        wipe_conntrack_from_skb_list();
    }

    // Accept the packet (let it continue)
    return NF_ACCEPT;
}

// Module initialization function
static int __init netfilter_post_routing_init(void)
{
    int ret;
    
    printk(KERN_INFO "Netfilter post routing hook module loaded\n");
    
    // Register the hook
    ret = nf_register_net_hook(&init_net, &post_routing_ops);
    if (ret < 0) {
        printk(KERN_ERR "Failed to register netfilter hook: %d\n", ret);
        return ret;
    }
    
    printk(KERN_INFO "Post routing hook registered successfully\n");
    return 0;
}

// Module cleanup function
static void __exit netfilter_post_routing_exit(void)
{
    // Unregister the hook
    nf_unregister_net_hook(&init_net, &post_routing_ops);

    // Clean up all entries in the RCU list
    cleanup_skb_list();
    
    // Wait for all RCU callbacks to complete
    rcu_barrier();

    printk(KERN_INFO "Netfilter post routing hook module unloaded\n");
}

// Register module entry and exit points
module_init(netfilter_post_routing_init);
module_exit(netfilter_post_routing_exit);