Clarified packet ownership transfer between a few functions.

Added a large number of missing calls to free_pkb().  In the case of UDP,
no received packets were ever freed, which lead to memory exhaustion
remarkably quickly once pxelinux started up.

In general, any function with _rx() in its name which accepts a pk_buff
*must* either call free_pkb() or pass the pkb to another _rx() function
(e.g. the next layer up the stack).  Since the UDP (and TCP) layers don't
pass packet buffers up to the higher-layer protocols (the
"applications"), they must free the packet buffer after calling the
application's newdata() method.

6 files changed, 63 insertions, 31 deletions

diff --git a/src/include/gpxe/netdevice.h b/src/include/gpxe/netdevice.h
index b8cf9aeec..60ce886d9 100644
--- a/src/include/gpxe/netdevice.h
+++ b/src/include/gpxe/netdevice.h
@@ -83,6 +83,7 @@ struct ll_protocol {
 	 *
 	 * This method should prepend in the link-layer header
 	 * (e.g. the Ethernet DIX header) and transmit the packet.
+	 * This method takes ownership of the packet buffer.
 	 */
 	int ( * tx ) ( struct pk_buff *pkb, struct net_device *netdev,
 		       struct net_protocol *net_protocol,
@@ -95,9 +96,10 @@ struct ll_protocol {
 	 *
 	 * This method should strip off the link-layer header
 	 * (e.g. the Ethernet DIX header) and pass the packet to
-	 * net_rx().
+	 * net_rx().  This method takes ownership of the packet
+	 * buffer.
 	 */
-	void ( * rx ) ( struct pk_buff *pkb, struct net_device *netdev );
+	int ( * rx ) ( struct pk_buff *pkb, struct net_device *netdev );
 	/**
 	 * Transcribe link-layer address
 	 *
@@ -206,8 +208,8 @@ extern int netdev_tx ( struct net_device *netdev, struct pk_buff *pkb );
 extern void netdev_rx ( struct net_device *netdev, struct pk_buff *pkb );
 extern int net_tx ( struct pk_buff *pkb, struct net_device *netdev,
 		    struct net_protocol *net_protocol, const void *ll_dest );
-extern void net_rx ( struct pk_buff *pkb, struct net_device *netdev,
-		     uint16_t net_proto, const void *ll_source );
+extern int net_rx ( struct pk_buff *pkb, struct net_device *netdev,
+		    uint16_t net_proto, const void *ll_source );
 extern int netdev_poll ( struct net_device *netdev );
 extern struct pk_buff * netdev_rx_dequeue ( struct net_device *netdev );
 extern struct net_device * alloc_netdev ( size_t priv_size );
diff --git a/src/net/ethernet.c b/src/net/ethernet.c
index 13b3c2dbc..c4b526f51 100644
--- a/src/net/ethernet.c
+++ b/src/net/ethernet.c
@@ -19,6 +19,7 @@
 #include <stdint.h>
 #include <string.h>
 #include <byteswap.h>
+#include <errno.h>
 #include <assert.h>
 #include <vsprintf.h>
 #include <gpxe/if_arp.h>
@@ -68,21 +69,22 @@ static int eth_tx ( struct pk_buff *pkb, struct net_device *netdev,
  * Strips off the Ethernet link-layer header and passes up to the
  * network-layer protocol.
  */
-static void eth_rx ( struct pk_buff *pkb, struct net_device *netdev ) {
+static int eth_rx ( struct pk_buff *pkb, struct net_device *netdev ) {
 	struct ethhdr *ethhdr = pkb->data;
 
 	/* Sanity check */
 	if ( pkb_len ( pkb ) < sizeof ( *ethhdr ) ) {
 		DBG ( "Ethernet packet too short (%d bytes)\n",
 		      pkb_len ( pkb ) );
-		return;
+		free_pkb ( pkb );
+		return -EINVAL;
 	}
 
 	/* Strip off Ethernet header */
 	pkb_pull ( pkb, sizeof ( *ethhdr ) );
 
 	/* Hand off to network-layer protocol */
-	net_rx ( pkb, netdev, ethhdr->h_protocol, ethhdr->h_source );
+	return net_rx ( pkb, netdev, ethhdr->h_protocol, ethhdr->h_source );
 }
 
 /**
diff --git a/src/net/ipv4.c b/src/net/ipv4.c
index 59fa9f82e..da16452fd 100644
--- a/src/net/ipv4.c
+++ b/src/net/ipv4.c
@@ -389,9 +389,8 @@ static int ipv4_rx ( struct pk_buff *pkb, struct net_device *netdev __unused,
 
 	/* Sanity check */
 	if ( pkb_len ( pkb ) < sizeof ( *iphdr ) ) {
-		DBG ( "IP datagram too short (%d bytes)\n",
-			pkb_len ( pkb ) );
-		return -EINVAL;
+		DBG ( "IP datagram too short (%d bytes)\n", pkb_len ( pkb ) );
+		goto err;
 	}
 
 	/* Print IP4 header for debugging */
@@ -400,14 +399,14 @@ static int 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 -EINVAL;
+		goto err;
 	}
 
 	/* Validate length of IP packet */
 	if ( ntohs ( iphdr->len ) > pkb_len ( pkb ) ) {
 		DBG ( "Inconsistent packet length %d\n",
 		      ntohs ( iphdr->len ) );
-		return -EINVAL;
+		goto err;
 	}
 
 	/* Verify the checksum */
@@ -447,6 +446,10 @@ static int ipv4_rx ( struct pk_buff *pkb, struct net_device *netdev __unused,
 
 	/* Send it to the transport layer */
 	return tcpip_rx ( pkb, iphdr->protocol, &src.st, &dest.st );
+
+ err:
+	free_pkb ( pkb );
+	return -EINVAL;
 }
 
 /** 
diff --git a/src/net/netdevice.c b/src/net/netdevice.c
index 825613cc8..634977f7a 100644
--- a/src/net/netdevice.c
+++ b/src/net/netdevice.c
@@ -97,8 +97,9 @@ int net_tx ( struct pk_buff *pkb, struct net_device *netdev,
  * @v netdev		Network device
  * @v net_proto		Network-layer protocol, in network-byte order
  * @v ll_source		Source link-layer address
+ * @ret rc		Return status code
  */
-void net_rx ( struct pk_buff *pkb, struct net_device *netdev,
+int net_rx ( struct pk_buff *pkb, struct net_device *netdev,
 	      uint16_t net_proto, const void *ll_source ) {
 	struct net_protocol *net_protocol;
 
@@ -106,10 +107,11 @@ void net_rx ( struct pk_buff *pkb, struct net_device *netdev,
 	for ( net_protocol = net_protocols ; net_protocol < net_protocols_end ;
 	      net_protocol++ ) {
 		if ( net_protocol->net_proto == net_proto ) {
-			net_protocol->rx ( pkb, netdev, ll_source );
-			break;
+			return net_protocol->rx ( pkb, netdev, ll_source );
 		}
 	}
+	free_pkb ( pkb );
+	return 0;
 }
 
 /**
diff --git a/src/net/tcpip.c b/src/net/tcpip.c
index 4a44a5564..33c9872c8 100644
--- a/src/net/tcpip.c
+++ b/src/net/tcpip.c
@@ -54,6 +54,7 @@ int tcpip_rx ( struct pk_buff *pkb, uint8_t tcpip_proto,
 	}
 
 	DBG ( "Unrecognised TCP/IP protocol %d\n", tcpip_proto );
+	free_pkb ( pkb );
 	return -EPROTONOSUPPORT;
 }
 
@@ -78,6 +79,7 @@ int tcpip_tx ( struct pk_buff *pkb, struct tcpip_protocol *tcpip_protocol,
 	}
 	
 	DBG ( "Unrecognised TCP/IP address family %d\n", st_dest->st_family );
+	free_pkb ( pkb );
 	return -EAFNOSUPPORT;
 }
 
diff --git a/src/net/udp.c b/src/net/udp.c
index 94a71524b..418479327 100644
--- a/src/net/udp.c
+++ b/src/net/udp.c
@@ -93,6 +93,8 @@ void udp_close ( struct udp_connection *conn ) {
  * callback. The callback may use the buffer space
  */
 int udp_senddata ( struct udp_connection *conn ) {
+	int rc;
+
 	conn->tx_pkb = alloc_pkb ( UDP_MAX_TXPKB );
 	if ( conn->tx_pkb == NULL ) {
 		DBG ( "UDP %p cannot allocate packet buffer of length %d\n",
@@ -100,8 +102,11 @@ int udp_senddata ( struct udp_connection *conn ) {
 		return -ENOMEM;
 	}
 	pkb_reserve ( conn->tx_pkb, UDP_MAX_HLEN );
-	return conn->udp_op->senddata ( conn, conn->tx_pkb->data, 
-					pkb_available ( conn->tx_pkb ) );
+	rc = conn->udp_op->senddata ( conn, conn->tx_pkb->data, 
+				      pkb_available ( conn->tx_pkb ) );
+	if ( conn->tx_pkb )
+		free_pkb ( conn->tx_pkb );
+	return rc;
 }
 		
 /**
@@ -122,13 +127,20 @@ int udp_senddata ( struct udp_connection *conn ) {
 int udp_sendto ( struct udp_connection *conn, struct sockaddr_tcpip *peer,
 		 const void *data, size_t len ) {
        	struct udp_header *udphdr;
+	struct pk_buff *pkb;
+
+	/* Take ownership of packet buffer back from the
+	 * udp_connection structure.
+	 */
+	pkb = conn->tx_pkb;
+	conn->tx_pkb = NULL;
 
 	/* Avoid overflowing TX buffer */
-	if ( len > pkb_available ( conn->tx_pkb ) )
-		len = pkb_available ( conn->tx_pkb );
+	if ( len > pkb_available ( pkb ) )
+		len = pkb_available ( pkb );
 
 	/* Copy payload */
-	memmove ( pkb_put ( conn->tx_pkb, len ), data, len );
+	memmove ( pkb_put ( pkb, len ), data, len );
 
 	/*
 	 * Add the UDP header
@@ -136,22 +148,22 @@ int udp_sendto ( struct udp_connection *conn, struct sockaddr_tcpip *peer,
 	 * Covert all 16- and 32- bit integers into network btye order before
 	 * sending it over the network
 	 */
-	udphdr = pkb_push ( conn->tx_pkb, sizeof ( *udphdr ) );
+	udphdr = pkb_push ( pkb, sizeof ( *udphdr ) );
 	udphdr->dest_port = peer->st_port;
 	udphdr->source_port = conn->local_port;
-	udphdr->len = htons ( pkb_len ( conn->tx_pkb ) );
+	udphdr->len = htons ( pkb_len ( pkb ) );
 	udphdr->chksum = 0;
 	udphdr->chksum = tcpip_chksum ( udphdr, sizeof ( *udphdr ) + len );
 
 	/* 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 ),
+	      "chksum %#04x\n", conn, pkb->data,
+	      pkb_len ( 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 );
+	return tcpip_tx ( pkb, &udp_protocol, peer );
 }
 
 /**
@@ -186,12 +198,14 @@ static int udp_rx ( struct pk_buff *pkb, struct sockaddr_tcpip *st_src,
 	struct udp_connection *conn;
 	unsigned int ulen;
 	uint16_t chksum;
+	int rc;
 
 	/* Sanity check */
 	if ( pkb_len ( pkb ) < sizeof ( *udphdr ) ) {
 		DBG ( "UDP received underlength packet %p+%#zx\n",
 		      pkb->data, pkb_len ( pkb ) );
-		return -EINVAL;
+		rc = -EINVAL;
+		goto done;
 	}
 
 	/* Dump debugging information */
@@ -205,7 +219,8 @@ static int udp_rx ( struct pk_buff *pkb, struct sockaddr_tcpip *st_src,
 	if ( ulen > pkb_len ( pkb ) ) {
 		DBG ( "UDP received truncated packet %p+%#zx\n",
 		      pkb->data, pkb_len ( pkb ) );
-		return -EINVAL;
+		rc = -EINVAL;
+		goto done;
 	}
 	pkb_unput ( pkb, ( pkb_len ( pkb ) - ulen ) );
 
@@ -215,7 +230,8 @@ static int udp_rx ( struct pk_buff *pkb, struct sockaddr_tcpip *st_src,
 	chksum = tcpip_chksum ( pkb->data, pkb_len ( pkb ) );
 	if ( chksum != 0xffff ) {
 		DBG ( "Bad checksum %#x\n", chksum );
-		return -EINVAL;
+		rc = -EINVAL;
+		goto done;
 	}
 #endif
 
@@ -237,13 +253,18 @@ static int udp_rx ( struct pk_buff *pkb, struct sockaddr_tcpip *st_src,
 		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 );
+		rc = conn->udp_op->newdata ( conn, pkb->data, pkb_len( pkb ),
+					     st_src, st_dest );
+		goto done;
 	}
 
 	DBG ( "No UDP connection listening on port %d\n",
 	      ntohs ( udphdr->dest_port ) );
-	return 0;
+	rc = 0;
+
+ done:
+	free_pkb ( pkb );
+	return rc;
 }
 
 struct tcpip_protocol udp_protocol  = {