[tftp] Make TFTP size requests abort transfer with an error
[people/cooldavid/gpxe.git] / src / net / udp / tftp.c
index 810202c..e8d73ab 100644 (file)
@@ -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 */
@@ -410,6 +412,42 @@ static int tftp_send_ack ( struct tftp_request *tftp ) {
        return xfer_deliver_iob_meta ( &tftp->socket, iobuf, &meta );
 }
 
+/**
+ * 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
  *
@@ -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 "
@@ -1120,6 +1173,26 @@ struct uri_opener tftp_uri_opener __uri_opener = {
        .open   = tftp_open,
 };
 
+/**
+ * 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
  *