网络层实现NAT简单例子程序

#include <asm/atomic.h>
#include <asm/byteorder.h>
#include <asm/checksum.h>
#include <asm/unaligned.h>
#include <linux/module.h>
#include <linux/init.h>
#include <linux/kernel.h>
#include <linux/net.h>
#include <linux/in.h>
#include <linux/types.h>
#include <linux/kthread.h>
#include <linux/wait.h>
#include <linux/skbuff.h>
#include <linux/string.h>
#include <linux/sysctl.h>
#include <linux/netfilter.h>
#include <linux/netfilter_ipv4.h>
#include <linux/ip.h>
#include <linux/workqueue.h>
#include <linux/jiffies.h>
#include <linux/route.h>
#include <linux/if_arp.h>
#include <linux/slab.h>
#include <linux/stddef.h>
#include <linux/mutex.h>
#include <linux/compiler.h>
#include <linux/icmp.h>
#include <linux/jhash.h>
#include <linux/list.h>
#include <linux/spinlock_types.h>
#include <net/net_namespace.h>
#include <net/sock.h>
#include <net/route.h>
#include <net/inet_connection_sock.h>
#include <net/request_sock.h>
#include <net/icmp.h>
#include <net/ip.h>

#include "main.h"

static DEFINE_SPINLOCK(hashtable_lock);
static struct list_head hashtable[HASHSIZE] __read_mostly;
static struct kmem_cache *cachep __read_mostly;

static u32 rand;

static u32 conn_hash(__be16 sport, __be16 dport, u8 which)
{
    if (which == CONNIN)
        return jhash_1word((__force u32)sport, rand) & HASHMASK;

    return jhash_1word((__force u32)dport, rand) & HASHMASK;
}

static struct conn_struct * __conn_get(__be32 saddr, __be32 daddr, __be16 sport, __be16 dport, u8 which)
{
    struct conn_struct *cp = NULL;
    u32 hash;

    hash = conn_hash(sport, dport, which);

    spin_lock_bh(&hashtable_lock);
    list_for_each_entry(cp, &hashtable[hash], c_list) {
        if (which == CONNIN) {
            if (cp->caddr == saddr && cp->vaddr == daddr && cp->cport == sport && cp->vport == dport) {
                cp->which = CONNIN;
                goto out;
            }
        } else {
            if (cp->daddr == saddr && cp->vaddr == daddr && cp->dport == sport && cp->cport == dport) {
                cp->which = CONNOUT;
                goto out;
            }
        }
    }
    cp = NULL;    /* note this, current cp pointer to the last member */
out:
    spin_unlock_bh(&hashtable_lock);
    return cp;
}

static struct conn_struct * conn_schedule(__be32 saddr, __be32 daddr, __be16 sport, __be16 dport)
{
    struct conn_struct *cp;
    u32 hash;

    cp = kmem_cache_zalloc(cachep, GFP_ATOMIC);
    if (cp == NULL)
        goto out;

    INIT_LIST_HEAD(&cp->c_list);
    cp->caddr = saddr;
    cp->vaddr = daddr;
    cp->daddr = 0x130AA8C0;        
    cp->cport = sport;
    cp->vport = dport;
    cp->dport = htons(7777);
    init_timer(&cp->timer);
    cp->which = CONNIN;
    
    hash = conn_hash(sport, dport, CONNIN);    
    spin_lock_bh(&hashtable_lock);
    list_add(&cp->c_list, &hashtable[hash]);
    spin_unlock_bh(&hashtable_lock);
out:
    return cp;
}

static struct conn_struct * conn_get(struct sk_buff *skb)
{
    struct conn_struct *cp = NULL;
    struct iphdr *iph;
    struct tcphdr *tcph;

    iph = ip_hdr(skb);
    if (iph == NULL) {
        err("ip_hdr");
        goto out;
    }
    
    tcph = (struct tcphdr *)(skb_network_header(skb) + ip_hdrlen(skb));
    if (tcph == NULL) {
        err("tcph is NULL.\n");
        goto out;
    }

    cp = __conn_get(iph->saddr, iph->daddr, tcph->source, tcph->dest, CONNIN);
    if (cp == NULL)
        cp = __conn_get(iph->saddr, iph->daddr, tcph->source, tcph->dest, CONNOUT);

    if (cp == NULL && tcph->syn == 1)     /* new syn packet */
        cp = conn_schedule(iph->saddr, iph->daddr, tcph->source, tcph->dest);

    return cp;
out:
    return NULL;
}

static struct rtable * get_output_route(__be32 rdaddr, u8 rtos)
{
    struct rtable *rt = NULL;
    
    struct flowi fl = { .oif = 0, .nl_u = { .ip4_u = { .daddr = rdaddr, .saddr = 0, .tos = rtos, } }, };

    if (ip_route_output_key(&init_net, &rt, &fl)) {
        err("ip_route_output_key");
        goto out;
    }
out:    
    return rt;
}

