From fa49e3e4c7183f6e276ecbdb4fa63792942c9ef7 Mon Sep 17 00:00:00 2001
From: Rajavardhan Gundi <rajavardhan.gundi@intel.com>
Date: Wed, 20 Feb 2019 12:33:30 +0530
Subject: [PATCH] usb: dfu: Signal completion of DFU

Generally when DFU is in progress, the system is not expected to
be doing anything else in addition. Hence, a completion signal
would help the system to know that DFU is over and it can proceed
towards next tasks.

Signed-off-by: Rajavardhan Gundi <rajavardhan.gundi@intel.com>
---
 include/usb/class/usb_dfu.h |  2 ++
 subsys/usb/class/Kconfig    |  8 +++++
 subsys/usb/class/usb_dfu.c  | 58 +++++++++++++++++++++++++++++++++++++
 3 files changed, 68 insertions(+)

diff --git a/include/usb/class/usb_dfu.h b/include/usb/class/usb_dfu.h
index 46d6223db36..3eb99a04ea2 100644
--- a/include/usb/class/usb_dfu.h
+++ b/include/usb/class/usb_dfu.h
@@ -119,4 +119,6 @@ enum dfu_state {
 	dfuERROR,
 };
 
+void wait_for_usb_dfu(void);
+
 #endif /* ZEPHYR_INCLUDE_USB_CLASS_USB_DFU_H_ */
diff --git a/subsys/usb/class/Kconfig b/subsys/usb/class/Kconfig
index 8dd0e9c0921..85c50561bf3 100644
--- a/subsys/usb/class/Kconfig
+++ b/subsys/usb/class/Kconfig
@@ -126,10 +126,18 @@ source "subsys/usb/class/hid/Kconfig"
 config USB_DFU_CLASS
 	bool "USB DFU Class Driver"
 	select MPU_ALLOW_FLASH_WRITE
+	select POLL
 	depends on IMG_MANAGER
 	help
 	  USB DFU class driver
 
+config USB_DFU_WAIT_DELAY_MS
+	int
+	depends on USB_DFU_CLASS
+	default 12000
+	help
+	  A thread can wait for a prescribed time (in ms) for DFU to begin
+
 config USB_DFU_MAX_XFER_SIZE
 	int
 	depends on USB_DFU_CLASS
diff --git a/subsys/usb/class/usb_dfu.c b/subsys/usb/class/usb_dfu.c
index f9a050c4c75..ab9e5e31bff 100644
--- a/subsys/usb/class/usb_dfu.c
+++ b/subsys/usb/class/usb_dfu.c
@@ -68,6 +68,15 @@ LOG_MODULE_REGISTER(usb_dfu);
 #define FIRMWARE_IMAGE_0_LABEL "image-1"
 #define FIRMWARE_IMAGE_1_LABEL "image-0"
 
+/* MCUBoot waits for CONFIG_USB_DFU_WAIT_DELAY_MS time in total to let DFU to
+ * be commenced. It intermittently checks every INTERMITTENT_CHECK_DELAY
+ * milliseconds to see if DFU has started.
+ */
+#define INTERMITTENT_CHECK_DELAY	50
+
+static struct k_poll_event dfu_event;
+static struct k_poll_signal dfu_signal;
+
 static struct k_work dfu_work;
 
 struct dfu_worker_data_t {
@@ -433,6 +442,8 @@ static int dfu_class_handle_req(struct usb_setup_packet *pSetup,
 		case dfuIDLE:
 			LOG_DBG("DFU_DNLOAD start");
 			dfu_reset_counters();
+			k_poll_signal_reset(&dfu_signal);
+
 			if (dfu_data.flash_area_id !=
 			    DT_FLASH_AREA_IMAGE_1_ID) {
 				dfu_data.status = errWRITE;
@@ -453,6 +464,7 @@ static int dfu_class_handle_req(struct usb_setup_packet *pSetup,
 			dfu_data_worker.worker_len  = pSetup->wLength;
 			if (dfu_data_worker.worker_len == 0) {
 				dfu_data.state = dfuMANIFEST_SYNC;
+				k_poll_signal_raise(&dfu_signal, 0);
 			}
 
 			memcpy(dfu_data_worker.buf, *data, pSetup->wLength);
@@ -741,6 +753,7 @@ static int usb_dfu_init(struct device *dev)
 	ARG_UNUSED(dev);
 
 	k_work_init(&dfu_work, dfu_work_handler);
+	k_poll_signal_init(&dfu_signal);
 
 #ifndef CONFIG_USB_COMPOSITE_DEVICE
 	dfu_config.usb_device_description = usb_get_device_descriptor();
@@ -772,4 +785,49 @@ static int usb_dfu_init(struct device *dev)
 	return 0;
 }
 
+/**
+ * @brief Function to check if DFU is started.
+ *
+ * @return  true if DNBUSY/DNLOAD_IDLE, false otherwise.
+ */
+static bool is_dfu_started(void)
+{
+	if ((dfu_data.state == dfuDNBUSY) ||
+	    (dfu_data.state == dfuDNLOAD_IDLE)) {
+		return true;
+	}
+
+	return false;
+}
+
+/**
+ * @brief Function to check and wait while the USB DFU is in progress.
+ *
+ * @return  N/A
+ */
+void wait_for_usb_dfu(void)
+{
+	/* Wait for a prescribed duration of time. If DFU hasn't started within
+	 * that time, stop waiting and proceed further.
+	 */
+	for (int time = 0;
+		time < (CONFIG_USB_DFU_WAIT_DELAY_MS/INTERMITTENT_CHECK_DELAY);
+		time++) {
+		if (is_dfu_started()) {
+			k_poll_event_init(&dfu_event, K_POLL_TYPE_SIGNAL,
+				K_POLL_MODE_NOTIFY_ONLY, &dfu_signal);
+
+			/* Wait till DFU is complete */
+			if (k_poll(&dfu_event, 1, K_FOREVER) != 0) {
+				LOG_DBG("USB DFU Error");
+			}
+
+			LOG_INF("USB DFU Completed");
+			break;
+		}
+
+		k_sleep(INTERMITTENT_CHECK_DELAY);
+	}
+}
+
 SYS_INIT(usb_dfu_init, APPLICATION, CONFIG_KERNEL_INIT_PRIORITY_DEVICE);