my_udp.c

一个基于linux的TCP/IP协议栈的实现

DEFINE_SNMP_STAT(struct udp_mib, myudp_statistics) __read_mostly;
struct hlist_head myudp_hash[UDP_HTABLE_SIZE];
DEFINE_RWLOCK(myudp_hash_lock);
#define MYUDP_INC_STATS_USER(field) 	SNMP_INC_STATS_USER(myudp_statistics, field)
extern struct icmp_mib *myicmp_statistics[2];
#define MYICMP_INC_STATS_BH(field)		SNMP_INC_STATS_BH(myicmp_statistics, field)
#define MYUDP_INC_STATS_BH(field)			SNMP_INC_STATS_BH(myudp_statistics, field)
static void myudp_flush_pending_frames(struct sock *sk)
{
	struct udp_sock *up = udp_sk(sk);

	if (up->pending) {
		up->len = 0;
		up->pending = 0;
		myip_flush_pending_frames(sk);
	}
}

static int myudp_push_pending_frames(struct sock *sk, struct udp_sock *up)
{
	struct inet_sock *inet = inet_sk(sk);
	struct flowi *fl = &inet->cork.fl;
	struct sk_buff *skb;
	struct udphdr *uh;
	int err = 0;

	if ((skb = skb_peek(&sk->sk_write_queue)) == NULL)
		goto out;

	uh = skb->h.uh;
	uh->source = fl->fl_ip_sport;
	uh->dest = fl->fl_ip_dport;
	uh->len = htons(up->len);
	uh->check = 0;

	if( sk->sk_no_check == UDP_CSUM_NOXMIT ){
		skb->ip_summed = CHECKSUM_NONE;
		goto send;
	}
	int queue_len;
	if( (queue_len = skb_queue_len(&sk->sk_write_queue)) == 1 ){
		if( skb->ip_summed == CHECKSUM_HW ){
			skb->csum = offsetof(struct udphdr, check);
			uh->check = ~csum_tcpudp_magic(fl->fl4_src, fl->fl4_dst,
					up->len, IPPROTO_UDP, 0);
		}else{
			skb->csum = csum_partial((char *)uh,
							sizeof(struct udphdr), skb->csum);
			uh->check = csum_tcpudp_magic(fl->fl4_src, fl->fl4_dst,
					up->len, IPPROTO_UDP, skb->csum);
			if (uh->check == 0)
				uh->check = -1;
		}
	}else{
		unsigned int csum = 0;
		if (skb->ip_summed == CHECKSUM_HW) {
			int offset = (unsigned char *)uh - skb->data;
			skb->csum = skb_checksum(skb, offset, skb->len - offset, 0);

			skb->ip_summed = CHECKSUM_NONE;
		}else{
			skb->csum = csum_partial((char *)uh,
					sizeof(struct udphdr), skb->csum);
		}

		skb_queue_walk(&sk->sk_write_queue, skb) {
			csum = csum_add(csum, skb->csum);
		}
		uh->check = csum_tcpudp_magic(fl->fl4_src, fl->fl4_dst,
				up->len, IPPROTO_UDP, csum);
		if (uh->check == 0)
			uh->check = -1;
	}
send:
	err = myip_push_pending_frames(sk);
out:
	up->len = 0;
	up->pending = 0;
	return err;
}

int myudp_disconnect(struct sock *sk, int flags)
{
	struct inet_sock *inet = inet_sk(sk);
	 
	sk->sk_state = TCP_CLOSE;
	inet->daddr = 0;
	inet->dport = 0;
	sk->sk_bound_dev_if = 0;
	if (!(sk->sk_userlocks & SOCK_BINDADDR_LOCK))
		inet_reset_saddr(sk);

	if (!(sk->sk_userlocks & SOCK_BINDPORT_LOCK)) {
		sk->sk_prot->unhash(sk);
		inet->sport = 0;
	}
	sk_dst_reset(sk);
	return 0;
}

static void myudp_close(struct sock *sk, long timeout)
{
	sk_common_release(sk);
}

int myudp_ioctl(struct sock *sk, int cmd, unsigned long arg)
{
	return 0;
}

unsigned int myudp_poll(struct file *file, struct socket *sock, poll_table *wait)
{
	return 0;
}

static int myudp_destroy_sock(struct sock *sk)
{
	lock_sock(sk);
	myudp_flush_pending_frames(sk);
	release_sock(sk);
	return 0;

}

