2 files changed, 125 insertions, 68 deletions

diff --git a/src/include/ipxe/tls.h b/src/include/ipxe/tls.h
index 6fcb69b..30bb1c4 100644
--- a/src/include/ipxe/tls.h
+++ b/src/include/ipxe/tls.h
@@ -52,6 +52,9 @@ struct tls_header {
 /** Change cipher content type */
 #define TLS_TYPE_CHANGE_CIPHER 20
 
+/** Change cipher spec magic byte */
+#define TLS_CHANGE_CIPHER_SPEC 1
+
 /** Alert content type */
 #define TLS_TYPE_ALERT 21
 
@@ -395,6 +398,8 @@ struct tls_connection {
 	struct io_buffer rx_header_iobuf;
 	/** List of received data buffers */
 	struct list_head rx_data;
+	/** Received handshake fragment */
+	struct io_buffer *rx_handshake;
 };
 
 /** RX I/O buffer size
diff --git a/src/net/tls.c b/src/net/tls.c
index 8996296..e774d9f 100644
--- a/src/net/tls.c
+++ b/src/net/tls.c
@@ -388,6 +388,7 @@ static void free_tls ( struct refcnt *refcnt ) {
 		list_del ( &iobuf->list );
 		free_iob ( iobuf );
 	}
+	free_iob ( tls->rx_handshake );
 	x509_chain_put ( tls->certs );
 	x509_chain_put ( tls->chain );
 	x509_root_put ( tls->root );
@@ -1731,20 +1732,27 @@ static int tls_send_finished ( struct tls_connection *tls ) {
  * Receive new Change Cipher record
  *
  * @v tls		TLS connection
- * @v data		Plaintext record
- * @v len		Length of plaintext record
+ * @v iobuf		I/O buffer
  * @ret rc		Return status code
  */
 static int tls_new_change_cipher ( struct tls_connection *tls,
-				   const void *data, size_t len ) {
+				   struct io_buffer *iobuf ) {
+	const struct {
+		uint8_t spec;
+	} __attribute__ (( packed )) *change = iobuf->data;
+	size_t len = iob_len ( iobuf );
 	int rc;
 
-	if ( ( len != 1 ) || ( *( ( uint8_t * ) data ) != 1 ) ) {
+	/* Sanity check */
+	if ( ( sizeof ( *change ) != len ) ||
+	     ( change->spec != TLS_CHANGE_CIPHER_SPEC ) ) {
 		DBGC ( tls, "TLS %p received invalid Change Cipher\n", tls );
-		DBGC_HD ( tls, data, len );
+		DBGC_HD ( tls, change, len );
 		return -EINVAL_CHANGE_CIPHER;
 	}
+	iob_pull ( iobuf, sizeof ( *change ) );
 
+	/* Change receive cipher spec */
 	if ( ( rc = tls_change_cipher ( tls, &tls->rx_cipherspec_pending,
 					&tls->rx_cipherspec ) ) != 0 ) {
 		DBGC ( tls, "TLS %p could not activate RX cipher: %s\n",
@@ -1760,25 +1768,27 @@ static int tls_new_change_cipher ( struct tls_connection *tls,
  * Receive new Alert record
  *
  * @v tls		TLS connection
- * @v data		Plaintext record
- * @v len		Length of plaintext record
+ * @v iobuf		I/O buffer
  * @ret rc		Return status code
  */
-static int tls_new_alert ( struct tls_connection *tls, const void *data,
-			   size_t len ) {
+static int tls_new_alert ( struct tls_connection *tls,
+			   struct io_buffer *iobuf ) {
 	const struct {
 		uint8_t level;
 		uint8_t description;
 		char next[0];
-	} __attribute__ (( packed )) *alert = data;
+	} __attribute__ (( packed )) *alert = iobuf->data;
+	size_t len = iob_len ( iobuf );
 
 	/* Sanity check */
 	if ( sizeof ( *alert ) != len ) {
 		DBGC ( tls, "TLS %p received overlength Alert\n", tls );
-		DBGC_HD ( tls, data, len );
+		DBGC_HD ( tls, alert, len );
 		return -EINVAL_ALERT;
 	}
+	iob_pull ( iobuf, sizeof ( *alert ) );
 
+	/* Handle alert */
 	switch ( alert->level ) {
 	case TLS_ALERT_WARNING:
 		DBGC ( tls, "TLS %p received warning alert %d\n",
@@ -2392,38 +2402,33 @@ static int tls_new_finished ( struct tls_connection *tls,
  * Receive new Handshake record
  *
  * @v tls		TLS connection
- * @v data		Plaintext record
- * @v len		Length of plaintext record
+ * @v iobuf		I/O buffer
  * @ret rc		Return status code
  */
 static int tls_new_handshake ( struct tls_connection *tls,
-			       const void *data, size_t len ) {
-	size_t remaining = len;
+			       struct io_buffer *iobuf ) {
+	size_t remaining;
 	int rc;
 
-	while ( remaining ) {
+	while ( ( remaining = iob_len ( iobuf ) ) ) {
 		const struct {
 			uint8_t type;
 			tls24_t length;
 			uint8_t payload[0];
-		} __attribute__ (( packed )) *handshake = data;
+		} __attribute__ (( packed )) *handshake = iobuf->data;
 		const void *payload;
 		size_t payload_len;
 		size_t record_len;
 
 		/* Parse header */
 		if ( sizeof ( *handshake ) > remaining ) {
-			DBGC ( tls, "TLS %p received underlength Handshake\n",
-			       tls );
-			DBGC_HD ( tls, data, remaining );
-			return -EINVAL_HANDSHAKE;
+			/* Leave remaining fragment unconsumed */
+			break;
 		}
 		payload_len = tls_uint24 ( &handshake->length );
 		if ( payload_len > ( remaining - sizeof ( *handshake ) ) ) {
-			DBGC ( tls, "TLS %p received overlength Handshake\n",
-			       tls );
-			DBGC_HD ( tls, data, len );
-			return -EINVAL_HANDSHAKE;
+			/* Leave remaining fragment unconsumed */
+			break;
 		}
 		payload = &handshake->payload;
 		record_len = ( sizeof ( *handshake ) + payload_len );
@@ -2470,65 +2475,87 @@ static int tls_new_handshake ( struct tls_connection *tls,
 		 * which are explicitly excluded).
 		 */
 		if ( handshake->type != TLS_HELLO_REQUEST )
-			tls_add_handshake ( tls, data, record_len );
+			tls_add_handshake ( tls, iobuf->data, record_len );
 
 		/* Abort on failure */
 		if ( rc != 0 )
 			return rc;
 
 		/* Move to next handshake record */
-		data += record_len;
-		remaining -= record_len;
+		iob_pull ( iobuf, record_len );
 	}
 
 	return 0;
 }
 
 /**
- * Receive new record
+ * Receive new unknown record
+ *
+ * @v tls		TLS connection
+ * @v iobuf		I/O buffer
+ * @ret rc		Return status code
+ */
+static int tls_new_unknown ( struct tls_connection *tls __unused,
+			     struct io_buffer *iobuf ) {
+
+	/* RFC4346 says that we should just ignore unknown record types */
+	iob_pull ( iobuf, iob_len ( iobuf ) );
+	return 0;
+}
+
+/**
+ * Receive new data record
  *
  * @v tls		TLS connection
- * @v type		Record type
  * @v rx_data		List of received data buffers
  * @ret rc		Return status code
  */
-static int tls_new_record ( struct tls_connection *tls, unsigned int type,
-			    struct list_head *rx_data ) {
+static int tls_new_data ( struct tls_connection *tls,
+			  struct list_head *rx_data ) {
 	struct io_buffer *iobuf;
-	int ( * handler ) ( struct tls_connection *tls, const void *data,
-			    size_t len );
 	int rc;
 
-	/* Deliver data records to the plainstream interface */
-	if ( type == TLS_TYPE_DATA ) {
-
-		/* Fail unless we are ready to receive data */
-		if ( ! tls_ready ( tls ) )
-			return -ENOTCONN;
-
-		/* Deliver each I/O buffer in turn */
-		while ( ( iobuf = list_first_entry ( rx_data, struct io_buffer,
-						     list ) ) ) {
-			list_del ( &iobuf->list );
-			if ( ( rc = xfer_deliver_iob ( &tls->plainstream,
-						       iobuf ) ) != 0 ) {
-				DBGC ( tls, "TLS %p could not deliver data: "
-				       "%s\n", tls, strerror ( rc ) );
-				return rc;
-			}
+	/* Fail unless we are ready to receive data */
+	if ( ! tls_ready ( tls ) )
+		return -ENOTCONN;
+
+	/* Deliver each I/O buffer in turn */
+	while ( ( iobuf = list_first_entry ( rx_data, struct io_buffer,
+					     list ) ) ) {
+		list_del ( &iobuf->list );
+		if ( ( rc = xfer_deliver_iob ( &tls->plainstream,
+					       iobuf ) ) != 0 ) {
+			DBGC ( tls, "TLS %p could not deliver data: "
+			       "%s\n", tls, strerror ( rc ) );
+			return rc;
 		}
-		return 0;
 	}
 
-	/* For all other records, merge into a single I/O buffer */
-	iobuf = iob_concatenate ( rx_data );
-	if ( ! iobuf ) {
-		DBGC ( tls, "TLS %p could not concatenate non-data record "
-		       "type %d\n", tls, type );
-		return -ENOMEM_RX_CONCAT;
-	}
+	return 0;
+}
+
+/**
+ * Receive new record
+ *
+ * @v tls		TLS connection
+ * @v type		Record type
+ * @v rx_data		List of received data buffers
+ * @ret rc		Return status code
+ */
+static int tls_new_record ( struct tls_connection *tls, unsigned int type,
+			    struct list_head *rx_data ) {
+	int ( * handler ) ( struct tls_connection *tls,
+			    struct io_buffer *iobuf );
+	struct io_buffer *tmp = NULL;
+	struct io_buffer **iobuf;
+	int rc;
 
-	/* Determine handler */
+	/* Deliver data records as-is to the plainstream interface */
+	if ( type == TLS_TYPE_DATA )
+		return tls_new_data ( tls, rx_data );
+
+	/* Determine handler and fragment buffer */
+	iobuf = &tmp;
 	switch ( type ) {
 	case TLS_TYPE_CHANGE_CIPHER:
 		handler = tls_new_change_cipher;
@@ -2538,19 +2565,44 @@ static int tls_new_record ( struct tls_connection *tls, unsigned int type,
 		break;
 	case TLS_TYPE_HANDSHAKE:
 		handler = tls_new_handshake;
+		iobuf = &tls->rx_handshake;
 		break;
 	default:
-		/* RFC4346 says that we should just ignore unknown
-		 * record types.
-		 */
-		handler = NULL;
-		DBGC ( tls, "TLS %p ignoring record type %d\n", tls, type );
+		DBGC ( tls, "TLS %p unknown record type %d\n", tls, type );
+		handler = tls_new_unknown;
 		break;
 	}
 
-	/* Handle record and free I/O buffer */
-	rc = ( handler ? handler ( tls, iobuf->data, iob_len ( iobuf ) ) : 0 );
-	free_iob ( iobuf );
+	/* Merge into a single I/O buffer */
+	if ( *iobuf )
+		list_add ( &(*iobuf)->list, rx_data );
+	*iobuf = iob_concatenate ( rx_data );
+	if ( ! ( *iobuf ) ) {
+		DBGC ( tls, "TLS %p could not concatenate non-data record "
+		       "type %d\n", tls, type );
+		rc = -ENOMEM_RX_CONCAT;
+		goto err_concatenate;
+	}
+
+	/* Handle record */
+	if ( ( rc = handler ( tls, *iobuf ) ) != 0 )
+		goto err_handle;
+
+	/* Discard I/O buffer if empty */
+	if ( ! iob_len ( *iobuf ) ) {
+		free_iob ( *iobuf );
+		*iobuf = NULL;
+	}
+
+	/* Sanity check */
+	assert ( tmp == NULL );
+
+	return 0;
+
+ err_handle:
+	free_iob ( *iobuf );
+	*iobuf = NULL;
+ err_concatenate:
 	return rc;
 }