static void packet_dnat(struct conn_struct *cp, struct sk_buff *skb)
{
    struct rtable *rt;
    struct iphdr *iph = ip_hdr(skb);
    struct tcphdr *tcph = (struct tcphdr *)(skb_network_header(skb) + ip_hdrlen(skb));

    rt = get_output_route(cp->daddr, RT_TOS(iph->tos));
    if (rt == NULL) {
        printk(KERN_INFO "rt is NULL.\n");
        goto out;
    }

    if (!skb_make_writable(skb, iph->ihl * 4 + tcph->doff * 4)) {
        err("skb_make_writable");
        goto out;
    }
    if (skb_cow(skb, rt->u.dst.dev->hard_header_len)) {
        err("skb_cow");
        goto out;
    }

    skb_dst_drop(skb);
    skb_dst_set(skb, &rt->u.dst);

    iph->saddr = cp->vaddr;
    iph->daddr = cp->daddr;
    tcph->dest = cp->dport;

    tcph->check = 0;
    skb->csum = skb_checksum(skb, ip_hdrlen(skb), skb->len - ip_hdrlen(skb), 0);
    tcph->check = csum_tcpudp_magic(iph->saddr, iph->daddr, skb->len - ip_hdrlen(skb), iph->protocol, skb->csum);
    ip_send_check(ip_hdr(skb));

    dst_output(skb);
out:
    return;
}

static void packet_snat(struct conn_struct *cp, struct sk_buff *skb)
{
    struct rtable *rt;
    struct iphdr *iph = ip_hdr(skb);
    struct tcphdr *tcph = (struct tcphdr *)(skb_network_header(skb) + ip_hdrlen(skb));

    rt = get_output_route(cp->caddr, RT_TOS(iph->tos));
    if (rt == NULL) {
        printk(KERN_INFO "rt is NULL.\n");
        goto out;
    }

    if (!skb_make_writable(skb, iph->ihl * 4 + tcph->doff * 4)) {
        err("skb_make_writable");
        goto out;
    }
    if (skb_cow(skb, rt->u.dst.dev->hard_header_len)) {
        err("skb_cow");
        goto out;
    }

    skb_dst_drop(skb);
    skb_dst_set(skb, &rt->u.dst);

    iph->saddr = cp->vaddr;
    iph->daddr = cp->caddr;
    tcph->source = cp->vport;
    tcph->dest = cp->cport;
    
    tcph->check = 0;
    skb->csum = skb_checksum(skb, ip_hdrlen(skb), skb->len - ip_hdrlen(skb), 0);
    tcph->check = csum_tcpudp_magic(iph->saddr, iph->daddr, skb->len - ip_hdrlen(skb), iph->protocol, skb->csum);
    ip_send_check(ip_hdr(skb));

    dst_output(skb);
out:
    return;
}

/* 192.168.10.250 network byteorder 0xFA0AA8C0 */
/* 192.168.10.188 network byteorder 0xBC0AA8C0 */
/* 192.168.10.19 network byteorder 0x130AA8C0 */

static unsigned int sg_vs_in(unsigned int hooknum,
            struct sk_buff *skb,
            const struct net_device *in,
            const struct net_device *out,
            int (*okfn)(struct sk_buff *))
{
    struct conn_struct *cp;
    struct tcphdr *tcph = (struct tcphdr *)(skb_network_header(skb) + ip_hdrlen(skb));

    if (tcph->dest != htons(8888) && tcph->source != htons(7777))
        goto accept;

    if ((skb = skb_share_check(skb, GFP_ATOMIC)) == NULL) {
        err("skb_share_check");
        goto out;
    }

    if (!pskb_may_pull(skb, sizeof(struct iphdr) + sizeof(struct tcphdr))) {
        err("pskb_may_pull");
        goto out;
    }

    cp = conn_get(skb);    
    if (cp == NULL)
        goto accept;

    if (cp->which == CONNIN)
        packet_dnat(cp, skb);
    else
        packet_snat(cp, skb);
    
    return NF_STOLEN;
accept:
    return NF_ACCEPT;
out:
    return NF_DROP;
}

static struct nf_hook_ops sg_nf_ops = {
    .hook        = sg_vs_in,    
    .owner        = THIS_MODULE,
    .pf        = PF_INET,
    .hooknum = NF_INET_LOCAL_IN,        
};

static void conn_hashtable_init(void)
{
    int i;

    for (i = 0; i < HASHSIZE; i++)
        INIT_LIST_HEAD(&hashtable[i]);
}

static int main_init(void)
{
    int rc;

    get_random_bytes(&rand, sizeof(rand));
    conn_hashtable_init();

    cachep = kmem_cache_create("ipnat", sizeof(struct conn_struct), 0, SLAB_HWCACHE_ALIGN, NULL);
    if (cachep == NULL) {
        err("kmem_cache_create");
        goto out;
    }

    rc = nf_register_hook(&sg_nf_ops);
    if (rc < 0) {
        err("nf_register_hook");
        goto out;
    }
    
    printk(KERN_INFO "NAT module register.\n");
    return 0;
out:
    return -ENOMEM;
}

static void conn_cleanup(void)
{
    struct conn_struct *cp, *nxt;
    int i;

    spin_lock_bh(&hashtable_lock);
        
    for (i = 0; i < HASHSIZE; i++) {
        list_for_each_entry_safe(cp, nxt, &hashtable[i], c_list) {
            list_del(&cp->c_list);
            del_timer_sync(&cp->timer);
            kmem_cache_free(cachep, cp);
        }
    }

    spin_unlock_bh(&hashtable_lock);
}

static void main_exit(void)
{
    nf_unregister_hook(&sg_nf_ops);
    conn_cleanup();        
    kmem_cache_destroy(cachep);

    printk(KERN_INFO "NAT module unregister.\n");
}

module_init(main_init);
module_exit(main_exit);
MODULE_LICENSE("GPL");