Added features that will be required for PXE UDP support.

Introduced struct sockaddr_tcpip, to simplify code that deals with
both IPv4 and IPv6 addresses.

Reorganised parts of tcpip.c and udp.c.

6 files changed, 294 insertions, 387 deletions

diff --git a/src/net/ipv4.c b/src/net/ipv4.c
index 147d267..d1c770a 100644
--- a/src/net/ipv4.c
+++ b/src/net/ipv4.c
@@ -128,7 +128,8 @@ static void ipv4_dump ( struct iphdr *iphdr __unused ) {
  * @v timer	Retry timer
  * @v over	If asserted, the timer is greater than @c MAX_TIMEOUT 
  */
-void ipv4_frag_expired ( struct retry_timer *timer __unused , int over ) {
+static void ipv4_frag_expired ( struct retry_timer *timer __unused,
+				int over ) {
 	if ( over ) {
 		DBG ( "Fragment reassembly timeout" );
 		/* Free the fragment buffer */
@@ -140,7 +141,7 @@ void ipv4_frag_expired ( struct retry_timer *timer __unused , int over ) {
  *
  * @v fragbug	Fragment buffer
  */
-void free_fragbuf ( struct frag_buffer *fragbuf ) {
+static void free_fragbuf ( struct frag_buffer *fragbuf ) {
 	if ( fragbuf ) {
 		free_dma ( fragbuf, sizeof ( *fragbuf ) );
 	}
@@ -152,7 +153,7 @@ void free_fragbuf ( struct frag_buffer *fragbuf ) {
  * @v pkb		Packet buffer, fragment of the datagram
  * @ret frag_pkb	Reassembled packet, or NULL
  */
-struct pk_buff * ipv4_reassemble ( struct pk_buff * pkb ) {
+static struct pk_buff * ipv4_reassemble ( struct pk_buff * pkb ) {
 	struct iphdr *iphdr = pkb->data;
 	struct frag_buffer *fragbuf;
 	
@@ -234,8 +235,8 @@ struct pk_buff * ipv4_reassemble ( struct pk_buff * pkb ) {
  *
  * This function calculates the tcpip 
  */
-void ipv4_tx_csum ( struct pk_buff *pkb, struct tcpip_protocol *tcpip ) {
-
+static void ipv4_tx_csum ( struct pk_buff *pkb,
+			   struct tcpip_protocol *tcpip ) {
 	struct iphdr *iphdr = pkb->data;
 	struct ipv4_pseudo_header pshdr;
 	uint16_t *csum = ( ( ( void * ) iphdr ) + sizeof ( *iphdr )
@@ -255,8 +256,8 @@ void ipv4_tx_csum ( struct pk_buff *pkb, struct tcpip_protocol *tcpip ) {
 /**
  * Calculate the transport-layer checksum while processing packets
  */
-uint16_t ipv4_rx_csum ( struct pk_buff *pkb __unused,
-			uint8_t trans_proto __unused ) {
+static uint16_t ipv4_rx_csum ( struct pk_buff *pkb __unused,
+			       uint8_t trans_proto __unused ) {
 	/** 
 	 * This function needs to be implemented. Until then, it will return
 	 * 0xffffffff every time
@@ -265,91 +266,20 @@ uint16_t ipv4_rx_csum ( struct pk_buff *pkb __unused,
 }
 
 /**
- * Transmit packet constructed by uIP
- *
- * @v pkb		Packet buffer
- * @ret rc		Return status code
- *
- */
-int ipv4_uip_tx ( struct pk_buff *pkb ) {
-	struct iphdr *iphdr = pkb->data;
-	struct ipv4_miniroute *miniroute;
-	struct net_device *netdev = NULL;
-	struct in_addr next_hop;
-	struct in_addr source;
-	uint8_t ll_dest_buf[MAX_LL_ADDR_LEN];
-	const uint8_t *ll_dest = ll_dest_buf;
-	int rc;
-
-	/* Use routing table to identify next hop and transmitting netdev */
-	next_hop = iphdr->dest;
-	list_for_each_entry ( miniroute, &miniroutes, list ) {
-		if ( ( ( ( iphdr->dest.s_addr ^ miniroute->address.s_addr ) &
-			 miniroute->netmask.s_addr ) == 0 ) ||
-		     ( miniroute->gateway.s_addr != INADDR_NONE ) ) {
-			netdev = miniroute->netdev;
-			source = miniroute->address;
-			if ( miniroute->gateway.s_addr != INADDR_NONE )
-				next_hop = miniroute->gateway;
-			break;
-		}
-	}
-
-	/* Abort if no network device identified */
-	if ( ! netdev ) {
-		DBG ( "No route to %s\n", inet_ntoa ( iphdr->dest ) );
-		rc = -EHOSTUNREACH;
-		goto err;
-	}
-
-	/* Determine link-layer destination address */
-	if ( next_hop.s_addr == INADDR_BROADCAST ) {
-		/* Broadcast address */
-		ll_dest = netdev->ll_protocol->ll_broadcast;
-	} else if ( IN_MULTICAST ( next_hop.s_addr ) ) {
-		/* Special case: IPv4 multicast over Ethernet.	This
-		 * code may need to be generalised once we find out
-		 * what happens for other link layers.
-		 */
-		uint8_t *next_hop_bytes = ( uint8_t * ) &next_hop;
-		ll_dest_buf[0] = 0x01;
-		ll_dest_buf[0] = 0x00;
-		ll_dest_buf[0] = 0x5e;
-		ll_dest_buf[3] = next_hop_bytes[1] & 0x7f;
-		ll_dest_buf[4] = next_hop_bytes[2];
-		ll_dest_buf[5] = next_hop_bytes[3];
-	} else {
-		/* Unicast address: resolve via ARP */
-		if ( ( rc = arp_resolve ( netdev, &ipv4_protocol, &next_hop,
-					  &source, ll_dest_buf ) ) != 0 ) {
-			DBG ( "No ARP entry for %s\n",
-			      inet_ntoa ( iphdr->dest ) );
-			goto err;
-		}
-	}
-	
-	/* Hand off to link layer */
-	return net_tx ( pkb, netdev, &ipv4_protocol, ll_dest );
-
- err:
-	free_pkb ( pkb );
-	return rc;
-}
-
-/**
- * Transmit IP packet (without uIP)
+ * Transmit IP packet
  *
  * @v pkb		Packet buffer
  * @v tcpip		Transport-layer protocol
- * @v dest		Destination network-layer address
+ * @v st_dest		Destination network-layer address
  * @ret rc		Status
  *
  * This function expects a transport-layer segment and prepends the IP header
  */
-int ipv4_tx ( struct pk_buff *pkb, struct tcpip_protocol *tcpip,
-	      struct sockaddr* sock ) {
-	struct in_addr *dest = &sock->sin.sin_addr;
+static int ipv4_tx ( struct pk_buff *pkb,
+		     struct tcpip_protocol *tcpip_protocol,
+		     struct sockaddr_tcpip *st_dest ) {
 	struct iphdr *iphdr = pkb_push ( pkb, sizeof ( *iphdr ) );
+	struct sockaddr_in *sin_dest = ( ( struct sockaddr_in * ) st_dest );
 	struct ipv4_miniroute *miniroute;
 	struct net_device *netdev = NULL;
 	struct in_addr next_hop;
@@ -364,10 +294,10 @@ int ipv4_tx ( struct pk_buff *pkb, struct tcpip_protocol *tcpip,
 	iphdr->ident = htons ( next_ident++ );
 	iphdr->frags = 0;
 	iphdr->ttl = IP_TTL;
-	iphdr->protocol = tcpip->trans_proto;
+	iphdr->protocol = tcpip_protocol->tcpip_proto;
 
 	/* Copy destination address */
-	iphdr->dest = *dest;
+	iphdr->dest = sin_dest->sin_addr;
 
 	/**
 	 * All fields in the IP header filled in except the source network
@@ -375,8 +305,6 @@ int ipv4_tx ( struct pk_buff *pkb, struct tcpip_protocol *tcpip,
 	 * requires the source network address). As the pseudo header requires
 	 * the source address as well and the transport-layer checksum is
 	 * updated after routing.
-	 *
-	 * Continue processing as in ipv4_uip_tx()
 	 */
 
 	/* Use routing table to identify next hop and transmitting netdev */
@@ -400,8 +328,8 @@ int ipv4_tx ( struct pk_buff *pkb, struct tcpip_protocol *tcpip,
 	}
 
 	/* Calculate the transport layer checksum */
-	if ( tcpip->csum_offset > 0 ) {
-		ipv4_tx_csum ( pkb, tcpip );
+	if ( tcpip_protocol->csum_offset > 0 ) {
+		ipv4_tx_csum ( pkb, tcpip_protocol );
 	}
 
 	/* Calculate header checksum, in network byte order */
@@ -446,42 +374,7 @@ int ipv4_tx ( struct pk_buff *pkb, struct tcpip_protocol *tcpip,
 }
 
 /**
- * Process incoming IP packets
- *
- * @v pkb		Packet buffer
- * @v netdev		Network device
- * @v ll_source		Link-layer source address
- * @ret rc		Return status code
- *
- * This handles IP packets by handing them off to the uIP protocol
- * stack.
- */
-static int ipv4_uip_rx ( struct pk_buff *pkb,
-			 struct net_device *netdev __unused,
-			 const void *ll_source __unused ) {
-
-	/* Transfer to uIP buffer.  Horrendously space-inefficient,
-	 * but will do as a proof-of-concept for now.
-	 */
-	uip_len = pkb_len ( pkb );
-	memcpy ( uip_buf, pkb->data, uip_len );
-	free_pkb ( pkb );
-
-	/* Hand to uIP for processing */
-	uip_input ();
-	if ( uip_len > 0 ) {
-		pkb = alloc_pkb ( MAX_LL_HEADER_LEN + uip_len );
-		if ( ! pkb )
-			return -ENOMEM;
-		pkb_reserve ( pkb, MAX_LL_HEADER_LEN );
-		memcpy ( pkb_put ( pkb, uip_len ), uip_buf, uip_len );
-		ipv4_uip_tx ( pkb );
-	}
-	return 0;
-}
-
-/**
- * Process incoming packets (without uIP)
+ * Process incoming packets
  *
  * @v pkb	Packet buffer
  * @v netdev	Network device
@@ -490,18 +383,20 @@ static int ipv4_uip_rx ( struct pk_buff *pkb,
  * This function expects an IP4 network datagram. It processes the headers 
  * and sends it to the transport layer.
  */
-void ipv4_rx ( struct pk_buff *pkb, struct net_device *netdev __unused,
-			const void *ll_source __unused ) {
+static int ipv4_rx ( struct pk_buff *pkb, struct net_device *netdev __unused,
+		     const void *ll_source __unused ) {
 	struct iphdr *iphdr = pkb->data;
-	struct in_addr *src = &iphdr->src;
-	struct in_addr *dest = &iphdr->dest;
+	union {
+		struct sockaddr_in sin;
+		struct sockaddr_tcpip st;
+	} src, dest;
 	uint16_t chksum;
 
 	/* Sanity check */
 	if ( pkb_len ( pkb ) < sizeof ( *iphdr ) ) {
 		DBG ( "IP datagram too short (%d bytes)\n",
 			pkb_len ( pkb ) );
-		return;
+		return -EINVAL;
 	}
 
 	/* Print IP4 header for debugging */
@@ -510,14 +405,14 @@ void ipv4_rx ( struct pk_buff *pkb, struct net_device *netdev __unused,
 	/* Validate version and header length */
 	if ( iphdr->verhdrlen != 0x45 ) {
 		DBG ( "Bad version and header length %x\n", iphdr->verhdrlen );
-		return;
+		return -EINVAL;
 	}
 
 	/* Validate length of IP packet */
-	if ( ntohs ( iphdr->len ) != pkb_len ( pkb ) ) {
+	if ( ntohs ( iphdr->len ) > pkb_len ( pkb ) ) {
 		DBG ( "Inconsistent packet length %d\n",
-					ntohs ( iphdr->len ) );
-		return;
+		      ntohs ( iphdr->len ) );
+		return -EINVAL;
 	}
 
 	/* Verify the checksum */
@@ -533,7 +428,7 @@ void ipv4_rx ( struct pk_buff *pkb, struct net_device *netdev __unused,
 		 */
 		pkb = ipv4_reassemble ( pkb );
 		if ( !pkb ) {
-			return;
+			return 0;
 		}
 	}
 
@@ -543,11 +438,19 @@ void ipv4_rx ( struct pk_buff *pkb, struct net_device *netdev __unused,
 	 * 3. Check the service field
 	 */
 
+	/* Construct socket addresses */
+	memset ( &src, 0, sizeof ( src ) );
+	src.sin.sin_family = AF_INET;
+	src.sin.sin_addr = iphdr->src;
+	memset ( &dest, 0, sizeof ( dest ) );
+	dest.sin.sin_family = AF_INET;
+	dest.sin.sin_addr = iphdr->dest;
+
 	/* Strip header */
 	pkb_pull ( pkb, sizeof ( *iphdr ) );
 
 	/* Send it to the transport layer */
-	tcpip_rx ( pkb, iphdr->protocol, src, dest );
+	return tcpip_rx ( pkb, iphdr->protocol, &src.st, &dest.st );
 }
 
 /** 
@@ -601,11 +504,7 @@ struct net_protocol ipv4_protocol = {
 	.name = "IP",
 	.net_proto = htons ( ETH_P_IP ),
 	.net_addr_len = sizeof ( struct in_addr ),
-#if USE_UIP
-	.rx = ipv4_uip_rx,
-#else
 	.rx = ipv4_rx,
-#endif
 	.ntoa = ipv4_ntoa,
 };
 
@@ -613,10 +512,9 @@ NET_PROTOCOL ( ipv4_protocol );
 
 /** IPv4 TCPIP net protocol */
 struct tcpip_net_protocol ipv4_tcpip_protocol = {
-	.net_protocol = &ipv4_protocol,
+	.name = "IPv4",
 	.sa_family = AF_INET,
 	.tx = ipv4_tx,
-	.tx_csum = ipv4_tx_csum,
 };
 
 TCPIP_NET_PROTOCOL ( ipv4_tcpip_protocol );
diff --git a/src/net/ipv6.c b/src/net/ipv6.c
index 9825a61..a38ec21 100644
--- a/src/net/ipv6.c
+++ b/src/net/ipv6.c
@@ -12,8 +12,9 @@
 /**
  * Transmit IP6 packets
  */
-int ipv6_tx ( struct pk_buff *pkb __unused, uint16_t trans_proto __unused,
-	      struct in6_addr *dest __unused) {
+static int ipv6_tx ( struct pk_buff *pkb,
+		     struct tcpip_protocol *tcpip_protocol,
+		     struct sockaddr_tcpip *st_dest ) {
 	return -ENOSYS;
 }
 
@@ -22,13 +23,10 @@ int ipv6_tx ( struct pk_buff *pkb __unused, uint16_t trans_proto __unused,
  *
  * Placeholder function. Should rewrite in net/ipv6.c
  */
-void ipv6_rx ( struct pk_buff *pkb __unused,
-	       struct net_device *netdev __unused,
-	       const void *ll_source __unused ) {
-}
-
-void ipv6_tx_csum ( struct pk_buff *pkb, struct tcpip_protocol *tcpip ) {
-	return;
+static int ipv6_rx ( struct pk_buff *pkb __unused,
+		     struct net_device *netdev __unused,
+		     const void *ll_source __unused ) {
+	return -ENOSYS;
 }
 
 static const char * ipv6_ntoa ( const void *net_addr ) {
@@ -49,10 +47,9 @@ NET_PROTOCOL ( ipv6_protocol );
 
 /** IPv6 TCPIP net protocol */
 struct tcpip_net_protocol ipv6_tcpip_protocol = {
-	.net_protocol = &ipv6_protocol,
+	.name = "IPv6",
 	.sa_family = AF_INET6,
 	.tx = ipv6_tx,
-	.tx_csum = ipv6_tx_csum,
 };
 
 TCPIP_NET_PROTOCOL ( ipv6_tcpip_protocol );
diff --git a/src/net/tcp.c b/src/net/tcp.c
index 1c80132..86d27b5 100644
--- a/src/net/tcp.c
+++ b/src/net/tcp.c
@@ -379,7 +379,8 @@ void tcp_init_conn ( struct tcp_connection *conn ) {
  * peer. It sends a SYN packet to peer. When the connection is established, the
  * TCP stack calls the connected() callback function.
  */
-int tcp_connectto ( struct tcp_connection *conn, struct sockaddr *peer ) {
+int tcp_connectto ( struct tcp_connection *conn,
+		    struct sockaddr_tcpip *peer ) {
 	int rc;
 
 	/* A connection can only be established from the CLOSED state */
@@ -393,7 +394,7 @@ int tcp_connectto ( struct tcp_connection *conn, struct sockaddr *peer ) {
 	if ( ( rc = tcp_listen ( conn, conn->local_port ) ) != 0 ) {
 		return rc;
 	}
-	memcpy ( &conn->sa, peer, sizeof ( *peer ) );
+	memcpy ( &conn->peer, peer, sizeof ( conn->peer ) );
 
 	/* Send a SYN packet and transition to TCP_SYN_SENT */
 	conn->snd_una = ( ( ( uint32_t ) random() ) << 16 ) & random();
@@ -407,7 +408,7 @@ int tcp_connectto ( struct tcp_connection *conn, struct sockaddr *peer ) {
 }
 
 int tcp_connect ( struct tcp_connection *conn ) {
-	return tcp_connectto ( conn, &conn->sa );
+	return tcp_connectto ( conn, &conn->peer );
 }
 
 /**
@@ -542,7 +543,7 @@ int tcp_senddata ( struct tcp_connection *conn ) {
  * This function sends data to the peer socket address
  */
 int tcp_send ( struct tcp_connection *conn, const void *data, size_t len ) {
-	struct sockaddr *peer = &conn->sa;
+	struct sockaddr_tcpip *peer = &conn->peer;
 	struct pk_buff *pkb = conn->tx_pkb;
 	int slen;
 
@@ -557,18 +558,7 @@ int tcp_send ( struct tcp_connection *conn, const void *data, size_t len ) {
 	/* Source port, assumed to be in network byte order in conn */
 	tcphdr->src = conn->local_port;
 	/* Destination port, assumed to be in network byte order in peer */
-	switch ( peer->sa_family ) {
-	case AF_INET:
-		tcphdr->dest = peer->sin.sin_port;
-		break;
-	case AF_INET6:
-		tcphdr->dest = peer->sin6.sin6_port;
-		break;
-	default:
-		DBG ( "Family type %d not supported\n", 
-					peer->sa_family );
-		return -EAFNOSUPPORT;
-	}
+	tcphdr->dest = peer->st_port;
 	tcphdr->seq = htonl ( conn->snd_una );
 	tcphdr->ack = htonl ( conn->rcv_nxt );
 	/* Header length, = 0x50 (without TCP options) */
@@ -597,7 +587,9 @@ int tcp_send ( struct tcp_connection *conn, const void *data, size_t len ) {
  * @v pkb	Packet buffer
  * @v partial	Partial checksum
  */
-void tcp_rx ( struct pk_buff *pkb, uint16_t partial ) {
+static int tcp_rx ( struct pk_buff *pkb,
+		    struct sockaddr_tcpip *st_src __unused,
+		    struct sockaddr_tcpip *st_dest __unused ) {
 	struct tcp_connection *conn;
 	struct tcp_header *tcphdr;
 	uint32_t acked, toack;
@@ -606,7 +598,7 @@ void tcp_rx ( struct pk_buff *pkb, uint16_t partial ) {
 	/* Sanity check */
 	if ( pkb_len ( pkb ) < sizeof ( *tcphdr ) ) {
 		DBG ( "Packet too short (%d bytes)\n", pkb_len ( pkb ) );
-		return;
+		return -EINVAL;
 	}
 
 	/* Process TCP header */
@@ -616,7 +608,7 @@ void tcp_rx ( struct pk_buff *pkb, uint16_t partial ) {
 	hlen = ( ( tcphdr->hlen & TCP_MASK_HLEN ) / 16 ) * 4;
 	if ( hlen != sizeof ( *tcphdr ) ) {
 		DBG ( "Bad header length (%d bytes)\n", hlen );
-		return;
+		return -EINVAL;
 	}
 	
 	/* TODO: Verify checksum */
@@ -629,7 +621,7 @@ void tcp_rx ( struct pk_buff *pkb, uint16_t partial ) {
 	}
 	
 	DBG ( "No connection found on port %d\n", ntohs ( tcphdr->dest ) );
-	return;
+	return 0;
 
   found_conn:
 	/* Set the advertised window */
@@ -642,7 +634,7 @@ void tcp_rx ( struct pk_buff *pkb, uint16_t partial ) {
 	case TCP_CLOSED:
 		DBG ( "tcp_rx(): Invalid state %s\n",
 				tcp_states[conn->tcp_state] );
-		return;
+		return -EINVAL;
 	case TCP_LISTEN:
 		if ( tcphdr->flags & TCP_SYN ) {
 			tcp_trans ( conn, TCP_SYN_RCVD );
@@ -687,7 +679,7 @@ void tcp_rx ( struct pk_buff *pkb, uint16_t partial ) {
 		if ( tcphdr->flags & TCP_RST ) {
 			tcp_trans ( conn, TCP_LISTEN );
 			conn->tcp_op->closed ( conn, CONN_RESTART );
-			return;
+			return 0;
 		}
 		if ( tcphdr->flags & TCP_ACK ) {
 			tcp_trans ( conn, TCP_ESTABLISHED );
@@ -697,7 +689,7 @@ void tcp_rx ( struct pk_buff *pkb, uint16_t partial ) {
 			 */
 			conn->snd_una = tcphdr->ack - 1;
 			conn->tcp_op->connected ( conn );
-			return;
+			return 0;
 		}
 		/* Unexpected packet */
 		goto unexpected;
@@ -744,7 +736,7 @@ void tcp_rx ( struct pk_buff *pkb, uint16_t partial ) {
 	case TCP_CLOSING:
 		if ( tcphdr->flags & TCP_ACK ) {
 			tcp_trans ( conn, TCP_TIME_WAIT );
-			return;
+			return 0;
 		}
 		/* Unexpected packet */
 		goto unexpected;
@@ -757,7 +749,7 @@ void tcp_rx ( struct pk_buff *pkb, uint16_t partial ) {
 	case TCP_LAST_ACK:
 		if ( tcphdr->flags & TCP_ACK ) {
 			tcp_trans ( conn, TCP_CLOSED );
-			return;
+			return 0;
 		}
 		/* Unexpected packet */
 		goto unexpected;
@@ -791,7 +783,7 @@ void tcp_rx ( struct pk_buff *pkb, uint16_t partial ) {
 		acked = ntohl ( tcphdr->ack ) - conn->snd_una;
 		if ( acked < 0 ) { /* TODO: Replace all uint32_t arith */
 			DBG ( "Previously ACKed (%d)\n", tcphdr->ack );
-			return;
+			return 0;
 		}
 		/* Advance snd stream */
 		conn->snd_una += acked;
@@ -802,7 +794,7 @@ void tcp_rx ( struct pk_buff *pkb, uint16_t partial ) {
 		/* Invoke the senddata() callback function */
 		tcp_senddata ( conn );
 	}
-	return;
+	return 0;
 
   send_tcp_nomsg:
 	free_pkb ( conn->tx_pkb );
@@ -812,20 +804,20 @@ void tcp_rx ( struct pk_buff *pkb, uint16_t partial ) {
 	if ( ( rc = tcp_send ( conn, TCP_NOMSG, TCP_NOMSG_LEN ) ) != 0 ) {
 		DBG ( "Error sending TCP message (rc = %d)\n", rc );
 	}
-	return;
+	return 0;
 
   unexpected:
 	DBG ( "Unexpected packet received in %d state with flags = %hd\n",
 			conn->tcp_state, tcphdr->flags & TCP_MASK_FLAGS );
 	free_pkb ( conn->tx_pkb );
-	return;
+	return -EINVAL;
 }
 
 /** TCP protocol */
 struct tcpip_protocol tcp_protocol = {
 	.name = "TCP",
 	.rx = tcp_rx,
-	.trans_proto = IP_TCP,
+	.tcpip_proto = IP_TCP,
 	.csum_offset = 16,
 };
 
diff --git a/src/net/tcpip.c b/src/net/tcpip.c
index 670d337..4a44a55 100644
--- a/src/net/tcpip.c
+++ b/src/net/tcpip.c
@@ -1,94 +1,86 @@
 #include <stdint.h>
 #include <string.h>
 #include <errno.h>
-#include <malloc.h>
 #include <byteswap.h>
-#include <gpxe/in.h>
-#include <gpxe/ip.h>
-#include <gpxe/ip6.h>
 #include <gpxe/pkbuff.h>
 #include <gpxe/tables.h>
-#include <gpxe/netdevice.h>
 #include <gpxe/tcpip.h>
 
 /** @file
  *
  * Transport-network layer interface
  *
- * This file contains functions and utilities for the transport-network layer interface
+ * This file contains functions and utilities for the
+ * TCP/IP transport-network layer interface
  */
 
-/** Registered network-layer protocols that support TCPIP */
-static struct tcpip_net_protocol tcpip_net_protocols[0] __table_start ( tcpip_net_protocols );
-static struct tcpip_net_protocol tcpip_net_protocols_end[0] __table_end ( tcpip_net_protocols );
+/** Registered network-layer protocols that support TCP/IP */
+static struct tcpip_net_protocol
+tcpip_net_protocols[0] __table_start ( tcpip_net_protocols );
+static struct tcpip_net_protocol
+tcpip_net_protocols_end[0] __table_end ( tcpip_net_protocols );
 
-/** Registered transport-layer protocols that support TCPIP */
-static struct tcpip_protocol tcpip_protocols[0] __table_start ( tcpip_protocols );
-static struct tcpip_protocol tcpip_protocols_end[0] __table_end ( tcpip_protocols );
+/** Registered transport-layer protocols that support TCP/IP */
+static struct tcpip_protocol
+tcpip_protocols[0]__table_start ( tcpip_protocols );
+static struct tcpip_protocol
+tcpip_protocols_end[0] __table_end ( tcpip_protocols );
 
-/** Process a received packet
+/** Process a received TCP/IP packet
  *
  * @v pkb		Packet buffer
- * @v trans_proto	Transport-layer protocol number
- * @v src		Source network-layer address
- * @v dest		Destination network-layer address
+ * @v tcpip_proto	Transport-layer protocol number
+ * @v st_src		Partially-filled source address
+ * @v st_dest		Partially-filled destination address
+ * @ret rc		Return status code
  *
- * This function expects a transport-layer segment from the network-layer
+ * This function expects a transport-layer segment from the network
+ * layer.  The network layer should fill in as much as it can of the
+ * source and destination addresses (i.e. it should fill in the
+ * address family and the network-layer addresses, but leave the ports
+ * and the rest of the structures as zero).
  */
-void tcpip_rx ( struct pk_buff *pkb, uint8_t trans_proto, struct in_addr *src,
-		struct in_addr *dest ) {
+int tcpip_rx ( struct pk_buff *pkb, uint8_t tcpip_proto, 
+	       struct sockaddr_tcpip *st_src,
+	       struct sockaddr_tcpip *st_dest ) {
 	struct tcpip_protocol *tcpip;
 
-	/* Identify the transport layer protocol */
-	for ( tcpip = tcpip_protocols; tcpip <= tcpip_protocols_end; ++tcpip ) {
-		if ( tcpip->trans_proto == trans_proto ) {
-			tcpip->rx ( pkb, src, dest );
+	/* Hand off packet to the appropriate transport-layer protocol */
+	for ( tcpip = tcpip_protocols; tcpip < tcpip_protocols_end; tcpip++ ) {
+		if ( tcpip->tcpip_proto == tcpip_proto ) {
+			DBG ( "TCP/IP received %s packet\n", tcpip->name );
+			return tcpip->rx ( pkb, st_src, st_dest );
 		}
 	}
+
+	DBG ( "Unrecognised TCP/IP protocol %d\n", tcpip_proto );
+	return -EPROTONOSUPPORT;
 }
 
-/** Transmit a transport-layer segment
+/** Transmit a TCP/IP packet
  *
  * @v pkb		Packet buffer
- * @v trans_proto	Transport-layer protocol
- * @v sock		Destination socket address
- * @ret			Status
+ * @v tcpip_protocol	Transport-layer protocol
+ * @v st_dest		Destination address
+ * @ret rc		Return status code
  */
-int tcpip_tx ( struct pk_buff *pkb, struct tcpip_protocol *tcpip,
-	       struct sockaddr *sock ) {
-
-#if 0 /* This is the right thing to do */
-
+int tcpip_tx ( struct pk_buff *pkb, struct tcpip_protocol *tcpip_protocol,
+	       struct sockaddr_tcpip *st_dest ) {
 	struct tcpip_net_protocol *tcpip_net;
 
-	/* Identify the network layer protocol */
-	for ( tcpip_net = tcpip_net_protocols; 
-			tcpip_net <= tcpip_net_protocols_end; ++tcpip_net ) {
-		if ( tcpip_net->sa_family == sock->sa_family ) {
-			DBG ( "Packet sent to %s module\n", tcpip_net->net_protocol->name );
-			return tcpip_net->tx ( pkb, tcpip, sock );
+	/* Hand off packet to the appropriate network-layer protocol */
+	for ( tcpip_net = tcpip_net_protocols ;
+	      tcpip_net < tcpip_net_protocols_end ; tcpip_net++ ) {
+		if ( tcpip_net->sa_family == st_dest->st_family ) {
+			DBG ( "TCP/IP sending %s packet\n", tcpip_net->name );
+			return tcpip_net->tx ( pkb, tcpip_protocol, st_dest );
 		}
 	}
-	DBG ( "No suitable network layer protocol found for sa_family %s\n",
-						( sock->sa_family );
-	return -EAFNOSUPPORT;
-}
-
-#else
-
-	/* Identify the network layer protocol and send it using xxx_tx() */
-	switch ( sock->sa_family ) {
-	case AF_INET: /* IPv4 network family */
-		return ipv4_tx ( pkb, tcpip, sock );
-	case AF_INET6: /* IPv6 network family */
-		return ipv6_tx ( pkb, tcpip, sock );
-	}
-	DBG ( "Network family %d not supported", sock->sa_family );
+	
+	DBG ( "Unrecognised TCP/IP address family %d\n", st_dest->st_family );
 	return -EAFNOSUPPORT;
 }
 
-#endif
-
 /**
  * Calculate continued TCP/IP checkum
  *
diff --git a/src/net/udp.c b/src/net/udp.c
index 1f6a899..eb640ad 100644
--- a/src/net/udp.c
+++ b/src/net/udp.c
@@ -2,95 +2,120 @@
 #include <string.h>
 #include <assert.h>
 #include <byteswap.h>
-#include <latch.h>
 #include <errno.h>
-#include <gpxe/in.h>
-#include <gpxe/ip.h>
-#include <gpxe/ip6.h>
-#include <gpxe/udp.h>
-#include <gpxe/init.h>
+#include <gpxe/tcpip.h>
 #include <gpxe/pkbuff.h>
 #include <gpxe/netdevice.h>
-#include <gpxe/tcpip.h>
+#include <gpxe/udp.h>
 
 /** @file
  *
  * UDP protocol
  */
 
+struct tcpip_protocol udp_protocol;
+
 /**
  * List of registered UDP connections
  */
 static LIST_HEAD ( udp_conns );
 
 /**
- * Some utility functions
+ * Bind UDP connection to local port
+ *
+ * @v conn		UDP connection
+ * @v local_port	Local port, in network byte order
+ * @ret rc		Return status code
  */
-static inline void copy_sockaddr ( struct sockaddr *source, struct sockaddr *dest ) {
-	memcpy ( dest, source, sizeof ( *dest ) );
-}
+int udp_bind ( struct udp_connection *conn, uint16_t local_port ) {
+	struct udp_connection *existing;
 
-static inline uint16_t * dest_port ( struct sockaddr *sock ) {
-	switch ( sock->sa_family ) {
-	case AF_INET:
-		return &sock->sin.sin_port;
-	case AF_INET6:
-		return &sock->sin6.sin6_port;
+	list_for_each_entry ( existing, &udp_conns, list ) {
+		if ( existing->local_port == local_port )
+			return -EADDRINUSE;
 	}
-	return NULL;
+	conn->local_port = local_port;
+	return 0;
 }
 
 /**
- * Dump the UDP header
+ * Bind UDP connection to all local ports
  *
- * @v udphdr	UDP header
+ * @v conn		UDP connection
+ *
+ * A promiscuous UDP connection will receive packets with any
+ * destination UDP port.  This is required in order to support the PXE
+ * UDP API.
+ *
+ * If the promiscuous connection is not the only UDP connection, the
+ * behaviour is undefined.
  */
-void udp_dump ( struct udp_header *udphdr ) {
-
-	/* Print UDP header for debugging */
-	DBG ( "UDP header at %p + %#zx\n", udphdr, sizeof ( *udphdr ) );
-	DBG ( "\tSource Port = %d\n", ntohs ( udphdr->source_port ) );
-	DBG ( "\tDestination Port = %d\n", ntohs ( udphdr->dest_port ) );
-	DBG ( "\tLength = %d\n", ntohs ( udphdr->len ) );
-	DBG ( "\tChecksum = %x\n", ntohs ( udphdr->chksum ) );
-	DBG ( "\tChecksum located at %p\n", &udphdr->chksum );
+void udp_bind_promisc ( struct udp_connection *conn ) {
+	conn->local_port = 0;
 }
 
 /**
- * Open a UDP connection
+ * Connect UDP connection to remote host and port
  *
- * @v conn      UDP connection
- * @v peer      Destination socket address
+ * @v conn		UDP connection
+ * @v peer		Destination socket address
  *
  * This function stores the socket address within the connection
  */
-void udp_connect ( struct udp_connection *conn, struct sockaddr *peer ) {
-	copy_sockaddr ( peer, &conn->sa );
-
-	/* Not sure if this should add the connection to udp_conns; If it does,
-	 * uncomment the following code
-	 */
-//	list_add ( &conn->list, &udp_conns );
+void udp_connect ( struct udp_connection *conn, struct sockaddr_tcpip *peer ) {
+	memcpy ( &conn->peer, peer, sizeof ( conn->peer ) );
 }
 
 /**
- * Initialize a UDP connection
+ * Open a local port
+ *
+ * @v conn		UDP connection
+ * @v local_port	Local port, in network byte order, or zero
+ * @ret rc		Return status code
  *
- * @v conn      UDP connection
- * @v udp_op	UDP operations
+ * Opens the UDP connection and binds to a local port.  If no local
+ * port is specified, the first available port will be used.
  */
-void udp_init ( struct udp_connection *conn, struct udp_operations *udp_op ) {
-	conn->local_port = 0;
-	conn->tx_pkb = NULL;
-	if ( udp_op != NULL ) {
-		conn->udp_op = udp_op;
+int udp_open ( struct udp_connection *conn, uint16_t local_port ) {
+	static uint16_t try_port = 1024;
+	int rc;
+
+	/* If no port specified, find the first available port */
+	if ( ! local_port ) {
+		for ( ; try_port ; try_port++ ) {
+			if ( try_port < 1024 )
+				continue;
+			if ( udp_open ( conn, htons ( try_port ) ) == 0 )
+				return 0;
+		}
+		return -EADDRINUSE;
 	}
+
+	/* Attempt bind to local port */
+	if ( ( rc = udp_bind ( conn, local_port ) ) != 0 )
+		return rc;
+
+	/* Add to UDP connection list */
+	list_add ( &conn->list, &udp_conns );
+	DBG ( "UDP opened %p on port %d\n", conn, ntohs ( local_port ) );
+
+	return 0;
+}
+
+/**
+ * Close a UDP connection
+ *
+ * @v conn		UDP connection
+ */
+void udp_close ( struct udp_connection *conn ) {
+	list_del ( &conn->list );
+	DBG ( "UDP closed %p\n", conn );
 }
 
 /**
  * User request to send data via a UDP connection
  *
- * @v conn	UDP connection
+ * @v conn		UDP connection
  *
  * This function allocates buffer space and invokes the function's senddata()
  * callback. The callback may use the buffer space
@@ -98,8 +123,8 @@ void udp_init ( struct udp_connection *conn, struct udp_operations *udp_op ) {
 int udp_senddata ( struct udp_connection *conn ) {
 	conn->tx_pkb = alloc_pkb ( UDP_MAX_TXPKB );
 	if ( conn->tx_pkb == NULL ) {
-		DBG ( "Error allocating packet buffer of length %d\n",
-							UDP_MAX_TXPKB );
+		DBG ( "UDP %p cannot allocate packet buffer of length %d\n",
+		      conn, UDP_MAX_TXPKB );
 		return -ENOMEM;
 	}
 	pkb_reserve ( conn->tx_pkb, UDP_MAX_HLEN );
@@ -111,19 +136,25 @@ int udp_senddata ( struct udp_connection *conn ) {
 /**
  * Transmit data via a UDP connection to a specified address
  *
- * @v conn      UDP connection
- * @v peer	Destination address
- * @v data      Data to send
- * @v len       Length of data
+ * @v conn		UDP connection
+ * @v peer		Destination address
+ * @v data		Data to send
+ * @v len		Length of data
+ * @ret rc		Return status code
  *
- * This function fills up the UDP headers and sends the data. Discover the
- * network protocol through the sa_family field in the destination socket
- * address.
+ * This function fills up the UDP headers and sends the data.  It may
+ * be called only from within the context of an application's
+ * senddata() method; if the application wishes to send data it must
+ * call udp_senddata() and wait for its senddata() method to be
+ * called.
  */
-int udp_sendto ( struct udp_connection *conn, struct sockaddr *peer,
+int udp_sendto ( struct udp_connection *conn, struct sockaddr_tcpip *peer,
 		 const void *data, size_t len ) {
-       	struct udp_header *udphdr;		/* UDP header */
-	uint16_t *dest;
+       	struct udp_header *udphdr;
+
+	/* Avoid overflowing TX buffer */
+	if ( len > pkb_available ( conn->tx_pkb ) )
+		len = pkb_available ( conn->tx_pkb );
 
 	/* Copy payload */
 	memmove ( pkb_put ( conn->tx_pkb, len ), data, len );
@@ -135,104 +166,77 @@ int udp_sendto ( struct udp_connection *conn, struct sockaddr *peer,
 	 * sending it over the network
 	 */
 	udphdr = pkb_push ( conn->tx_pkb, sizeof ( *udphdr ) );
-	if ( (dest = dest_port ( peer ) ) == NULL ) {
-		DBG ( "Network family %d not supported\n", peer->sa_family );
-		return -EAFNOSUPPORT;
-	}
-	udphdr->dest_port = *dest;
+	udphdr->dest_port = peer->st_port;
 	udphdr->source_port = conn->local_port;
 	udphdr->len = htons ( pkb_len ( conn->tx_pkb ) );
 	udphdr->chksum = 0;
 	udphdr->chksum = tcpip_chksum ( udphdr, sizeof ( *udphdr ) + len );
 
-	/**
-	 * Dump the contents of the UDP header
-	 */
-	udp_dump ( udphdr );
+	/* Dump debugging information */
+	DBG ( "UDP %p transmitting %p+%#zx len %#x src %d dest %d "
+	      "chksum %#04x\n", conn, conn->tx_pkb->data,
+	      pkb_len ( conn->tx_pkb ), ntohs ( udphdr->len ),
+	      ntohs ( udphdr->source_port ), ntohs ( udphdr->dest_port ),
+	      ntohs ( udphdr->chksum ) );
 
 	/* Send it to the next layer for processing */
 	return tcpip_tx ( conn->tx_pkb, &udp_protocol, peer );
 }
 
 /**
- * Transmit data via a UDP connection to a specified address
- *
- * @v conn      UDP connection
- * @v data      Data to send
- * @v len       Length of data
- */
-int udp_send ( struct udp_connection *conn, const void *data, size_t len ) {
-	return udp_sendto ( conn, &conn->sa, data, len );
-}
-
-/**
- * Close a UDP connection
- *
- * @v conn      UDP connection
- */
-void udp_close ( struct udp_connection *conn ) {
-	list_del ( &conn->list );
-}
-
-/**
- * Open a local port
+ * Transmit data via a UDP connection
  *
  * @v conn		UDP connection
- * @v local_port	Local port on which to open connection
+ * @v data		Data to send
+ * @v len		Length of data
+ * @ret rc		Return status code
  *
- * This does not support the 0 port option correctly yet
+ * This function fills up the UDP headers and sends the data.  It may
+ * be called only from within the context of an application's
+ * senddata() method; if the application wishes to send data it must
+ * call udp_senddata() and wait for its senddata() method to be
+ * called.
  */
-int udp_open ( struct udp_connection *conn, uint16_t local_port ) {
-	struct udp_connection *connr;
-	uint16_t min_port = 0xffff;
-
-	/* Iterate through udp_conns to see if local_port is available */
-	list_for_each_entry ( connr, &udp_conns, list ) {
-		if ( connr->local_port == local_port ) {
-			return -EISCONN;
-		}
-		if ( min_port > connr->local_port ) {
-			min_port = connr->local_port;
-		}
-	}
-	/* This code is buggy. I will update it soon :) */
-	conn->local_port = local_port == 0 ? min_port > 1024 ? 1024 :
-						min_port + 1 : local_port;
-
-	/* Add the connection to the list of listening connections */
-	list_add ( &conn->list, &udp_conns );
-	return 0;
+int udp_send ( struct udp_connection *conn, const void *data, size_t len ) {
+	return udp_sendto ( conn, &conn->peer, data, len );
 }
 
 /**
  * Process a received packet
  *
- * @v pkb	       Packet buffer
- * @v src_net_addr      Source network address
- * @v dest_net_addr     Destination network address
+ * @v pkb		Packet buffer
+ * @v st_src		Partially-filled source address
+ * @v st_dest		Partially-filled destination address
+ * @ret rc		Return status code
  */
-void udp_rx ( struct pk_buff *pkb, struct in_addr *src_net_addr __unused,
-			struct in_addr *dest_net_addr __unused ) {
+static int udp_rx ( struct pk_buff *pkb, struct sockaddr_tcpip *st_src,
+		    struct sockaddr_tcpip *st_dest ) {
 	struct udp_header *udphdr = pkb->data;
 	struct udp_connection *conn;
-	uint16_t ulen;
+	unsigned int ulen;
 	uint16_t chksum;
 
-	udp_dump ( udphdr );
-
-	/* Validate the packet and the UDP length */
+	/* Sanity check */
 	if ( pkb_len ( pkb ) < sizeof ( *udphdr ) ) {
-		DBG ( "UDP packet too short (%d bytes)\n",
-		      pkb_len ( pkb ) );
-		return;
+		DBG ( "UDP received underlength packet %p+%#zx\n",
+		      pkb->data, pkb_len ( pkb ) );
+		return -EINVAL;
 	}
 
+	/* Dump debugging information */
+	DBG ( "UDP received %p+%#zx len %#x src %d dest %d chksum %#04x\n",
+	      pkb->data, pkb_len ( pkb ), ntohs ( udphdr->len ),
+	      ntohs ( udphdr->source_port ), ntohs ( udphdr->dest_port ),
+	      ntohs ( udphdr->chksum ) );
+
+	/* Check length and trim any excess */
 	ulen = ntohs ( udphdr->len );
-	if ( ulen != pkb_len ( pkb ) ) {
-		DBG ( "Inconsistent UDP packet length (%d bytes)\n",
-		      pkb_len ( pkb ) );
-		return;
+	if ( ulen > pkb_len ( pkb ) ) {
+		DBG ( "UDP received truncated packet %p+%#zx\n",
+		      pkb->data, pkb_len ( pkb ) );
+		return -EINVAL;
 	}
+	pkb_unput ( pkb, ( pkb_len ( pkb ) - ulen ) );
 
 	/* Verify the checksum */
 #warning "Don't we need to take the pseudo-header into account here?"
@@ -240,32 +244,49 @@ void udp_rx ( struct pk_buff *pkb, struct in_addr *src_net_addr __unused,
 	chksum = tcpip_chksum ( pkb->data, pkb_len ( pkb ) );
 	if ( chksum != 0xffff ) {
 		DBG ( "Bad checksum %#x\n", chksum );
-		return;
+		return -EINVAL;
 	}
 #endif
 
-	/* Todo: Check if it is a broadcast or multicast address */
+	/* Complete the socket addresses */
+	st_src->st_port = udphdr->source_port;
+	st_dest->st_port = udphdr->dest_port;
 
 	/* Demux the connection */
 	list_for_each_entry ( conn, &udp_conns, list ) {
-		if ( conn->local_port == udphdr->dest_port ) {
-			goto conn;
+		if ( conn->local_port &&
+		     ( conn->local_port != udphdr->dest_port ) ) {
+			/* Bound to local port and local port doesn't match */
+			continue;
 		}
-	}
-	return;
+		if ( conn->peer.st_family &&
+		     ( memcmp ( &conn->peer, st_src,
+				sizeof ( conn->peer ) ) != 0 ) ) {
+			/* Connected to remote port and remote port
+			 * doesn't match
+			 */
+			continue;
+		}
+		
+		/* Strip off the UDP header */
+		pkb_pull ( pkb, sizeof ( *udphdr ) );
 
-	conn:
-	/** Strip off the UDP header */
-	pkb_pull ( pkb, sizeof ( *udphdr ) );
+		DBG ( "UDP delivering to %p\n", conn );
+		
+		/* Call the application's callback */
+		return conn->udp_op->newdata ( conn, pkb->data, pkb_len( pkb ),
+					       st_src, st_dest );
+	}
 
-	/** Call the application's callback */
-	conn->udp_op->newdata ( conn, pkb->data, ulen - sizeof ( *udphdr ) );
+	DBG ( "No UDP connection listening on port %d\n",
+	      ntohs ( udphdr->dest_port ) );
+	return 0;
 }
 
 struct tcpip_protocol udp_protocol  = {
 	.name = "UDP",
 	.rx = udp_rx,
-	.trans_proto = IP_UDP,
+	.tcpip_proto = IP_UDP,
 	.csum_offset = 6,
 };
 
diff --git a/src/net/udp/dhcp.c b/src/net/udp/dhcp.c
index b214901..c2a9097 100644
--- a/src/net/udp/dhcp.c
+++ b/src/net/udp/dhcp.c
@@ -491,9 +491,12 @@ static void dhcp_done ( struct dhcp_session *dhcp, int rc ) {
 }
 
 /** Address for transmitting DHCP requests */
-static struct sockaddr sa_dhcp_server = {
-	.sa_family = AF_INET,
+static union {
+	struct sockaddr_tcpip st;
+	struct sockaddr_in sin;
+} sa_dhcp_server = {
 	.sin = {
+		.sin_family = AF_INET,
 		.sin_addr.s_addr = INADDR_BROADCAST,
 		.sin_port = htons ( BOOTPS_PORT ),
 	},
@@ -548,7 +551,7 @@ static void dhcp_senddata ( struct udp_connection *conn,
 	}
 
 	/* Transmit the packet */
-	if ( ( rc = udp_sendto ( conn, &sa_dhcp_server,
+	if ( ( rc = udp_sendto ( conn, &sa_dhcp_server.st,
 				 dhcppkt.dhcphdr, dhcppkt.len ) ) != 0 ) {
 		DBG ( "Could not transmit UDP packet\n" );
 		return;
@@ -588,9 +591,12 @@ static void dhcp_timer_expired ( struct retry_timer *timer, int fail ) {
  * @v udp		UDP connection
  * @v data		Received data
  * @v len		Length of received data
+ * @v st_src		Partially-filled source address
+ * @v st_dest		Partially-filled destination address
  */
-static void dhcp_newdata ( struct udp_connection *conn,
-			   void *data, size_t len ) {
+static int dhcp_newdata ( struct udp_connection *conn, void *data, size_t len,
+			  struct sockaddr_tcpip *st_src __unused,
+			  struct sockaddr_tcpip *st_dest __unused ) {
 	struct dhcp_session *dhcp = udp_to_dhcp ( conn );
 	struct dhcphdr *dhcphdr = data;
 	struct dhcp_option_block *options;
@@ -600,14 +606,14 @@ static void dhcp_newdata ( struct udp_connection *conn,
 	if ( dhcphdr->xid != dhcp->xid ) {
 		DBG ( "DHCP wrong transaction ID (wanted %08lx, got %08lx)\n",
 		      ntohl ( dhcphdr->xid ), ntohl ( dhcp->xid ) );
-		return;
+		return 0;
 	};
 
 	/* Parse packet and create options structure */
 	options = dhcp_parse ( dhcphdr, len );
 	if ( ! options ) {
 		DBG ( "Could not parse DHCP packet\n" );
-		return;
+		return -EINVAL;
 	}
 
 	/* Determine message type */
@@ -643,10 +649,11 @@ static void dhcp_newdata ( struct udp_connection *conn,
 	} else {
 		dhcp_done ( dhcp, 0 );
 	}
-	return;
+	return 0;
 
  out_discard:
 	free_dhcp_options ( options );
+	return 0;
 }
 
 /** DHCP UDP operations */