1 files changed, 38 insertions, 16 deletions

diff --git a/cmd/setexpr.c b/cmd/setexpr.c
index e111b8ba98a..c45fa85c887 100644
--- a/cmd/setexpr.c
+++ b/cmd/setexpr.c
@@ -35,9 +35,37 @@ struct expr_arg {
 	};
 };
 
+/**
+ * arg_set_str() - copy string to expression argument
+ *
+ * The string is truncated to 64 KiB plus NUL terminator.
+ *
+ * @p:		pointer to string
+ * @argp:	pointer to expression argument
+ * Return:	0 on success, -ENOMEM if out of memory
+ */
+static int arg_set_str(void *p, struct expr_arg *argp)
+{
+	int len;
+	char *str;
+
+	/* Maximum string length of 64 KiB plus NUL terminator */
+	len = strnlen((char *)p, SZ_64K) + 1;
+	str = malloc(len);
+	if (!str) {
+		printf("Out of memory\n");
+		return -ENOMEM;
+	}
+	memcpy(str, p, len);
+	str[len - 1] = '\0';
+	argp->sval = str;
+	return 0;
+}
+
 static int get_arg(char *s, int w, struct expr_arg *argp)
 {
 	struct expr_arg arg;
+	int ret;
 
 	/*
 	 * If the parameter starts with a '*' then assume it is a pointer to
@@ -47,8 +75,6 @@ static int get_arg(char *s, int w, struct expr_arg *argp)
 		ulong *p;
 		ulong addr;
 		ulong val;
-		int len;
-		char *str;
 
 		addr = hextoul(&s[1], NULL);
 		switch (w) {
@@ -66,18 +92,10 @@ static int get_arg(char *s, int w, struct expr_arg *argp)
 			break;
 		case CMD_DATA_SIZE_STR:
 			p = map_sysmem(addr, SZ_64K);
-
-			/* Maximum string length of 64KB plus terminator */
-			len = strnlen((char *)p, SZ_64K) + 1;
-			str = malloc(len);
-			if (!str) {
-				printf("Out of memory\n");
-				return -ENOMEM;
-			}
-			memcpy(str, p, len);
-			str[len - 1] = '\0';
+			ret = arg_set_str(p, &arg);
 			unmap_sysmem(p);
-			arg.sval = str;
+			if (ret)
+				return ret;
 			break;
 		case 4:
 			p = map_sysmem(addr, sizeof(u32));
@@ -93,9 +111,13 @@ static int get_arg(char *s, int w, struct expr_arg *argp)
 			break;
 		}
 	} else {
-		if (w == CMD_DATA_SIZE_STR)
-			return -EINVAL;
-		arg.ival = hextoul(s, NULL);
+		if (w == CMD_DATA_SIZE_STR) {
+			ret = arg_set_str(s, &arg);
+			if (ret)
+				return ret;
+		} else {
+			arg.ival = hextoul(s, NULL);
+		}
 	}
 	*argp = arg;