1 files changed, 23 insertions, 11 deletions

diff --git a/drivers/dfu/dfu_mtd.c b/drivers/dfu/dfu_mtd.c
index cce9ce0845..c7075f12ec 100644
--- a/drivers/dfu/dfu_mtd.c
+++ b/drivers/dfu/dfu_mtd.c
@@ -12,6 +12,7 @@
 #include <mtd.h>
 #include <jffs2/load_kernel.h>
 #include <linux/err.h>
+#include <linux/ctype.h>
 
 static bool mtd_is_aligned_with_block_size(struct mtd_info *mtd, u64 size)
 {
@@ -270,9 +271,9 @@ static unsigned int dfu_polltimeout_mtd(struct dfu_entity *dfu)
 	return DFU_DEFAULT_POLL_TIMEOUT;
 }
 
-int dfu_fill_entity_mtd(struct dfu_entity *dfu, char *devstr, char *s)
+int dfu_fill_entity_mtd(struct dfu_entity *dfu, char *devstr, char **argv, int argc)
 {
-	char *st;
+	char *s;
 	struct mtd_info *mtd;
 	int ret, part;
 
@@ -284,22 +285,33 @@ int dfu_fill_entity_mtd(struct dfu_entity *dfu, char *devstr, char *s)
 	dfu->dev_type = DFU_DEV_MTD;
 	dfu->data.mtd.info = mtd;
 	dfu->max_buf_size = mtd->erasesize;
+	if (argc < 1)
+		return -EINVAL;
 
-	st = strsep(&s, " ");
-	if (!strcmp(st, "raw")) {
+	if (!strcmp(argv[0], "raw")) {
+		if (argc != 3)
+			return -EINVAL;
 		dfu->layout = DFU_RAW_ADDR;
-		dfu->data.mtd.start = hextoul(s, &s);
-		s++;
-		dfu->data.mtd.size = hextoul(s, &s);
-	} else if ((!strcmp(st, "part")) || (!strcmp(st, "partubi"))) {
+		dfu->data.mtd.start = hextoul(argv[1], &s);
+		if (*s)
+			return -EINVAL;
+		dfu->data.mtd.size = hextoul(argv[2], &s);
+		if (*s)
+			return -EINVAL;
+	} else if ((!strcmp(argv[0], "part")) || (!strcmp(argv[0], "partubi"))) {
 		char mtd_id[32];
 		struct mtd_device *mtd_dev;
 		u8 part_num;
 		struct part_info *pi;
 
+		if (argc != 2)
+			return -EINVAL;
+
 		dfu->layout = DFU_RAW_ADDR;
 
-		part = dectoul(s, &s);
+		part = dectoul(argv[1], &s);
+		if (*s)
+			return -EINVAL;
 
 		sprintf(mtd_id, "%s,%d", devstr, part - 1);
 		printf("using id '%s'\n", mtd_id);
@@ -314,10 +326,10 @@ int dfu_fill_entity_mtd(struct dfu_entity *dfu, char *devstr, char *s)
 
 		dfu->data.mtd.start = pi->offset;
 		dfu->data.mtd.size = pi->size;
-		if (!strcmp(st, "partubi"))
+		if (!strcmp(argv[0], "partubi"))
 			dfu->data.mtd.ubi = 1;
 	} else {
-		printf("%s: Memory layout (%s) not supported!\n", __func__, st);
+		printf("%s: Memory layout (%s) not supported!\n", __func__, argv[0]);
 		return -1;
 	}