1 files changed, 23 insertions, 21 deletions

diff --git a/common/cmd_usb_mass_storage.c b/common/cmd_usb_mass_storage.c
index ccf7195946..f583cafa35 100644
--- a/common/cmd_usb_mass_storage.c
+++ b/common/cmd_usb_mass_storage.c
@@ -8,51 +8,53 @@
 #include <common.h>
 #include <command.h>
 #include <g_dnl.h>
+#include <usb.h>
 #include <usb_mass_storage.h>
 
 int do_usb_mass_storage(cmd_tbl_t *cmdtp, int flag,
 			       int argc, char * const argv[])
 {
-	char *ep;
-	unsigned int dev_num = 0, offset = 0, part_size = 0;
-	int rc;
+	if (argc < 3)
+		return CMD_RET_USAGE;
 
-	struct ums_board_info *ums_info;
-	static char *s = "ums";
-
-	if (argc < 2) {
-		printf("usage: ums <dev> - e.g. ums 0\n");
-		return 0;
-	}
-
-	dev_num = (int)simple_strtoul(argv[1], &ep, 16);
+	const char *usb_controller = argv[1];
+	const char *mmc_devstring  = argv[2];
 
+	unsigned int dev_num = (unsigned int)(simple_strtoul(mmc_devstring,
+				NULL, 0));
 	if (dev_num) {
-		puts("\nSet eMMC device to 0! - e.g. ums 0\n");
+		error("Set eMMC device to 0! - e.g. ums 0");
 		goto fail;
 	}
 
-	board_usb_init();
-	ums_info = board_ums_init(dev_num, offset, part_size);
+	unsigned int controller_index = (unsigned int)(simple_strtoul(
+					usb_controller,	NULL, 0));
+	if (board_usb_init(controller_index, USB_INIT_DEVICE)) {
+		error("Couldn't init USB controller.");
+		goto fail;
+	}
 
+	struct ums_board_info *ums_info = board_ums_init(dev_num, 0, 0);
 	if (!ums_info) {
-		printf("MMC: %d -> NOT available\n", dev_num);
+		error("MMC: %d -> NOT available", dev_num);
 		goto fail;
 	}
-	rc = fsg_init(ums_info);
+
+	int rc = fsg_init(ums_info);
 	if (rc) {
-		printf("cmd ums: fsg_init failed\n");
+		error("fsg_init failed");
 		goto fail;
 	}
 
-	g_dnl_register(s);
+	g_dnl_register("ums");
 
 	while (1) {
 		/* Handle control-c and timeouts */
 		if (ctrlc()) {
-			printf("The remote end did not respond in time.\n");
+			error("The remote end did not respond in time.");
 			goto exit;
 		}
+
 		usb_gadget_handle_interrupts();
 		/* Check if USB cable has been detached */
 		if (fsg_main_thread(NULL) == EIO)
@@ -68,5 +70,5 @@ fail:
 
 U_BOOT_CMD(ums, CONFIG_SYS_MAXARGS, 1, do_usb_mass_storage,
 	"Use the UMS [User Mass Storage]",
-	"ums - User Mass Storage Gadget"
+	"<USB_controller> <mmc_dev>"
 );