1 files changed, 50 insertions, 25 deletions

diff --git a/common/iomux.c b/common/iomux.c
index 15bf533885..b9088aa3b5 100644
--- a/common/iomux.c
+++ b/common/iomux.c
@@ -15,18 +15,26 @@ void iomux_printdevs(const int console)
 	int i;
 	struct stdio_dev *dev;
 
-	for (i = 0; i < cd_count[console]; i++) {
-		dev = console_devices[console][i];
+	for_each_console_dev(i, console, dev)
 		printf("%s ", dev->name);
-	}
 	printf("\n");
 }
 
+int iomux_match_device(struct stdio_dev **set, const int n, struct stdio_dev *sdev)
+{
+	int i;
+
+	for (i = 0; i < n; i++)
+		if (sdev == set[i])
+			return i;
+	return -ENOENT;
+}
+
 /* This tries to preserve the old list if an error occurs. */
 int iomux_doenv(const int console, const char *arg)
 {
 	char *console_args, *temp, **start;
-	int i, j, k, io_flag, cs_idx, repeat;
+	int i, j, io_flag, cs_idx, repeat;
 	struct stdio_dev **cons_set, **old_set;
 	struct stdio_dev *dev;
 
@@ -75,15 +83,8 @@ int iomux_doenv(const int console, const char *arg)
 		return 1;
 	}
 
-	switch (console) {
-	case stdin:
-		io_flag = DEV_FLAGS_INPUT;
-		break;
-	case stdout:
-	case stderr:
-		io_flag = DEV_FLAGS_OUTPUT;
-		break;
-	default:
+	io_flag = stdio_file_to_flags(console);
+	if (io_flag < 0) {
 		free(start);
 		free(console_args);
 		free(cons_set);
@@ -103,14 +104,8 @@ int iomux_doenv(const int console, const char *arg)
 		/*
 		 * Prevent multiple entries for a device.
 		 */
-		 repeat = 0;
-		 for (k = 0; k < cs_idx; k++) {
-			if (dev == cons_set[k]) {
-				repeat++;
-				break;
-			}
-		 }
-		 if (repeat)
+		 repeat = iomux_match_device(cons_set, cs_idx, dev);
+		 if (repeat >= 0)
 			continue;
 		/*
 		 * Try assigning the specified device.
@@ -136,10 +131,7 @@ int iomux_doenv(const int console, const char *arg)
 
 	/* Stop dropped consoles */
 	for (i = 0; i < repeat; i++) {
-		for (j = 0; j < cs_idx; j++) {
-			if (old_set[i] == cons_set[j])
-				break;
-		}
+		j = iomux_match_device(cons_set, cs_idx, old_set[i]);
 		if (j == cs_idx)
 			console_stop(console, old_set[i]);
 	}
@@ -147,4 +139,37 @@ int iomux_doenv(const int console, const char *arg)
 	free(old_set);
 	return 0;
 }
+
+int iomux_replace_device(const int console, const char *old, const char *new)
+{
+	struct stdio_dev *dev;
+	char *arg = NULL;	/* Initial empty list */
+	int size = 1;		/* For NUL terminator */
+	int i, ret;
+
+	for_each_console_dev(i, console, dev) {
+		const char *name = strcmp(dev->name, old) ? dev->name : new;
+		char *tmp;
+
+		/* Append name with a ',' (comma) separator */
+		tmp = realloc(arg, size + strlen(name) + 1);
+		if (!tmp) {
+			free(arg);
+			return -ENOMEM;
+		}
+
+		strcat(tmp, ",");
+		strcat(tmp, name);
+
+		arg = tmp;
+		size = strlen(tmp) + 1;
+	}
+
+	ret = iomux_doenv(console, arg);
+	if (ret)
+		ret = -EINVAL;
+
+	free(arg);
+	return ret;
+}
 #endif /* CONSOLE_MUX */