练习代码 协议族

#include <linux/module.h>
#include <linux/init.h>
#include <linux/kernel.h>
#include <linux/socket.h>
#include <linux/net.h>
#include <net/sock.h>
#include <linux/skbuff.h>

#define err(msg) printk(KERN_ALERT "%s\n", msg)

#define AF_DRIVER 30

static int priv_family_create(struct socket *sock, int protocol);
static int socket_release(struct socket *sock);
static int socket_sendmsg(struct kiocb *iocb, struct socket *sock,
                        struct msghdr *m, size_t total_len);
static void proto_close(struct sock *sk, long timeout);
static int proto_sendmsg(struct kiocb *iocb, struct sock *sk, struct msghdr *msg, size_t len);
static int proto_recvmsg(struct kiocb *iocb, struct sock *sk, struct msghdr *msg,
                        size_t len, int noblock, int flags, int *addr_len);

static struct net_proto_family priv_family =
{
        .family = AF_DRIVER,
        .owner = THIS_MODULE,
        .create = priv_family_create,
};

static struct proto_ops priv_socket_ops =
{
        .family = AF_DRIVER,
        .owner = THIS_MODULE,
        .release = socket_release,
        .sendmsg = socket_sendmsg,
        .recvmsg = sock_common_recvmsg,
};

static struct proto priv_proto_ops =
{
        .name = "privproto",
        .owner = THIS_MODULE,
        .close = proto_close,
        .sendmsg = proto_sendmsg,
        .recvmsg = proto_recvmsg,
        .obj_size = sizeof(struct sock),
};

static struct sk_buff_head buffqueue; /* goatsucker: store data buffer queue */

static int priv_family_create(struct socket *sock, int protocol)
{
        struct sock *sk;
        int err;

        sock->state = SS_UNCONNECTED;
        sock->ops = &priv_socket_ops;
        err = -ENOBUFS;
        sk = sk_alloc(AF_DRIVER, GFP_KERNEL, &priv_proto_ops, 1);
        if (!sk)
        {
                err("sk_alloc");
                goto alloc_err;
        }

        sock_init_data(sock, sk); /* goatsucker: init sk_refcnt = 1 */
        sk->sk_family = AF_DRIVER;
        sk->sk_protocol = protocol;

        return 0;
alloc_err:
        return err;
}

/* goatsucker: socket_release/socket_sendmsg/socket_recvmsg is bsd layer represent */

static int socket_release(struct socket *sock)
{
        struct sock *sk = sock->sk;

        if (sk)
        {
                sock->sk = NULL;
                sk->sk_prot->close(sk, 0);
        }

        return 0;
}

static int socket_sendmsg(struct kiocb *iocb, struct socket *sock, struct msghdr *m, size_t total_len)
{
        struct sock *sk = sock->sk;

        return sk->sk_prot->sendmsg(iocb, sk, m, total_len);
}

/* goatsucker: proto_close/proto_sendmsg/proto_recvmsg is transmit layer represent */

static void priv_destroy_sock(struct sock *sk)
{
        skb_queue_purge(&sk->sk_receive_queue);
        skb_queue_purge(&sk->sk_error_queue);
        skb_queue_purge(&sk->sk_write_queue);
        sock_put(sk);
}

static void proto_close(struct sock *sk, long timeout)
{
        lock_sock(sk);
        sock_hold(sk); /* goatsucker: increment sk->sk_refcnt */
        sock_orphan(sk);
        release_sock(sk);

        local_bh_disable();
        bh_lock_sock(sk);
        priv_destroy_sock(sk);
        bh_unlock_sock(sk);
        local_bh_enable();
        sock_put(sk);
}

static int proto_sendmsg(struct kiocb *iocb, struct sock *sk, struct msghdr *msg, size_t len)
{
        struct sk_buff *skb;

        skb = alloc_skb(len, GFP_KERNEL);
        if (!skb)
        {
                err("alloc_skb");
                return -1;
        }

        if (memcpy_fromiovec(skb_put(skb, len), msg->msg_iov, len) != 0)
        {
                err("memcpy_fromiovec");
                goto memcpy_err;
        }
        skb_queue_tail(&buffqueue, skb);

        return 0;
memcpy_err:
        kfree_skb(skb);
        return -EFAULT;
}

static int proto_recvmsg(struct kiocb *iocb, struct sock *sk, struct msghdr *msg,
                        size_t len, int noblock, int flags, int *addr_len)
{
        struct sk_buff *skb;
        int copied;

        skb = skb_dequeue(&buffqueue);
        if (skb)
        {
                if (memcpy_toiovec(msg->msg_iov, skb->data, skb->len) != 0)
                {
                        err("memcpy_toiovec");
                        kfree_skb(skb);
                        return -EFAULT;
                }

                copied = skb->len;
                kfree_skb(skb);
                return copied;
        }

        return 0;
}

static int privproto_init(void)
{
        skb_queue_head_init(&buffqueue);

        if (proto_register(&priv_proto_ops, 1) != 0)
        {
                err("proto_register");
                return -1;
        }

        if (sock_register(&priv_family) != 0)
        {
                err("sock_register");
                goto sock_err;
        }

        return 0;
sock_err:
        proto_unregister(&priv_proto_ops);
        return -1;
}

static void privproto_exit(void)
{
        sock_unregister(AF_DRIVER);
        proto_unregister(&priv_proto_ops);
        skb_queue_purge(&buffqueue);
}

module_init(privproto_init);
module_exit(privproto_exit);

MODULE_LICENSE("GPL");