bootstd: Add a function to update a command line

The Linux command line consists of a number of words with optional values.
At present U-Boot allows this to be changed using the bootargs environment
variable.

But this is quite painful, since the command line can be very long.

Add a function which can adjust a single field in the command line, so
that it is easier to make changes before booting.

Signed-off-by: Simon Glass <sjg@chromium.org>
Reviewed-by: Bin Meng <bmeng.cn@gmail.com>

boot/bootflow.c

@@ -610,3 +610,194 @@ static int on_bootargs(const char *name, const char *value, enum env_op op,
 }
 U_BOOT_ENV_CALLBACK(bootargs, on_bootargs);
 #endif
+
+/**
+ * copy_in() - Copy a string into a cmdline buffer
+ *
+ * @buf: Buffer to copy into
+ * @end: End of buffer (pointer to char after the end)
+ * @arg: String to copy from
+ * @len: Number of chars to copy from @arg (note that this is not usually the
+ * sane as strlen(arg) since the string may contain following arguments)
+ * @new_val: Value to put after arg, or BOOTFLOWCL_EMPTY to use an empty value
+ * with no '=' sign
+ * Returns: Number of chars written to @buf
+ */
+static int copy_in(char *buf, char *end, const char *arg, int len,
+		   const char *new_val)
+{
+	char *to = buf;
+
+	/* copy the arg name */
+	if (to + len >= end)
+		return -E2BIG;
+	memcpy(to, arg, len);
+	to += len;
+
+	if (new_val == BOOTFLOWCL_EMPTY) {
+		/* no value */
+	} else {
+		bool need_quote = strchr(new_val, ' ');
+		len = strlen(new_val);
+
+		/* need space for value, equals sign and maybe two quotes */
+		if (to + 1 + (need_quote ? 2 : 0) + len >= end)
+			return -E2BIG;
+		*to++ = '=';
+		if (need_quote)
+			*to++ = '"';
+		memcpy(to, new_val, len);
+		to += len;
+		if (need_quote)
+			*to++ = '"';
+	}
+
+	return to - buf;
+}
+
+int cmdline_set_arg(char *buf, int maxlen, const char *cmdline,
+		    const char *set_arg, const char *new_val, int *posp)
+{
+	bool found_arg = false;
+	const char *from;
+	char *to, *end;
+	int set_arg_len;
+	char empty = '\0';
+	int ret;
+
+	from = cmdline ?: ∅
+
+	/* check if the value has quotes inside */
+	if (new_val && new_val != BOOTFLOWCL_EMPTY && strchr(new_val, '"'))
+		return -EBADF;
+
+	set_arg_len = strlen(set_arg);
+	for (to = buf, end = buf + maxlen; *from;) {
+		const char *val, *arg_end, *val_end, *p;
+		bool in_quote;
+
+		if (to >= end)
+			return -E2BIG;
+		while (*from == ' ')
+			from++;
+		if (!*from)
+			break;
+
+		/* find the end of this arg */
+		val = NULL;
+		arg_end = NULL;
+		val_end = NULL;
+		in_quote = false;
+		for (p = from;; p++) {
+			if (in_quote) {
+				if (!*p)
+					return -EINVAL;
+				if (*p == '"')
+					in_quote = false;
+				continue;
+			}
+			if (*p == '=') {
+				arg_end = p;
+				val = p + 1;
+			} else if (*p == '"') {
+				in_quote = true;
+			} else if (!*p || *p == ' ') {
+				val_end = p;
+				if (!arg_end)
+					arg_end = p;
+				break;
+			}
+		}
+		/*
+		 * At this point val_end points to the end of the value, or the
+		 * last char after the arg name, if there is no label.
+		 * arg_end is the char after the arg name
+		 * val points to the value, or NULL if there is none
+		 * char after the value.
+		 *
+		 *        fred=1234
+		 *        ^   ^^   ^
+		 *      from  ||   |
+		 *           / \    \
+		 *    arg_end  val   val_end
+		 */
+		log_debug("from %s arg_end %ld val %ld val_end %ld\n", from,
+			  (long)(arg_end - from), (long)(val - from),
+			  (long)(val_end - from));
+
+		if (to != buf) {
+			if (to >= end)
+				return -E2BIG;
+			*to++ = ' ';
+		}
+
+		/* if this is the target arg, update it */
+		if (!strncmp(from, set_arg, arg_end - from)) {
+			if (!buf) {
+				bool has_quote = val_end[-1] == '"';
+
+				/*
+				 * exclude any start/end quotes from
+				 * calculations
+				 */
+				if (!val)
+					val = val_end;
+				*posp = val - cmdline + has_quote;
+				return val_end - val - 2 * has_quote;
+			}
+			found_arg = true;
+			if (!new_val) {
+				/* delete this arg */
+				from = val_end + (*val_end == ' ');
+				log_debug("delete from: %s\n", from);
+				if (to != buf)
+					to--; /* drop the space we added */
+				continue;
+			}
+
+			ret = copy_in(to, end, from, arg_end - from, new_val);
+			if (ret < 0)
+				return ret;
+			to += ret;
+
+		/* if not the target arg, copy it unchanged */
+		} else if (to) {
+			int len;
+
+			len = val_end - from;
+			if (to + len >= end)
+				return -E2BIG;
+			memcpy(to, from, len);
+			to += len;
+		}
+		from = val_end;
+	}
+
+	/* If we didn't find the arg, add it */
+	if (!found_arg) {
+		/* trying to delete something that is not there */
+		if (!new_val || !buf)
+			return -ENOENT;
+		if (to >= end)
+			return -E2BIG;
+
+		/* add a space to separate it from the previous arg */
+		if (to != buf && to[-1] != ' ')
+			*to++ = ' ';
+		ret = copy_in(to, end, set_arg, set_arg_len, new_val);
+		log_debug("ret=%d, to: %s buf: %s\n", ret, to, buf);
+		if (ret < 0)
+			return ret;
+		to += ret;
+	}
+
+	/* delete any trailing space */
+	if (to > buf && to[-1] == ' ')
+		to--;
+
+	if (to >= end)
+		return -E2BIG;
+	*to++ = '\0';
+
+	return to - buf;
+}