static int myudp_setsockopt(struct sock *sk, int level, int optname, 
							char __user *optval, int optlen)
{
	struct udp_sock *up = udp_sk(sk);
	int val;
	int err = 0;

	if( level != SOL_UDP )
		return myip_setsockopt(sk, level, optname, optval, optlen);

	//tmp code, FIXME!!
	return -ENOPROTOOPT;
}

static int myudp_getsockopt(struct sock *sk, int level, int optname, 
							char __user *optval, int __user *optlen)
{
	return 0;
}

int myudp_sendmsg(struct kiocb *iocb, struct sock *sk, struct msghdr *msg,
						size_t len)
{
	struct inet_sock *inet = inet_sk(sk);
	struct udp_sock *up = udp_sk(sk);
	int ulen = len;
	struct ipcm_cookie ipc;
	struct rtable *rt = NULL;
	int free = 0;
	int connected = 0;
	u32 daddr, faddr, saddr;
	u16 dport;
	u8  tos;
	int err;
	int corkreq = up->corkflag || msg->msg_flags&MSG_MORE;

	if( len > 0xFFFF )
		return -EMSGSIZE;
	if( msg->msg_flags & MSG_OOB )
		return -EOPNOTSUPP;

	ipc.opt = NULL;

	if (up->pending) {
		lock_sock(sk);
		if( likely(up->pending) ){
			if (unlikely(up->pending != MY_AF_INET)) {
				release_sock(sk);
				return -EINVAL;
			}
 			goto do_append_data;
		}
		release_sock(sk);
	}
	ulen += sizeof(struct udphdr);

	if( msg->msg_name ){
		struct sockaddr_in * usin = (struct sockaddr_in*)msg->msg_name;
		if (msg->msg_namelen < sizeof(*usin))
			return -EINVAL;
		if (usin->sin_family != MY_AF_INET) {
			if (usin->sin_family != AF_UNSPEC)
				return -EAFNOSUPPORT;
		}

		daddr = usin->sin_addr.s_addr;
		dport = usin->sin_port;
		if (dport == 0)
			return -EINVAL;
	}else{
		if (sk->sk_state != TCP_ESTABLISHED)
			return -EDESTADDRREQ;
		daddr = inet->daddr;
		dport = inet->dport;
		connected = 1;
  	}

	ipc.addr = inet->saddr;
	ipc.oif = sk->sk_bound_dev_if;
	if (msg->msg_controllen) {
		err = myip_cmsg_send(msg, &ipc);
		if (err)
			return err;
		if (ipc.opt)
			free = 1;
		connected = 0;
	}
	if (!ipc.opt)
		ipc.opt = inet->opt;

	saddr = ipc.addr;
	ipc.addr = faddr = daddr;

	if (ipc.opt && ipc.opt->srr) {
		if (!daddr)
			return -EINVAL;
		faddr = ipc.opt->faddr;
		connected = 0;
	}
	tos = RT_TOS(inet->tos);
	if (sock_flag(sk, SOCK_LOCALROUTE) || (msg->msg_flags & MSG_DONTROUTE) || 
					(ipc.opt && ipc.opt->is_strictroute)) {
		tos |= RTO_ONLINK;
		connected = 0;
	}
	if (MULTICAST(daddr)) {
		if (!ipc.oif)
			ipc.oif = inet->mc_index;
		if (!saddr)
			saddr = inet->mc_addr;
		connected = 0;
	}

	if( connected )
		rt = (struct rtable*)sk_dst_check(sk, 0);

	if (rt == NULL) {
		struct flowi fl = { .oif = ipc.oif,
				.nl_u = { .ip4_u =
						{ .daddr = faddr,
								.saddr = saddr,
								.tos = tos } },
				.proto = MY_IPPROTO_UDP,
				.uli_u = { .ports =
						{ .sport = inet->sport,
								.dport = dport } } };
		err = myip_route_output_flow( &rt, &fl, sk, !(msg->msg_flags&MSG_DONTWAIT) );
		if (err)
			goto out;
		err = -EACCES;
		if( (rt->rt_flags & RTCF_BROADCAST) && !sock_flag(sk, SOCK_BROADCAST) )
			goto out;
		if( connected )
			sk_dst_set( sk, dst_clone(&rt->u.dst) );
	}
	if( msg->msg_flags & MSG_CONFIRM )
		goto do_confirm;

back_from_confirm:
	saddr = rt->rt_src;
	if (!ipc.addr)
		daddr = ipc.addr = rt->rt_dst;
	lock_sock(sk);
	if( unlikely(up->pending) ){
		release_sock(sk);
		LIMIT_NETDEBUG(KERN_DEBUG "udp cork app bug 2\n");
		err = -EINVAL;
		goto out;
	}

	inet->cork.fl.fl4_dst = daddr;
	inet->cork.fl.fl_ip_dport = dport;
	inet->cork.fl.fl4_src = saddr;
	inet->cork.fl.fl_ip_sport = inet->sport;
	up->pending = MY_AF_INET;

do_append_data:
	up->len += ulen;
	err = myip_append_data(sk, myip_generic_getfrag, msg->msg_iov, ulen, 
					sizeof(struct udphdr), &ipc, rt, 
					corkreq ? msg->msg_flags|MSG_MORE : msg->msg_flags);
	if (err)
		myudp_flush_pending_frames(sk);
	else if (!corkreq)
		err = myudp_push_pending_frames(sk, up);
	release_sock(sk);

out:
	ip_rt_put(rt);
	if (free)
		kfree(ipc.opt);
	if (!err) {
		MYUDP_INC_STATS_USER(UDP_MIB_OUTDATAGRAMS);
		return len;
	}
	return err;

do_confirm:
	dst_confirm(&rt->u.dst);
	if (!(msg->msg_flags&MSG_PROBE) || len)
		goto back_from_confirm;
	err = 0;
	goto out;
}

