[bzimage] Allow initrds to be rearranged in place

At present, loading a bzImage via iPXE requires enough RAM to hold two
copies of each initrd file.  Remove this constraint by rearranging the
initrds in place.

Signed-off-by: Michael Brown <mcb30@ipxe.org>

src/arch/i386/image/bzimage.c

 #include <assert.h>
 #include <realmode.h>
 #include <bzimage.h>
+#include <initrd.h>
 #include <ipxe/uaccess.h>
 #include <ipxe/image.h>
 #include <ipxe/segment.h>
  * @v image		bzImage image
  * @v bzimg		bzImage context
  * @v cmdline		Kernel command line
- * @ret rc		Return status code
  */
-static int bzimage_set_cmdline ( struct image *image,
-				 struct bzimage_context *bzimg,
-				 const char *cmdline ) {
+static void bzimage_set_cmdline ( struct image *image,
+				  struct bzimage_context *bzimg,
+				  const char *cmdline ) {
 	size_t cmdline_len;
 
 	/* Copy command line down to real-mode portion */
 	copy_to_user ( bzimg->rm_kernel, bzimg->rm_cmdline,
 		       cmdline, cmdline_len );
 	DBGC ( image, "bzImage %p command line \"%s\"\n", image, cmdline );
-
-	return 0;
 }
 
 /**
  * @v image		bzImage image
  * @v initrd		initrd image
  * @v address		Address at which to load, or UNULL
- * @ret len		Length of loaded image, rounded up to 4 bytes
+ * @ret len		Length of loaded image, rounded up to INITRD_ALIGN
  */
 static size_t bzimage_load_initrd ( struct image *image,
 				    struct image *initrd,
 	struct cpio_header cpio;
         size_t offset = 0;
 	size_t name_len;
+	size_t pad_len;
 
 	/* Do not include kernel image itself as an initrd */
 	if ( initrd == image )
 
 	/* Create cpio header before non-prebuilt images */
 	if ( filename && filename[0] ) {
-		DBGC ( image, "bzImage %p inserting initrd %p as %s\n",
-		       image, initrd, filename );
 		cmdline = strchr ( filename, ' ' );
 		name_len = ( ( cmdline ? ( ( size_t ) ( cmdline - filename ) )
 			       : strlen ( filename ) ) + 1 /* NUL */ );
 
 	/* Copy in initrd image body */
 	if ( address )
-		memcpy_user ( address, offset, initrd->data, 0, initrd->len );
-	offset += initrd->len;
+		memmove_user ( address, offset, initrd->data, 0, initrd->len );
 	if ( address ) {
-		DBGC ( image, "bzImage %p has initrd %p at [%lx,%lx)\n",
-		       image, initrd, user_to_phys ( address, 0 ),
-		       user_to_phys ( address, offset ) );
+		DBGC ( image, "bzImage %p initrd %p [%#08lx,%#08lx,%#08lx)"
+		       "%s%s\n", image, initrd, user_to_phys ( address, 0 ),
+		       user_to_phys ( address, offset ),
+		       user_to_phys ( address, ( offset + initrd->len ) ),
+		       ( filename ? " " : "" ), ( filename ? filename : "" ) );
+		DBGC2_MD5A ( image, user_to_phys ( address, offset ),
+			     user_to_virt ( address, offset ), initrd->len );
 	}
+	offset += initrd->len;
+
+	/* Round up to multiple of INITRD_ALIGN and zero-pad */
+	pad_len = ( ( -offset ) & ( INITRD_ALIGN - 1 ) );
+	if ( address )
+		memset_user ( address, offset, 0, pad_len );
+	offset += pad_len;
 
-	/* Round up to 4-byte boundary */
-	offset = ( ( offset + 0x03 ) & ~0x03 );
 	return offset;
 }
 
 /**
- * Load initrds, if any
+ * Check that initrds can be loaded
  *
  * @v image		bzImage image
  * @v bzimg		bzImage context
  * @ret rc		Return status code
  */
-static int bzimage_load_initrds ( struct image *image,
-				  struct bzimage_context *bzimg ) {
+static int bzimage_check_initrds ( struct image *image,
+				   struct bzimage_context *bzimg ) {
 	struct image *initrd;
-	size_t total_len = 0;
-	physaddr_t address;
+	userptr_t bottom;
+	size_t len = 0;
 	int rc;
 
-	/* Add up length of all initrd images */
-	for_each_image ( initrd )
-		total_len += bzimage_load_initrd ( image, initrd, UNULL );
+	/* Calculate total loaded length of initrds */
+	for_each_image ( initrd ) {
 
-	/* Give up if no initrd images found */
-	if ( ! total_len )
-		return 0;
+		/* Skip kernel */
+		if ( initrd == image )
+			continue;
+
+		/* Calculate length */
+		len += bzimage_load_initrd ( image, initrd, UNULL );
 
-	/* Find a suitable start address.  Try 1MB boundaries,
-	 * starting from the downloaded kernel image itself and
-	 * working downwards until we hit an available region.
+		DBGC ( image, "bzImage %p initrd %p from [%#08lx,%#08lx)%s%s\n",
+		       image, initrd, user_to_phys ( initrd->data, 0 ),
+		       user_to_phys ( initrd->data, initrd->len ),
+		       ( initrd->cmdline ? " " : "" ),
+		       ( initrd->cmdline ? initrd->cmdline : "" ) );
+		DBGC2_MD5A ( image, user_to_phys ( initrd->data, 0 ),
+			     user_to_virt ( initrd->data, 0 ), initrd->len );
+	}
+
+	/* Calculate lowest usable address */
+	bottom = userptr_add ( bzimg->pm_kernel, bzimg->pm_sz );
+
+	/* Check that total length fits within space available for
+	 * reshuffling.  This is a conservative check, since CPIO
+	 * headers are not present during reshuffling, but this
+	 * doesn't hurt and keeps the code simple.
 	 */
-	for ( address = ( user_to_phys ( image->data, 0 ) & ~0xfffff ) ; ;
-	      address -= 0x100000 ) {
-		/* Check that we're not going to overwrite the
-		 * kernel itself.  This check isn't totally
-		 * accurate, but errs on the side of caution.
-		 */
-		if ( address <= ( BZI_LOAD_HIGH_ADDR + image->len ) ) {
-			DBGC ( image, "bzImage %p could not find a location "
-			       "for initrd\n", image );
-			return -ENOBUFS;
-		}
-		/* Check that we are within the kernel's range */
-		if ( ( address + total_len - 1 ) > bzimg->mem_limit )
-			continue;
-		/* Prepare and verify segment */
-		if ( ( rc = prep_segment ( phys_to_user ( address ), 0,
-					   total_len ) ) != 0 )
-			continue;
-		/* Use this address */
-		break;
+	if ( ( rc = initrd_reshuffle_check ( len, bottom ) ) != 0 ) {
+		DBGC ( image, "bzImage %p failed reshuffle check: %s\n",
+		       image, strerror ( rc ) );
+		return rc;
 	}
 
-	/* Record initrd location */
-	bzimg->ramdisk_image = address;
-	bzimg->ramdisk_size = total_len;
+	/* Check that total length fits within kernel's memory limit */
+	if ( user_to_phys ( bottom, len ) > bzimg->mem_limit ) {
+		DBGC ( image, "bzImage %p not enough space for initrds\n",
+		       image );
+		return -ENOBUFS;
+	}
+
+	return 0;
+}
+
+/**
+ * Load initrds, if any
+ *
+ * @v image		bzImage image
+ * @v bzimg		bzImage context
+ */
+static void bzimage_load_initrds ( struct image *image,
+				   struct bzimage_context *bzimg ) {
+	struct image *initrd;
+	struct image *highest = NULL;
+	struct image *other;
+	userptr_t top;
+	userptr_t dest;
+	size_t len;
+
+	/* Reshuffle initrds into desired order */
+	initrd_reshuffle ( userptr_add ( bzimg->pm_kernel, bzimg->pm_sz ) );
 
-	/* Construct initrd */
-	DBGC ( image, "bzImage %p constructing initrd at [%lx,%lx)\n",
-	       image, address, ( address + total_len ) );
+	/* Find highest initrd */
 	for_each_image ( initrd ) {
-		address += bzimage_load_initrd ( image, initrd,
-						 phys_to_user ( address ) );
+		if ( ( highest == NULL ) ||
+		     ( userptr_sub ( initrd->data, highest->data ) > 0 ) ) {
+			highest = initrd;
+		}
 	}
 
-	return 0;
+	/* Do nothing if there are no initrds */
+	if ( ! highest )
+		return;
+
+	/* Find highest usable address */
+	top = userptr_add ( highest->data,
+			    ( ( highest->len + INITRD_ALIGN - 1 ) &
+			      ~( INITRD_ALIGN - 1 ) ) );
+	if ( user_to_phys ( top, 0 ) > bzimg->mem_limit )
+		top = phys_to_user ( bzimg->mem_limit );
+	DBGC ( image, "bzImage %p loading initrds from %#08lx downwards\n",
+	       image, user_to_phys ( top, 0 ) );
+
+	/* Load initrds in order */
+	for_each_image ( initrd ) {
+
+		/* Calculate cumulative length of following
+		 * initrds (including padding).
+		 */
+		len = 0;
+		for_each_image ( other ) {
+			if ( other == initrd )
+				len = 0;
+			len += bzimage_load_initrd ( image, other, UNULL );
+		}
+
+		/* Load initrd at this address */
+		dest = userptr_add ( top, -len );
+		bzimage_load_initrd ( image, initrd, dest );
+
+		/* Record initrd location */
+		if ( ! bzimg->ramdisk_image ) {
+			bzimg->ramdisk_image = user_to_phys ( dest, 0 );
+			bzimg->ramdisk_size = len;
+		}
+	}
 }
 
 /**
 		return rc;
 	}
 
+	/* Parse command line for bootloader parameters */
+	if ( ( rc = bzimage_parse_cmdline ( image, &bzimg, cmdline ) ) != 0)
+		return rc;
+
+	/* Check that initrds can be loaded */
+	if ( ( rc = bzimage_check_initrds ( image, &bzimg ) ) != 0 )
+		return rc;
+
+	/* Remove kernel from image list (without invalidating image pointer) */
+	unregister_image ( image_get ( image ) );
+
 	/* Load segments */
 	memcpy_user ( bzimg.rm_kernel, 0, image->data,
 		      0, bzimg.rm_filesz );
 	memcpy_user ( bzimg.pm_kernel, 0, image->data,
 		      bzimg.rm_filesz, bzimg.pm_sz );
 
-	/* Update and write out header */
-	bzimage_update_header ( image, &bzimg, bzimg.rm_kernel );
-
-	/* Parse command line for bootloader parameters */
-	if ( ( rc = bzimage_parse_cmdline ( image, &bzimg, cmdline ) ) != 0)
-		return rc;
-
 	/* Store command line */
-	if ( ( rc = bzimage_set_cmdline ( image, &bzimg, cmdline ) ) != 0 )
-		return rc;
+	bzimage_set_cmdline ( image, &bzimg, cmdline );
+
+	/* Prepare for exiting.  Must do this before loading initrds,
+	 * since loading the initrds will corrupt the external heap.
+	 */
+	shutdown_boot();
 
 	/* Load any initrds */
-	if ( ( rc = bzimage_load_initrds ( image, &bzimg ) ) != 0 )
-		return rc;
+	bzimage_load_initrds ( image, &bzimg );
 
 	/* Update kernel header */
 	bzimage_update_header ( image, &bzimg, bzimg.rm_kernel );
 
-	/* Prepare for exiting */
-	shutdown_boot();
-
 	DBGC ( image, "bzImage %p jumping to RM kernel at %04x:0000 "
 	       "(stack %04x:%04zx)\n", image, ( bzimg.rm_kernel_seg + 0x20 ),
 	       bzimg.rm_kernel_seg, bzimg.rm_heap );