From c6b4e506d1a5240fbb21ec99a4141e09169eb542 Mon Sep 17 00:00:00 2001
From: Michal Moskal <michal@moskal.me>
Date: Wed, 26 Jul 2017 18:14:26 +0100
Subject: [PATCH] Send restart request to uf2d upon usb disconnect

---
 brick/kernel/f_mass_storage.c | 26 +++++++++++++++++++++++++-
 brick/kernel/storage_common.c |  1 +
 brick/uf2daemon/fat.c         | 11 +++++++++++
 3 files changed, 37 insertions(+), 1 deletion(-)

diff --git a/brick/kernel/f_mass_storage.c b/brick/kernel/f_mass_storage.c
index 6e48802d..2c6d4b05 100644
--- a/brick/kernel/f_mass_storage.c
+++ b/brick/kernel/f_mass_storage.c
@@ -2410,7 +2410,28 @@ static void fsg_disable(struct usb_function *f)
 	fsg->common->prev_fsg = fsg->common->fsg;
 	fsg->common->fsg = fsg;
 	fsg->common->new_config = 0;
-	raise_exception(fsg->common, FSG_STATE_CONFIG_CHANGE);
+
+	DBG(common, "fsg_disable filp=%p\n", fsg->common->luns[0].filp);
+	raise_exception(fsg->common, FSG_STATE_CONFIG_CHANGE_DISABLE);
+}
+
+static void shutdown_server(void) {
+	DBG(fsg_common, "shutdown_server filp=%p\n", fsg_common->luns[0].filp);
+	if (fsg_common->luns[0].filp) {
+		uint32_t buf[512 / 4];
+		loff_t file_offset_tmp = 512 * 50000; // make sure we're outside of FS area
+		int i;
+
+		// this should shut down the nbd server, so that the caches are flushed
+		memset(buf, 0, sizeof(buf));
+		buf[0] = 0x20da6d81;
+		buf[1] = 0x747e09d4;
+
+		fsg_common->luns[0].filp->f_flags |= O_SYNC;
+		for (i = 0; i < 2; ++i)
+			vfs_write(fsg_common->luns[0].filp,
+							(char*)buf, 512, &file_offset_tmp);
+	}
 }
 
 
@@ -2532,6 +2553,9 @@ static void handle_exception(struct fsg_common *common)
 		/*		SS_RESET_OCCURRED;  */
 		break;
 
+	case FSG_STATE_CONFIG_CHANGE_DISABLE:
+		shutdown_server();
+		// fall-through
 	case FSG_STATE_CONFIG_CHANGE:
 		rc = do_set_config(common, new_config);
 		if (common->ep0_req_tag != exception_req_tag)
diff --git a/brick/kernel/storage_common.c b/brick/kernel/storage_common.c
index 868d8ee8..6b0dde67 100644
--- a/brick/kernel/storage_common.c
+++ b/brick/kernel/storage_common.c
@@ -334,6 +334,7 @@ enum fsg_state {
 	FSG_STATE_ABORT_BULK_OUT,
 	FSG_STATE_RESET,
 	FSG_STATE_INTERFACE_CHANGE,
+	FSG_STATE_CONFIG_CHANGE_DISABLE,
 	FSG_STATE_CONFIG_CHANGE,
 	FSG_STATE_DISCONNECT,
 	FSG_STATE_EXIT,
diff --git a/brick/uf2daemon/fat.c b/brick/uf2daemon/fat.c
index ed43045f..4d98835d 100644
--- a/brick/uf2daemon/fat.c
+++ b/brick/uf2daemon/fat.c
@@ -703,12 +703,23 @@ void restartProgram() {
         exit(0); // causes parent to eject MSD etc
 }
 
+int numWrites = 0;
 static WriteState wrState;
 void write_block(uint32_t block_no, uint8_t *data) {
     WriteState *state = &wrState;
 
     UF2_Block *bl = (void *)data;
 
+    if (bl->magicStart0 == 0x20da6d81 && bl->magicStart1 == 0x747e09d4) {
+        DBG("restart req, #wr=%d", numWrites);
+        if (numWrites) {
+            exit(0);
+        }
+        return;
+    }
+
+    numWrites++;
+
     if (!is_uf2_block(bl)) {
         return;
     }