static __inline__ int __myudp_checksum_complete(struct sk_buff *skb)
{
	return __skb_checksum_complete(skb);
}

static __inline__ int myudp_checksum_complete(struct sk_buff *skb)
{
	return skb->ip_summed != CHECKSUM_UNNECESSARY &&
		__myudp_checksum_complete(skb);
}

static int myudp_recvmsg(struct kiocb *iocb, struct sock *sk, struct msghdr *msg,
						       size_t len, int noblock, int flags, int *addr_len)
{
	struct inet_sock *inet = inet_sk(sk);
  	struct sockaddr_in *sin = (struct sockaddr_in *)msg->msg_name;
  	struct sk_buff *skb;
  	int copied, err;

	if( addr_len )
		*addr_len=sizeof(*sin);

	if( flags & MSG_ERRQUEUE )
		return myip_recv_error(sk, msg, len);

try_again:
	skb = skb_recv_datagram(sk, flags, noblock, &err);
	if (!skb)
		goto out;

  	copied = skb->len - sizeof(struct udphdr);
	if (copied > len) {
		copied = len;
		msg->msg_flags |= MSG_TRUNC;
	}

	if( skb->ip_summed == CHECKSUM_UNNECESSARY ){
		err = skb_copy_datagram_iovec(skb, sizeof(struct udphdr), msg->msg_iov,
						copied);
	}else if( msg->msg_flags & MSG_TRUNC ){
		if (__myudp_checksum_complete(skb))
			goto csum_copy_err;
		err = skb_copy_datagram_iovec(skb, sizeof(struct udphdr), msg->msg_iov,
						copied);
	} else {
		err = skb_copy_and_csum_datagram_iovec(skb, sizeof(struct udphdr), msg->msg_iov);
		if (err == -EINVAL)
			goto csum_copy_err;
	}

		if (err)
		goto out_free;

	sock_recv_timestamp(msg, sk, skb);

	if (sin){
		sin->sin_family = MY_AF_INET;
		sin->sin_port = skb->h.uh->source;
		sin->sin_addr.s_addr = skb->nh.iph->saddr;
		memset(sin->sin_zero, 0, sizeof(sin->sin_zero));
  	}
	if (inet->cmsg_flags)
		myip_cmsg_recv(msg, skb);
	err = copied;
	if (flags & MSG_TRUNC)
		err = skb->len - sizeof(struct udphdr);
  
out_free:
  	skb_free_datagram(sk, skb);
out:
  	return err;

csum_copy_err:
	MYUDP_INC_STATS_BH(UDP_MIB_INERRORS);

	skb_kill_datagram(sk, skb, flags);

	if (noblock)
		return -EAGAIN;	
	goto try_again;
}

static int myudp_sendpage(struct sock *sk, struct page *page, int offset,
							size_t size, int flags)
{
	return 0;
}

static int myudp_queue_rcv_skb(struct sock * sk, struct sk_buff *skb)
{
	struct udp_sock *up = udp_sk(sk);

	//if (!xfrm4_policy_check(sk, XFRM_POLICY_IN, skb)) {
	//	kfree_skb(skb);
	//	return -1;
	//}
	nf_reset(skb);