diff --git a/src/core/image.c b/src/core/image.c index bf9bb7fa..666ee3de 100644 --- a/src/core/image.c +++ b/src/core/image.c @@ -194,6 +194,10 @@ int register_image ( struct image *image ) { */ void unregister_image ( struct image *image ) { + /* Do nothing unless image is registered */ + if ( ! ( image->flags & IMAGE_REGISTERED ) ) + return; + DBGC ( image, "IMAGE %s unregistered\n", image->name ); list_del ( &image->list ); image->flags &= ~IMAGE_REGISTERED; @@ -259,23 +263,13 @@ int image_probe ( struct image *image ) { */ int image_exec ( struct image *image ) { struct image *saved_current_image; - struct image *replacement; + struct image *replacement = NULL; struct uri *old_cwuri; int rc; /* Sanity check */ assert ( image->flags & IMAGE_REGISTERED ); - /* Check that this image can be selected for execution */ - if ( ( rc = image_select ( image ) ) != 0 ) - return rc; - - /* Check that image is trusted (if applicable) */ - if ( require_trusted_images && ! ( image->flags & IMAGE_TRUSTED ) ) { - DBGC ( image, "IMAGE %s is not trusted\n", image->name ); - return -EACCES_UNTRUSTED; - } - /* Switch current working directory to be that of the image itself */ old_cwuri = uri_get ( cwuri ); churi ( image->uri ); @@ -289,6 +283,17 @@ int image_exec ( struct image *image ) { */ current_image = image_get ( image ); + /* Check that this image can be selected for execution */ + if ( ( rc = image_select ( image ) ) != 0 ) + goto err; + + /* Check that image is trusted (if applicable) */ + if ( require_trusted_images && ! ( image->flags & IMAGE_TRUSTED ) ) { + DBGC ( image, "IMAGE %s is not trusted\n", image->name ); + rc = -EACCES_UNTRUSTED; + goto err; + } + /* Record boot attempt */ syslog ( LOG_NOTICE, "Executing \"%s\"\n", image->name ); @@ -317,6 +322,11 @@ int image_exec ( struct image *image ) { if ( replacement ) assert ( replacement->flags & IMAGE_REGISTERED ); + err: + /* Unregister image if applicable */ + if ( image->flags & IMAGE_AUTO_UNREGISTER ) + unregister_image ( image ); + /* Drop temporary reference to the original image */ image_put ( image ); diff --git a/src/hci/commands/image_cmd.c b/src/hci/commands/image_cmd.c index 1ae33074..38d814cd 100644 --- a/src/hci/commands/image_cmd.c +++ b/src/hci/commands/image_cmd.c @@ -38,19 +38,24 @@ FILE_LICENCE ( GPL2_OR_LATER ); struct imgsingle_options { /** Image name */ const char *name; + /** Free image after execution */ + int autofree; }; /** "img{single}" option list */ static struct option_descriptor imgsingle_opts[] = { OPTION_DESC ( "name", 'n', required_argument, struct imgsingle_options, name, parse_string ), + OPTION_DESC ( "autofree", 'a', no_argument, + struct imgsingle_options, autofree, parse_flag ), }; /** "img{single}" command descriptor */ static struct command_descriptor imgsingle_cmd = COMMAND_DESC ( struct imgsingle_options, imgsingle_opts, 1, MAX_ARGUMENTS, - "[--name ] [...]" ); + "[--name ] [--autofree] " + " [...]" ); /** An "img{single}" family command descriptor */ struct imgsingle_descriptor { @@ -134,6 +139,10 @@ static int imgsingle_exec ( int argc, char **argv, } } + /* Set the auto-unregister flag, if applicable */ + if ( opts.autofree ) + image->flags |= IMAGE_AUTO_UNREGISTER; + /* Carry out command action, if applicable */ if ( desc->action ) { if ( ( rc = desc->action ( image ) ) != 0 ) { @@ -160,7 +169,7 @@ static int imgsingle_exec ( int argc, char **argv, static struct command_descriptor imgfetch_cmd = COMMAND_DESC ( struct imgsingle_options, imgsingle_opts, 1, MAX_ARGUMENTS, - "[--name ] [...]" ); + "[--name ] [--autofree] [...]" ); /** "imgfetch" family command descriptor */ struct imgsingle_descriptor imgfetch_desc = { @@ -202,7 +211,7 @@ static int imgselect_exec ( int argc, char **argv ) { static struct command_descriptor imgexec_cmd = COMMAND_DESC ( struct imgsingle_options, imgsingle_opts, 0, MAX_ARGUMENTS, - "[--name ] [ [...]]" ); + "[--autofree] [ [...]]" ); /** "imgexec" family command descriptor */ struct imgsingle_descriptor imgexec_desc = { diff --git a/src/include/ipxe/image.h b/src/include/ipxe/image.h index ac97137b..6022dce6 100644 --- a/src/include/ipxe/image.h +++ b/src/include/ipxe/image.h @@ -67,6 +67,9 @@ struct image { /** Image is trusted */ #define IMAGE_TRUSTED 0x0004 +/** Image will be automatically unregistered after execution */ +#define IMAGE_AUTO_UNREGISTER 0x0008 + /** An executable image type */ struct image_type { /** Name of this image type */ diff --git a/src/usr/imgmgmt.c b/src/usr/imgmgmt.c index 2c74f486..ce3bc90f 100644 --- a/src/usr/imgmgmt.c +++ b/src/usr/imgmgmt.c @@ -146,6 +146,8 @@ void imgstat ( struct image *image ) { printf ( " [TRUSTED]" ); if ( image->flags & IMAGE_SELECTED ) printf ( " [SELECTED]" ); + if ( image->flags & IMAGE_AUTO_UNREGISTER ) + printf ( " [AUTOFREE]" ); if ( image->cmdline ) printf ( " \"%s\"", image->cmdline ); printf ( "\n" );