include/bootflow.h

@@ -444,4 +444,44 @@ int bootflow_menu_apply_theme(struct expo *exp, ofnode node);
 int bootflow_menu_run(struct bootstd_priv *std, bool text_mode,
 		      struct bootflow **bflowp);
 
+#define BOOTFLOWCL_EMPTY	((void *)1)
+
+/**
+ * cmdline_set_arg() - Update or read an argument in a cmdline string
+ *
+ * Handles updating a single arg in a cmdline string, returning it in a supplied
+ * buffer; also reading an arg from a cmdline string
+ *
+ * When updating, consecutive spaces are squashed as are spaces at the start and
+ * end.
+ *
+ * @buf: Working buffer to use (initial contents are ignored). Use NULL when
+ * reading
+ * @maxlen: Length of working buffer. Use 0 when reading
+ * @cmdline: Command line to update, in the form:
+ *
+ *	fred mary= jane=123 john="has spaces"
+ *
+ * @set_arg: Argument to set or read (may or may not exist)
+ * @new_val: Value for the new argument. May not include quotes (") but may
+ * include embedded spaces, in which case it will be quoted when added to the
+ * command line. Use NULL to delete the argument from @cmdline, BOOTFLOWCL_EMPTY
+ * to set it to an empty value (no '=' sign after arg), "" to add an '=' sign
+ * but with an empty value. Use NULL when reading.
+ * @posp: Ignored when setting an argument; when getting an argument, returns
+ * the start position of its value in @cmdline, after the first quote, if any
+ *
+ * Return:
+ * For updating:
+ *	length of new buffer (including \0 terminator) on success, -ENOENT if
+ *	@new_val is NULL and @set_arg does not exist in @from, -EINVAL if a
+ *	quoted arg-value in @from is not terminated with a quote, -EBADF if
+ *	@new_val has spaces but does not start and end with quotes (or it has
+ *	quotes in the middle of the string), -E2BIG if @maxlen is too small
+ * For reading:
+ *	length of arg value (excluding quotes), -ENOENT if not found
+ */
+int cmdline_set_arg(char *buf, int maxlen, const char *cmdline,
+		    const char *set_arg, const char *new_val, int *posp);
+
 #endif

