diff --git a/doc/system_description/firmware.md b/doc/system_description/firmware.md
index 7fea0aa..d8bcaba 100644
--- a/doc/system_description/firmware.md
+++ b/doc/system_description/firmware.md
@@ -18,6 +18,54 @@ If your available `objcopy` and `size` commands is anything other than
 the default `llvm-objcopy` and `llvm-size` define `OBJCOPY` and `SIZE`
 to whatever they're called on your system.
 
+## Firmware state machine
+
+States:
+
+- `initial` - At start.
+- `init_loading` - Reset app digest, size, `USS` and load address.
+- `loading` - Expect more app data or a reset by `LoadApp()`.
+- `run` - Computes CDI and starts the device app.
+
+Commands in state `initial`:
+
+| *command*          | *next state*   |
+|--------------------|----------------|
+| NameVersion()      | unchanged      |
+| GetUDI()           | unchanged      |
+| LoadApp(size, uss) | `init_loading` |
+|                    |                |
+
+Commands in `init_loading`:
+
+| *command*          | *next state*   |
+|--------------------|----------------|
+| NameVersion()      | unchanged      |
+| GetUDI()           | unchanged      |
+| LoadApp(size, uss) | `init_loading` |
+| LoadAppData(data)  | `loading`      |
+|                    |                |
+
+Commands in `loading`:
+
+| *command*          | *next state*                     |
+|--------------------|----------------------------------|
+| NameVersion()      | unchanged                        |
+| GetUDI()           | unchanged                        |
+| LoadApp(size, uss) | `init_loading`                   |
+| LoadAppData(data)  | `loading` or `run` on last chunk |
+|                    |                                  |
+
+Behaviour:
+
+- NameVersion: identifies stick.
+- GetUDI: returns the Unique Device ID with vendor id, product id,
+  revision, serial number.
+- LoadApp(size, uss): Start loading an app with this size and user
+  supplied secret.
+- LoadAppData(data): Load chunk of data of app. When last data chunk
+  is received we compute and return the digest.
+
 ## Using QEMU
 
 Checkout the `tk1` branch of [our version of the
diff --git a/hw/application_fpga/fw/tk1/main.c b/hw/application_fpga/fw/tk1/main.c
index 7e9d463..b2fec00 100644
--- a/hw/application_fpga/fw/tk1/main.c
+++ b/hw/application_fpga/fw/tk1/main.c
@@ -60,9 +60,10 @@ static void print_digest(uint8_t *md)
 }
 
 // CDI = blake2s(uds, blake2s(app), uss)
