[tftp] Make TFTP size requests abort transfer with an error

pxenv_tftp_get_fsize is an API call that PXE clients can call to
obtain the size of a remote file. It is implemented by starting a TFTP
transfer with pxe_tftp_open, waiting for the response and then
stopping the transfer with pxe_tftp_close(). This leaves the session
hanging on the TFTP server and it will try to resend the packet
repeatedly (verified with tftpd-hpa) until it times out.

This patch adds a method "tftpsize" that will abort the transfer after
the first packet is received from the server. This will terminate the
session on the server and is the same behaviour as Intel's PXE ROM
exhibits.

Together with a qemu patch to handle the ERROR packet (submitted to
qemu's mailing list), this resolves a specific issue where booting
pxegrub with qemu's TFTP server would be slow or hang.

I've tested this against qemu's tftp server and against my normal boot
infrastructure (tftpd-hpa). Booting pxegrub and loading extra files
now produces a trace similar to Intel's PXE client and there are no
spurious retransmits from tftpd any more.

Signed-off-by: Thomas Horsten <thomas@horsten.com>
Signed-off-by: Milan Plzik <milan.plzik@gmail.com>
Signed-off-by: Stefan Hajnoczi <stefanha@gmail.com>
Signed-off-by: Marty Connor <mdc@etherboot.org>

2 files changed, 81 insertions, 5 deletions

diff --git a/src/arch/i386/interface/pxe/pxe_tftp.c b/src/arch/i386/interface/pxe/pxe_tftp.c
index c1d831bda..0e3ca3c59 100644
--- a/src/arch/i386/interface/pxe/pxe_tftp.c
+++ b/src/arch/i386/interface/pxe/pxe_tftp.c
@@ -165,7 +165,8 @@ static struct xfer_interface_operations pxe_tftp_xfer_ops = {
  * @ret rc		Return status code
  */
 static int pxe_tftp_open ( uint32_t ipaddress, unsigned int port,
-			   const unsigned char *filename, size_t blksize ) {
+			   const unsigned char *filename, size_t blksize,
+			   int sizeonly ) {
 	char uri_string[PXE_TFTP_URI_LEN];
 	struct in_addr address;
 	int rc;
@@ -185,7 +186,8 @@ static int pxe_tftp_open ( uint32_t ipaddress, unsigned int port,
 	if ( blksize < TFTP_DEFAULT_BLKSIZE )
 		blksize = TFTP_DEFAULT_BLKSIZE;
 	snprintf ( uri_string, sizeof ( uri_string ),
-		   "tftp://%s:%d%s%s?blksize=%zd",
+		   "tftp%s://%s:%d%s%s?blksize=%zd",
+		   sizeonly ? "size" : "",
 		   inet_ntoa ( address ), ntohs ( port ),
 		   ( ( filename[0] == '/' ) ? "" : "/" ), filename, blksize );
 	DBG ( " %s", uri_string );
@@ -254,7 +256,8 @@ PXENV_EXIT_t pxenv_tftp_open ( struct s_PXENV_TFTP_OPEN *tftp_open ) {
 	if ( ( rc = pxe_tftp_open ( tftp_open->ServerIPAddress,
 				    tftp_open->TFTPPort,
 				    tftp_open->FileName,
-				    tftp_open->PacketSize ) ) != 0 ) {
+				    tftp_open->PacketSize,
+				    0) ) != 0 ) {
 		tftp_open->Status = PXENV_STATUS ( rc );
 		return PXENV_EXIT_FAILURE;
 	}
@@ -488,7 +491,7 @@ PXENV_EXIT_t pxenv_tftp_read_file ( struct s_PXENV_TFTP_READ_FILE
 
 	/* Open TFTP file */
 	if ( ( rc = pxe_tftp_open ( tftp_read_file->ServerIPAddress, 0,
-				    tftp_read_file->FileName, 0 ) ) != 0 ) {
+				    tftp_read_file->FileName, 0, 0 ) ) != 0 ) {
 		tftp_read_file->Status = PXENV_STATUS ( rc );
 		return PXENV_EXIT_FAILURE;
 	}
@@ -558,7 +561,7 @@ PXENV_EXIT_t pxenv_tftp_get_fsize ( struct s_PXENV_TFTP_GET_FSIZE
 
 	/* Open TFTP file */
 	if ( ( rc = pxe_tftp_open ( tftp_get_fsize->ServerIPAddress, 0,
-				    tftp_get_fsize->FileName, 0 ) ) != 0 ) {
+				    tftp_get_fsize->FileName, 0, 1 ) ) != 0 ) {
 		tftp_get_fsize->Status = PXENV_STATUS ( rc );
 		return PXENV_EXIT_FAILURE;
 	}
diff --git a/src/net/udp/tftp.c b/src/net/udp/tftp.c
index 810202cfc..e8d73abdb 100644
--- a/src/net/udp/tftp.c
+++ b/src/net/udp/tftp.c
@@ -133,6 +133,8 @@ enum {
 	TFTP_FL_RRQ_MULTICAST = 0x0004,
 	/** Perform MTFTP recovery on timeout */
 	TFTP_FL_MTFTP_RECOVERY = 0x0008,
+	/** Only get filesize and then abort the transfer */
+	TFTP_FL_SIZEONLY = 0x0010,
 };
 
 /** Maximum number of MTFTP open requests before falling back to TFTP */
@@ -411,6 +413,42 @@ static int tftp_send_ack ( struct tftp_request *tftp ) {
 }
 
 /**
+ * Transmit ERROR (Abort)
+ *
+ * @v tftp		TFTP connection
+ * @v errcode		TFTP error code
+ * @v errmsg		Error message string
+ * @ret rc		Return status code
+ */
+static int tftp_send_error ( struct tftp_request *tftp, int errcode,
+			     const char *errmsg ) {
+	struct tftp_error *err;
+	struct io_buffer *iobuf;
+	struct xfer_metadata meta = {
+		.dest = ( struct sockaddr * ) &tftp->peer,
+	};
+	size_t msglen;
+
+	DBGC2 ( tftp, "TFTP %p sending ERROR %d: %s\n", tftp, errcode,
+		errmsg );
+
+	/* Allocate buffer */
+	msglen = sizeof ( *err ) + strlen ( errmsg ) + 1 /* NUL */;
+	iobuf = xfer_alloc_iob ( &tftp->socket, msglen );
+	if ( ! iobuf )
+		return -ENOMEM;
+
+	/* Build ERROR */
+	err = iob_put ( iobuf, msglen );
+	err->opcode = htons ( TFTP_ERROR );
+	err->errcode = htons ( errcode );
+	strcpy ( err->errmsg, errmsg );
+
+	/* ERR always goes to the peer recorded from the RRQ response */
+	return xfer_deliver_iob_meta ( &tftp->socket, iobuf, &meta );
+}
+
+/**
  * Transmit next relevant packet
  *
  * @v tftp		TFTP connection
@@ -732,6 +770,14 @@ static int tftp_rx_oack ( struct tftp_request *tftp, void *buf, size_t len ) {
 			goto done;
 	}
 
+	/* Abort request if only trying to determine file size */
+	if ( tftp->flags & TFTP_FL_SIZEONLY ) {
+		rc = 0;
+		tftp_send_error ( tftp, TFTP_ERR_UNKNOWN_TID, "TFTP Aborted" );
+		tftp_done ( tftp, rc );
+		return rc;
+	}
+
 	/* Request next data block */
 	tftp_send_packet ( tftp );
 
@@ -759,6 +805,13 @@ static int tftp_rx_data ( struct tftp_request *tftp,
 	size_t data_len;
 	int rc;
 
+	if ( tftp->flags & TFTP_FL_SIZEONLY ) {
+		/* If we get here then server doesn't support SIZE option */
+		rc = -ENOTSUP;
+		tftp_send_error ( tftp, TFTP_ERR_UNKNOWN_TID, "TFTP Aborted" );
+		goto done;
+	}
+
 	/* Sanity check */
 	if ( iob_len ( iobuf ) < sizeof ( *data ) ) {
 		DBGC ( tftp, "TFTP %p received underlength DATA packet "
@@ -1121,6 +1174,26 @@ struct uri_opener tftp_uri_opener __uri_opener = {
 };
 
 /**
+ * Initiate TFTP-size request
+ *
+ * @v xfer		Data transfer interface
+ * @v uri		Uniform Resource Identifier
+ * @ret rc		Return status code
+ */
+static int tftpsize_open ( struct xfer_interface *xfer, struct uri *uri ) {
+	return tftp_core_open ( xfer, uri, TFTP_PORT, NULL,
+				( TFTP_FL_RRQ_SIZES |
+				  TFTP_FL_SIZEONLY ) );
+
+}
+
+/** TFTP URI opener */
+struct uri_opener tftpsize_uri_opener __uri_opener = {
+	.scheme	= "tftpsize",
+	.open	= tftpsize_open,
+};
+
+/**
  * Initiate TFTM download
  *
  * @v xfer		Data transfer interface