test/boot/bootflow.c

@@ -683,3 +683,157 @@ static int bootflow_menu_theme(struct unit_test_state *uts)
 	return 0;
 }
 BOOTSTD_TEST(bootflow_menu_theme, UT_TESTF_DM | UT_TESTF_SCAN_FDT);
+
+/**
+ * check_arg() - Check both the normal case and the buffer-overflow case
+ *
+ * @uts: Unit-test state
+ * @expect_ret: Expected return value (i.e. buffer length)
+ * @expect_str: String expected to be returned
+ * @buf: Buffer to use
+ * @from: Original cmdline to update
+ * @arg: Argument to update (e.g. "console")
+ * @val: Value to set (e.g. "ttyS2") or NULL to delete the argument if present,
+ * "" to set it to an empty value (e.g. "console=") and BOOTFLOWCL_EMPTY to add
+ * it without any value ("initrd")
+ */
+static int check_arg(struct unit_test_state *uts, int expect_ret,
+		     const char *expect_str, char *buf, const char *from,
+		     const char *arg, const char *val)
+{
+	/* check for writing outside the reported bounds */
+	buf[expect_ret] = '[';
+	ut_asserteq(expect_ret,
+		    cmdline_set_arg(buf, expect_ret, from, arg, val, NULL));
+	ut_asserteq_str(expect_str, buf);
+	ut_asserteq('[', buf[expect_ret]);
+
+	/* do the test again but with one less byte in the buffer */
+	ut_asserteq(-E2BIG, cmdline_set_arg(buf, expect_ret - 1, from, arg,
+					    val, NULL));
+
+	return 0;
+}
+
+/* Test of bootflow_cmdline_set_arg() */
+static int test_bootflow_cmdline_set(struct unit_test_state *uts)
+{
+	char buf[50];
+	const int size = sizeof(buf);
+
+	/*
+	 * note that buffer-overflow tests are immediately each test case, just
+	 * top keep the code together
+	 */
+
+	/* add an arg that doesn't already exist, starting from empty */
+	ut_asserteq(-ENOENT, cmdline_set_arg(buf, size, NULL, "me", NULL,
+					     NULL));
+
+	ut_assertok(check_arg(uts, 3, "me", buf, NULL, "me", BOOTFLOWCL_EMPTY));
+	ut_assertok(check_arg(uts, 4, "me=", buf, NULL, "me", ""));
+	ut_assertok(check_arg(uts, 8, "me=fred", buf, NULL, "me", "fred"));
+
+	/* add an arg that doesn't already exist, starting from non-empty */
+	ut_assertok(check_arg(uts, 11, "arg=123 me", buf, "arg=123", "me",
+			      BOOTFLOWCL_EMPTY));
+	ut_assertok(check_arg(uts, 12, "arg=123 me=", buf, "arg=123", "me",
+			      ""));
+	ut_assertok(check_arg(uts, 16, "arg=123 me=fred", buf, "arg=123", "me",
+			      "fred"));
+
+	/* update an arg at the start */
+	ut_assertok(check_arg(uts, 1, "", buf, "arg=123", "arg", NULL));
+	ut_assertok(check_arg(uts, 4, "arg", buf, "arg=123", "arg",
+			      BOOTFLOWCL_EMPTY));
+	ut_assertok(check_arg(uts, 5, "arg=", buf, "arg=123", "arg", ""));
+	ut_assertok(check_arg(uts, 6, "arg=1", buf, "arg=123", "arg", "1"));
+	ut_assertok(check_arg(uts, 9, "arg=1234", buf, "arg=123", "arg",
+			      "1234"));
+
+	/* update an arg at the end */
+	ut_assertok(check_arg(uts, 5, "mary", buf, "mary arg=123", "arg",
+			      NULL));
+	ut_assertok(check_arg(uts, 9, "mary arg", buf, "mary arg=123", "arg",
+			      BOOTFLOWCL_EMPTY));
+	ut_assertok(check_arg(uts, 10, "mary arg=", buf, "mary arg=123", "arg",
+			      ""));
+	ut_assertok(check_arg(uts, 11, "mary arg=1", buf, "mary arg=123", "arg",
+			      "1"));
+	ut_assertok(check_arg(uts, 14, "mary arg=1234", buf, "mary arg=123",
+			      "arg", "1234"));
+
+	/* update an arg in the middle */
+	ut_assertok(check_arg(uts, 16, "mary=abc john=2", buf,
+			      "mary=abc arg=123 john=2", "arg", NULL));
+	ut_assertok(check_arg(uts, 20, "mary=abc arg john=2", buf,
+			      "mary=abc arg=123 john=2", "arg",
+			      BOOTFLOWCL_EMPTY));
+	ut_assertok(check_arg(uts, 21, "mary=abc arg= john=2", buf,
+			      "mary=abc arg=123 john=2", "arg", ""));
+	ut_assertok(check_arg(uts, 22, "mary=abc arg=1 john=2", buf,
+			      "mary=abc arg=123 john=2", "arg", "1"));
+	ut_assertok(check_arg(uts, 25, "mary=abc arg=1234 john=2", buf,
+			      "mary=abc arg=123 john=2", "arg", "1234"));
+
+	/* handle existing args with quotes */
+	ut_assertok(check_arg(uts, 16, "mary=\"abc\" john", buf,
+			      "mary=\"abc\" arg=123 john", "arg", NULL));
+
+	/* handle existing args with quoted spaces */
+	ut_assertok(check_arg(uts, 20, "mary=\"abc def\" john", buf,
+			      "mary=\"abc def\" arg=123 john", "arg", NULL));
+
+	ut_assertok(check_arg(uts, 34, "mary=\"abc def\" arg=123 john def=4",
+			      buf, "mary=\"abc def\" arg=123 john", "def",
+			      "4"));
+
+	/* quote at the start */
+	ut_asserteq(-EBADF, cmdline_set_arg(buf, size,
+					    "mary=\"abc def\" arg=\"123 456\"",
+					    "arg", "\"4 5 6", NULL));
+
+	/* quote at the end */
+	ut_asserteq(-EBADF, cmdline_set_arg(buf, size,
+					    "mary=\"abc def\" arg=\"123 456\"",
+					    "arg", "4 5 6\"", NULL));
+
+	/* quote in the middle */
+	ut_asserteq(-EBADF, cmdline_set_arg(buf, size,
+					    "mary=\"abc def\" arg=\"123 456\"",
+					    "arg", "\"4 \"5 6\"", NULL));
+
+	/* handle updating a quoted arg */
+	ut_assertok(check_arg(uts, 27, "mary=\"abc def\" arg=\"4 5 6\"", buf,
+			      "mary=\"abc def\" arg=\"123 456\"", "arg",
+			      "4 5 6"));
+
+	/* changing a quoted arg to a non-quoted arg */
+	ut_assertok(check_arg(uts, 23, "mary=\"abc def\" arg=789", buf,
+			      "mary=\"abc def\" arg=\"123 456\"", "arg",
+			      "789"));
+
+	/* changing a non-quoted arg to a quoted arg */
+	ut_assertok(check_arg(uts, 29, "mary=\"abc def\" arg=\"456 789\"", buf,
+			      "mary=\"abc def\" arg=123", "arg", "456 789"));
+
+	/* handling of spaces */
+	ut_assertok(check_arg(uts, 8, "arg=123", buf, " ", "arg", "123"));
+	ut_assertok(check_arg(uts, 8, "arg=123", buf, "   ", "arg", "123"));
+	ut_assertok(check_arg(uts, 13, "john arg=123", buf, " john  ", "arg",
+			      "123"));
+	ut_assertok(check_arg(uts, 13, "john arg=123", buf, " john  arg=123  ",
+			      "arg", "123"));
+	ut_assertok(check_arg(uts, 18, "john arg=123 mary", buf,
+			      " john  arg=123 mary ", "arg", "123"));
+
+	/* unchanged arg */
+	ut_assertok(check_arg(uts, 3, "me", buf, "me", "me", BOOTFLOWCL_EMPTY));
+
+	/* arg which starts with the same name */
+	ut_assertok(check_arg(uts, 28, "mary=abc johnathon=2 john=3", buf,
+			      "mary=abc johnathon=2 john=1", "john", "3"));
+
+	return 0;
+}
+BOOTSTD_TEST(test_bootflow_cmdline_set, 0);