[script] Allow "exit" to exit a script

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

4 files changed, 92 insertions, 45 deletions

diff --git a/src/core/exec.c b/src/core/exec.c
index 20b97a6d..60d8cb97 100644
--- a/src/core/exec.c
+++ b/src/core/exec.c
@@ -29,6 +29,7 @@ FILE_LICENCE ( GPL2_OR_LATER );
 #include <assert.h>
 #include <ipxe/tables.h>
 #include <ipxe/command.h>
+#include <ipxe/parseopt.h>
 #include <ipxe/settings.h>
 
 /** @file
@@ -37,9 +38,8 @@ FILE_LICENCE ( GPL2_OR_LATER );
  *
  */
 
-/* Avoid dragging in getopt.o unless a command really uses it */
-int optind;
-int nextchar;
+/** Shell exit flag */
+int shell_exit;
 
 /**
  * Execute command
@@ -270,6 +270,9 @@ int system ( const char *command ) {
 	int count;
 	int rc = 0;
 
+	/* Reset exit flag */
+	shell_exit = 0;
+
 	/* Perform variable expansion */
 	expcmd = expand_command ( command );
 	if ( ! expcmd )
@@ -294,6 +297,10 @@ int system ( const char *command ) {
 			argv[argc] = NULL;
 			rc = execv ( argv[0], argv );
 
+			/* Check exit flag */
+			if ( shell_exit )
+				break;
+
 			/* Handle terminator */
 			if ( terminator ( rc ) )
 				break;
@@ -307,7 +314,7 @@ int system ( const char *command ) {
 }
 
 /**
- * The "echo" command
+ * "echo" command
  *
  * @v argc		Argument count
  * @v argv		Argument list
@@ -328,3 +335,49 @@ struct command echo_command __command = {
 	.name = "echo",
 	.exec = echo_exec,
 };
+
+/** "exit" options */
+struct exit_options {};
+
+/** "exit" option list */
+static struct option_descriptor exit_opts[] = {};
+
+/** "exit" command descriptor */
+static struct command_descriptor exit_cmd =
+	COMMAND_DESC ( struct exit_options, exit_opts, 0, 1,
+		       "[<status>]", "" );
+
+/**
+ * "exit" command
+ *
+ * @v argc		Argument count
+ * @v argv		Argument list
+ * @ret rc		Return status code
+ */
+static int exit_exec ( int argc, char **argv ) {
+	struct exit_options opts;
+	unsigned int exit_code = 0;
+	int rc;
+
+	/* Parse options */
+	if ( ( rc = parse_options ( argc, argv, &exit_cmd, &opts ) ) != 0 )
+		return rc;
+
+	/* Parse exit status, if present */
+	if ( optind != argc ) {
+		if ( ( rc = parse_integer ( argv[optind], &exit_code ) ) != 0 )
+			return rc;
+	}
+
+	/* Set exit flag */
+	shell_exit = 1;
+
+	return exit_code;
+}
+
+/** "exit" command */
+struct command exit_command __command = {
+	.name = "exit",
+	.exec = exit_exec,
+};
+
diff --git a/src/hci/shell.c b/src/hci/shell.c
index 3fc25727..f9cd3af2 100644
--- a/src/hci/shell.c
+++ b/src/hci/shell.c
@@ -21,8 +21,10 @@ FILE_LICENCE ( GPL2_OR_LATER );
 #include <stdint.h>
 #include <stdlib.h>
 #include <stdio.h>
+#include <getopt.h>
 #include <readline/readline.h>
 #include <ipxe/command.h>
+#include <ipxe/parseopt.h>
 #include <ipxe/shell.h>
 
 /** @file
@@ -34,29 +36,13 @@ FILE_LICENCE ( GPL2_OR_LATER );
 /** The shell prompt string */
 static const char shell_prompt[] = "iPXE> ";
 
-/** Flag set in order to exit shell */
-static int exit_flag = 0;
-
-/** "exit" command body */
-static int exit_exec ( int argc, char **argv __unused ) {
-
-	if ( argc == 1 ) {
-		exit_flag = 1;
-	} else {
-		printf ( "Usage: exit\n"
-			 "Exits the command shell\n" );
-	}
-
-	return 0;
-}
-
-/** "exit" command definition */
-struct command exit_command __command = {
-	.name = "exit",
-	.exec = exit_exec,
-};
-
-/** "help" command body */
+/**
+ * "help" command
+ *
+ * @v argc		Argument count
+ * @v argv		Argument list
+ * @ret rc		Return status code
+ */
 static int help_exec ( int argc __unused, char **argv __unused ) {
 	struct command *command;
 	unsigned int hpos = 0;
@@ -78,7 +64,7 @@ static int help_exec ( int argc __unused, char **argv __unused ) {
 	return 0;
 }
 
-/** "help" command definition */
+/** "help" command */
 struct command help_command __command = {
 	.name = "help",
 	.exec = help_exec,
@@ -91,12 +77,11 @@ struct command help_command __command = {
 void shell ( void ) {
 	char *line;
 
-	exit_flag = 0;
-	while ( ! exit_flag ) {
+	do {
 		line = readline ( shell_prompt );
 		if ( line ) {
 			system ( line );
 			free ( line );
 		}
-	}
+	} while ( shell_exit == 0 );
 }
diff --git a/src/image/script.c b/src/image/script.c
index ba098c2c..f9d5a393 100644
--- a/src/image/script.c
+++ b/src/image/script.c
@@ -99,23 +99,20 @@ static int process_script ( int ( * process_line ) ( const char *line ),
 }
 
 /**
- * Terminate script processing if line processing failed
+ * Terminate script processing on shell exit or command failure
  *
  * @v rc		Line processing status
  * @ret terminate	Terminate script processing
  */
-static int terminate_on_failure ( int rc ) {
-	return ( rc != 0 );
-}
+static int terminate_on_exit_or_failure ( int rc ) {
 
-/**
- * Terminate script processing if line processing succeeded
- *
- * @v rc		Line processing status
- * @ret terminate	Terminate script processing
- */
-static int terminate_on_success ( int rc ) {
-	return ( rc == 0 );
+	/* Check and consume exit flag */
+	if ( shell_exit ) {
+		shell_exit = 0;
+		return 1;
+	}
+
+	return ( rc != 0 );
 }
 
 /**
@@ -164,7 +161,7 @@ static int script_exec ( struct image *image ) {
 	script = image;
 
 	/* Process script */
-	rc = process_script ( script_exec_line, terminate_on_failure );
+	rc = process_script ( script_exec_line, terminate_on_exit_or_failure );
 
 	/* Restore saved state, re-register image, and return */
 	script_offset = saved_offset;
@@ -253,6 +250,16 @@ static int goto_find_label ( const char *line ) {
 }
 
 /**
+ * Terminate script processing when label is found
+ *
+ * @v rc		Line processing status
+ * @ret terminate	Terminate script processing
+ */
+static int terminate_on_label_found ( int rc ) {
+	return ( rc == 0 );
+}
+
+/**
  * "goto" command
  *
  * @v argc		Argument count
@@ -280,7 +287,7 @@ static int goto_exec ( int argc, char **argv ) {
 	/* Find label */
 	saved_offset = script_offset;
 	if ( ( rc = process_script ( goto_find_label,
-				     terminate_on_success ) ) != 0 ) {
+				     terminate_on_label_found ) ) != 0 ) {
 		script_offset = saved_offset;
 		return rc;
 	}
diff --git a/src/include/ipxe/command.h b/src/include/ipxe/command.h
index a7d0ae46..2486e7aa 100644
--- a/src/include/ipxe/command.h
+++ b/src/include/ipxe/command.h
@@ -23,4 +23,6 @@ struct command {
 
 #define __command __table_entry ( COMMANDS, 01 )
 
+extern int shell_exit;
+
 #endif /* _IPXE_COMMAND_H */