-static void compute_cdi(uint8_t digest[32], uint8_t uss[32])
+static void compute_cdi(uint8_t digest[32], uint8_t use_uss, uint8_t uss[32])
 {
 	uint32_t local_cdi[8];
+	int len;
 
 	// To protect UDS we use a special firmware-only RAM for both
 	// the in parameter to blake2s and the blake2s context.
@@ -70,11 +71,16 @@ static void compute_cdi(uint8_t digest[32], uint8_t uss[32])
 	// Only word aligned access to UDS
 	wordcpy((void *)fw_ram, (void *)uds, 8);
 	memcpy((void *)fw_ram + 32, digest, 32);
-	memcpy((void *)fw_ram + 64, uss, 32);
+	if (use_uss != 0) {
+		memcpy((void *)fw_ram + 64, uss, 32);
+		len = 96;
+	} else {
+		len = 64;
+	}
 
-	blake2s_ctx *secure_ctx = (blake2s_ctx *)(fw_ram + 96);
+	blake2s_ctx *secure_ctx = (blake2s_ctx *)(fw_ram + len);
 
-	blake2s((void *)local_cdi, 32, NULL, 0, (const void *)fw_ram, 96,
+	blake2s((void *)local_cdi, 32, NULL, 0, (const void *)fw_ram, len,
 		secure_ctx);
 
 	// Write over the firmware-only RAM
@@ -84,6 +90,13 @@ static void compute_cdi(uint8_t digest[32], uint8_t uss[32])
 	wordcpy((void *)cdi, (void *)local_cdi, 8);
 }
 
+enum state {
+	FW_STATE_INITIAL,
+	FW_STATE_INIT_LOADING,
+	FW_STATE_LOADING,
+	FW_STATE_RUN
+};
+
 int main()
 {
 	uint32_t local_name0 = *name0;
@@ -94,179 +107,34 @@ int main()
 	uint8_t rsp[CMDLEN_MAXBYTES];
 	uint8_t *loadaddr = (uint8_t *)TK1_APP_ADDR;
 	int left = 0; // Bytes left to receive
+	uint8_t use_uss = FALSE;
 	uint8_t uss[32] = {0};
 	uint8_t digest[32] = {0};
+	enum state state = FW_STATE_INITIAL;
 
 	print_hw_version(local_name0, local_name1, local_ver);
 
 	for (;;) {
-		// blocking; fw flashing white while waiting for cmd
-		uint8_t in = readbyte_ledflash(LED_WHITE, 800000);
-
-		if (parseframe(in, &hdr) == -1) {
-			puts("Couldn't parse header\n");
-			continue;
-		}
-
-		memset(cmd, 0, CMDLEN_MAXBYTES);
-		// Read firmware command: Blocks!
-		read(cmd, hdr.len);
-
-		// Is it for us?
-		if (hdr.endpoint != DST_FW) {
-			puts("Message not meant for us\n");
-			continue;
-		}
-
-		// Reset response buffer
-		memset(rsp, 0, CMDLEN_MAXBYTES);
-
-		// Min length is 1 byte so this should always be here
-		switch (cmd[0]) {
-		case FW_CMD_NAME_VERSION:
-			puts("cmd: name-version\n");
-
-			if (hdr.len != 1) {
-				// Bad length - give them an empty response
-				fwreply(hdr, FW_RSP_NAME_VERSION, rsp);
-				break;
-			}
-
-			memcpy(rsp, (uint8_t *)&local_name0, 4);
-			memcpy(rsp + 4, (uint8_t *)&local_name1, 4);
-			memcpy(rsp + 8, (uint8_t *)&local_ver, 4);
-
-			fwreply(hdr, FW_RSP_NAME_VERSION, rsp);
+		switch (state) {
+		case FW_STATE_INITIAL:
 			break;
 
-		case FW_CMD_GET_UDI:
-			puts("FW_CMD_GET_UDI\n");
-			if (hdr.len != 1) {
-				// Bad cmd length
-				rsp[0] = STATUS_BAD;
-				fwreply(hdr, FW_RSP_GET_UDI, rsp);
-				break;
-			}
-			rsp[0] = STATUS_OK;
-			uint32_t udi_words[2];
-			wordcpy(udi_words, (void *)udi, 2);
-			memcpy(rsp + 1, udi_words, 2 * 4);
-			fwreply(hdr, FW_RSP_GET_UDI, rsp);
-			break;
-
-		case FW_CMD_LOAD_USS:
-			puts("cmd: load-uss\n");
-
-			if (hdr.len != 128) {
-				// Bad cmd length
-				rsp[0] = STATUS_BAD;
-				fwreply(hdr, FW_RSP_LOAD_USS, rsp);
-				break;
-			}
-
-			memcpy(uss, cmd + 1, 32);
-
-			rsp[0] = STATUS_OK;
-			fwreply(hdr, FW_RSP_LOAD_USS, rsp);
-			break;
-
-		case FW_CMD_LOAD_APP_SIZE:
-			puts("cmd: load-app-size\n");
-
-			if (hdr.len != 32) {
-				// Bad length
-				rsp[0] = STATUS_BAD;
-				fwreply(hdr, FW_RSP_LOAD_APP_SIZE, rsp);
-				break;
-			}
-
-			// cmd[1..4] contains the size.
-			uint32_t local_app_size = cmd[1] + (cmd[2] << 8) +
-						  (cmd[3] << 16) +
-						  (cmd[4] << 24);
-
-			puts("app size: ");
-			putinthex(local_app_size);
-			lf();
-
-			if (local_app_size > TK1_APP_MAX_SIZE) {
-				rsp[0] = STATUS_BAD;
-				fwreply(hdr, FW_RSP_LOAD_APP_SIZE, rsp);
-				break;
-			}
-
-			*app_size = local_app_size;
+		case FW_STATE_INIT_LOADING:
 			*app_addr = 0;
-			// Clear digest as GET_APP_DIGEST returns it even if it
-			// has not been calculated
-			memset(digest, 0, 32);
+			left = *app_size;
 
 			// Reset where to start loading the program
 			loadaddr = (uint8_t *)TK1_APP_ADDR;
-			left = *app_size;
-
-			rsp[0] = STATUS_OK;
-			fwreply(hdr, FW_RSP_LOAD_APP_SIZE, rsp);
 			break;
 
-		case FW_CMD_LOAD_APP_DATA:
-			puts("cmd: load-app-data\n");
-
-			if (hdr.len != 128 || *app_size == 0 ||
-			    *app_addr != 0) {
-				// Bad length, or app_size not yet set, or
-				// app_addr already set (fully loaded!)
-				rsp[0] = STATUS_BAD;
-				fwreply(hdr, FW_RSP_LOAD_APP_DATA, rsp);
-				break;
-			}
-
-			int nbytes;
-			if (left > 127) {
-				nbytes = 127;
-			} else {
-				nbytes = left;
-			}
-			memcpy(loadaddr, cmd + 1, nbytes);
-			loadaddr += nbytes;
-			left -= nbytes;
-
-			if (left == 0) {
-				puts("Fully loaded ");
-				putinthex(*app_size);
-				lf();
-
-				*app_addr = TK1_APP_ADDR;
-				// Get the Blake2S digest of the app - store it
-				// for later queries
-				blake2s_ctx ctx;
-
-				blake2s(digest, 32, NULL, 0,
-					(const void *)*app_addr, *app_size,
-					&ctx);
-				print_digest(digest);
-
-				// CDI = hash(uds, hash(app), uss)
-				compute_cdi(digest, uss);
-			}
-
-			rsp[0] = STATUS_OK;
-			fwreply(hdr, FW_RSP_LOAD_APP_DATA, rsp);
+		case FW_STATE_LOADING:
 			break;
 
-		case FW_CMD_RUN_APP:
-			puts("cmd: run-app\n");
+		case FW_STATE_RUN:
+			*app_addr = TK1_APP_ADDR;
 
-			if (hdr.len != 1 || *app_size == 0 || *app_addr == 0) {
-				// Bad cmd length, or app_size or app_addr are
-				// not yet set
-				rsp[0] = STATUS_BAD;
-				fwreply(hdr, FW_RSP_RUN_APP, rsp);
-				break;
-			}
-
-			rsp[0] = STATUS_OK;
-			fwreply(hdr, FW_RSP_RUN_APP, rsp);
+			// CDI = hash(uds, hash(app), uss)
+			compute_cdi(digest, use_uss, uss);
 
 			// Flip over to application mode
 			*switch_app = 1;
@@ -293,11 +161,160 @@ int main()
 			// clang-format on
 			break; // This is never reached!
 
-		case FW_CMD_GET_APP_DIGEST:
-			puts("cmd: get-app-digest\n");
+		default:
+			// Unknown state!?
+			puts("Unknown firmware state 0x");
+			puthex(state);
+			lf();
 
-			memcpy(rsp, &digest, 32);
-			fwreply(hdr, FW_RSP_GET_APP_DIGEST, rsp);
+			asm volatile("forever:;"
+				     "j forever;");
+			break; // Not reached
+		}
+
+		// blocking; fw flashing white while waiting for cmd
+		uint8_t in = readbyte_ledflash(LED_WHITE, 800000);
+
+		if (parseframe(in, &hdr) == -1) {
+			puts("Couldn't parse header\n");
+			continue;
+		}
+
+		memset(cmd, 0, CMDLEN_MAXBYTES);
+		// Now we know the size of the cmd frame, read it all
+		read(cmd, hdr.len);
+
+		// Is it for us?
+		if (hdr.endpoint != DST_FW) {
+			puts("Message not meant for us\n");
+			continue;
+		}
+
+		// Reset response buffer
+		memset(rsp, 0, CMDLEN_MAXBYTES);
+
+		// Min length is 1 byte so this should always be here
+		switch (cmd[0]) {
+		case FW_CMD_NAME_VERSION:
+			puts("cmd: name-version\n");
+			if (hdr.len != 1) {
+				// Bad length - give them an empty response
+				fwreply(hdr, FW_RSP_NAME_VERSION, rsp);
+				break;
+			}
+
+			memcpy(rsp, (uint8_t *)&local_name0, 4);
+			memcpy(rsp + 4, (uint8_t *)&local_name1, 4);
+			memcpy(rsp + 8, (uint8_t *)&local_ver, 4);
+			fwreply(hdr, FW_RSP_NAME_VERSION, rsp);
+			// state unchanged
+			break;
+
+		case FW_CMD_GET_UDI:
+			puts("cmd: get-udi\n");
+			if (hdr.len != 1) {
+				// Bad cmd length
+				rsp[0] = STATUS_BAD;
+				fwreply(hdr, FW_RSP_GET_UDI, rsp);
+				break;
+			}
+
+			rsp[0] = STATUS_OK;
+			uint32_t udi_words[2];
+			wordcpy(udi_words, (void *)udi, 2);
+			memcpy(rsp + 1, udi_words, 2 * 4);
+			fwreply(hdr, FW_RSP_GET_UDI, rsp);
+			// state unchanged
+			break;
+
+		case FW_CMD_LOAD_APP:
+			puts("cmd: load-app(size, uss)\n");
+			if (hdr.len != 128) {
+				// Bad length
+				rsp[0] = STATUS_BAD;
+				fwreply(hdr, FW_RSP_LOAD_APP, rsp);
+				break;
+			}
+
+			// cmd[1..4] contains the size.
+			uint32_t local_app_size = cmd[1] + (cmd[2] << 8) +
+						  (cmd[3] << 16) +
+						  (cmd[4] << 24);
+
+			puts("app size: ");
+			putinthex(local_app_size);
+			lf();
+
+			if (local_app_size > TK1_APP_MAX_SIZE) {
+				rsp[0] = STATUS_BAD;
+				fwreply(hdr, FW_RSP_LOAD_APP, rsp);
+				break;
+			}
+
+			*app_size = local_app_size;
+
+			// Do we have a USS at all?
+			if (cmd[5] != 0) {
+				// Yes
+				use_uss = TRUE;
+				memcpy(uss, cmd + 6, 32);
+			} else {
+				use_uss = FALSE;
+			}
+
+			rsp[0] = STATUS_OK;
+			fwreply(hdr, FW_RSP_LOAD_APP, rsp);
+
+			state = FW_STATE_INIT_LOADING;
+			break;
+
+		case FW_CMD_LOAD_APP_DATA:
+			puts("cmd: load-app-data\n");
+			if (hdr.len != 128 || (state != FW_STATE_INIT_LOADING &&
+					       state != FW_STATE_LOADING)) {
+				// Bad cmd length or state
+				rsp[0] = STATUS_BAD;
+				fwreply(hdr, FW_RSP_LOAD_APP_DATA, rsp);
+				break;
+			}
+
+			int nbytes;
+			if (left > 127) {
+				nbytes = 127;
+			} else {
+				nbytes = left;
+			}
+			memcpy(loadaddr, cmd + 1, nbytes);
+			loadaddr += nbytes;
+			left -= nbytes;
+
+			if (left == 0) {
+				puts("Fully loaded ");
+				putinthex(*app_size);
+				lf();
+
+				// Compute Blake2S digest of the app, storing
+				// it for FW_STATE_RUN
+				blake2s_ctx ctx;
+
+				blake2s(digest, 32, NULL, 0,
+					(const void *)TK1_APP_ADDR, *app_size,
+					&ctx);
+				print_digest(digest);
+
+				// And return the digest in final response
+				rsp[0] = STATUS_OK;
+				memcpy(&rsp[1], &digest, 32);
+				fwreply(hdr, FW_RSP_LOAD_APP_DATA_READY, rsp);
+
+				state = FW_STATE_RUN;
+				break;
+			}
+
+			rsp[0] = STATUS_OK;
+			fwreply(hdr, FW_RSP_LOAD_APP_DATA, rsp);
+
+			state = FW_STATE_LOADING;
 			break;
 
 		default:
diff --git a/hw/application_fpga/fw/tk1/proto.c b/hw/application_fpga/fw/tk1/proto.c
index 60e6d4a..3ac617c 100644
--- a/hw/application_fpga/fw/tk1/proto.c
+++ b/hw/application_fpga/fw/tk1/proto.c
@@ -71,12 +71,7 @@ void fwreply(struct frame_header hdr, enum fwcmd rspcode, uint8_t *buf)
 		nbytes = 32;
 		break;
 
-	case FW_RSP_LOAD_USS:
-		len = LEN_4;
-		nbytes = 4;
-		break;
-
-	case FW_RSP_LOAD_APP_SIZE:
+	case FW_RSP_LOAD_APP:
 		len = LEN_4;
 		nbytes = 4;
 		break;
@@ -86,12 +81,7 @@ void fwreply(struct frame_header hdr, enum fwcmd rspcode, uint8_t *buf)
 		nbytes = 4;
 		break;
 
-	case FW_RSP_RUN_APP:
-		len = LEN_4;
-		nbytes = 4;
-		break;
-
-	case FW_RSP_GET_APP_DIGEST:
+	case FW_RSP_LOAD_APP_DATA_READY:
 		len = LEN_128;
 		nbytes = 128;
 		break;
diff --git a/hw/application_fpga/fw/tk1/proto.h b/hw/application_fpga/fw/tk1/proto.h
index 6a29043..d51848c 100644
--- a/hw/application_fpga/fw/tk1/proto.h
+++ b/hw/application_fpga/fw/tk1/proto.h
@@ -26,21 +26,15 @@ enum cmdlen {
 
 // clang-format off
 enum fwcmd {
-	FW_CMD_NAME_VERSION	= 0x01,
-	FW_RSP_NAME_VERSION	= 0x02,
-	FW_CMD_LOAD_APP_SIZE	= 0x03,
-	FW_RSP_LOAD_APP_SIZE	= 0x04,
-	FW_CMD_LOAD_APP_DATA	= 0x05,
-	FW_RSP_LOAD_APP_DATA	= 0x06,
-	FW_CMD_RUN_APP		= 0x07,
-	FW_RSP_RUN_APP		= 0x08,
-	FW_CMD_GET_APP_DIGEST	= 0x09,
-	FW_CMD_LOAD_USS		= 0x0a,
-	FW_RSP_LOAD_USS		= 0x0b,
-	FW_CMD_GET_UDI		= 0x0c,
-	FW_RSP_GET_UDI		= 0x0d,
-	/* ... */
-	FW_RSP_GET_APP_DIGEST	= 0x10, // encoded as 0x10 for backwards compatibility
+	FW_CMD_NAME_VERSION		= 0x01,
+	FW_RSP_NAME_VERSION		= 0x02,
+	FW_CMD_LOAD_APP			= 0x03,
+	FW_RSP_LOAD_APP			= 0x04,
+	FW_CMD_LOAD_APP_DATA		= 0x05,
+	FW_RSP_LOAD_APP_DATA		= 0x06,
+	FW_RSP_LOAD_APP_DATA_READY	= 0x07,
+	FW_CMD_GET_UDI			= 0x08,
+	FW_RSP_GET_UDI			= 0x09,
 };
 // clang-format on
 
diff --git a/hw/application_fpga/fw/tk1/types.h b/hw/application_fpga/fw/tk1/types.h
index 49c24a0..9a1fcff 100644
--- a/hw/application_fpga/fw/tk1/types.h
+++ b/hw/application_fpga/fw/tk1/types.h
@@ -16,4 +16,7 @@ typedef unsigned long size_t;
 
 #define NULL ((char *)0)
 
+#define FALSE 0
+#define TRUE !FALSE
+
 #endif