From 66fdb153784c8dcb5bc17918c6b5e46e13c54216 Mon Sep 17 00:00:00 2001
From: Daniele Lacamera <dan@danielinux.net>
Date: Mon, 12 Jun 2023 01:39:15 +0200
Subject: [PATCH] Experiments with second disk

---
 src/cryptoengine.c  |  33 ++++++++
 src/cryptoengine.h  |   7 ++
 src/flash.h         |   2 +
 src/fsm.c           |   9 ++-
 src/msc_disk.c      | 193 ++++++++++++++++++++++++++++++++++++--------
 src/password_safe.c |   5 +-
 src/ui.c            |   5 +-
 7 files changed, 214 insertions(+), 40 deletions(-)

diff --git a/src/cryptoengine.c b/src/cryptoengine.c
index ac7223f..6c5d2fa 100644
--- a/src/cryptoengine.c
+++ b/src/cryptoengine.c
@@ -41,6 +41,39 @@ static ecc_key ecc;
 static struct vault_service svc_cache;
 static int svc_cache_id = -1;
 
+int flash_decrypt_read_sector(uint32_t sector, uint8_t *buf)
+{
+    if (cryptoengine_check_vault() < 0)
+        return -1;
+    if (fsm_get() < VAULT_MAIN_MENU)
+        return -1;
+    if ((sector * SPI_FLASH_SECTOR_SIZE) >= DRIVE_FLASH_SIZE)
+        return -1;
+    flash_read(DRIVE_FLASH_OFFSET + sector * SPI_FLASH_SECTOR_SIZE, buf,
+            SPI_FLASH_SECTOR_SIZE);
+    wc_Chacha_SetIV(&cha, hdr_cache.host_seed, sector);
+    wc_Chacha_Process(&cha, buf, buf, SPI_FLASH_SECTOR_SIZE);
+    return 0;
+}
+
+int flash_encrypt_write_sector(uint32_t sector, uint8_t *buf)
+{
+    if (cryptoengine_check_vault() < 0)
+        return -1;
+    if (fsm_get() < VAULT_MAIN_MENU)
+        return -1;
+    if ((sector * SPI_FLASH_SECTOR_SIZE) >= DRIVE_FLASH_SIZE)
+        return -1;
+    wc_Chacha_SetIV(&cha, hdr_cache.host_seed, sector);
+    wc_Chacha_Process(&cha, buf, buf, SPI_FLASH_SECTOR_SIZE);
+
+    flash_sector_erase(DRIVE_FLASH_OFFSET + sector * SPI_FLASH_SECTOR_SIZE);
+
+    flash_write(DRIVE_FLASH_OFFSET + sector * SPI_FLASH_SECTOR_SIZE, buf,
+            SPI_FLASH_SECTOR_SIZE);
+    return 0;
+}
+
 int flash_decrypt_read_svc(struct vault_service *out, uint32_t addr)
 {
     if (cryptoengine_check_vault() < 0)
diff --git a/src/cryptoengine.h b/src/cryptoengine.h
index eed5ad2..95c3d8d 100644
--- a/src/cryptoengine.h
+++ b/src/cryptoengine.h
@@ -32,6 +32,9 @@
 #define VAULT_FLASH_OFFSET 0x00000000
 #define VAULT_FLASH_SIZE   (512 * 1024)
 
+#define DRIVE_FLASH_OFFSET (0x00080000)
+#define DRIVE_FLASH_SIZE   (512 * 1024)
+
 #define VAULT_DIGEST_SIZE WC_SHA512_DIGEST_SIZE
 #define SHA_PAYLOAD_SIZE (sizeof(struct vault_header) - (PK_SIGNATURE_SIZE + VAULT_DIGEST_SIZE))
 
@@ -116,4 +119,8 @@ int cryptoengine_service_count(uint16_t *n_srv, uint16_t *first_avail);
 int cryptoengine_fill_vault_status(struct vault_status *vst);
 int flash_decrypt_read_svc(struct vault_service *out, uint32_t addr);
 
+/* In msc_disk.c */
+void disk_crypto_activate(int status);
+void disk1_format(void);
+
 #endif
diff --git a/src/flash.h b/src/flash.h
index e73efae..925843e 100644
--- a/src/flash.h
+++ b/src/flash.h
@@ -28,5 +28,7 @@ int flash_sector_erase(uint32_t address);
 int flash_read(uint32_t address, void *data, int len);
 int flash_write(uint32_t address, const void *data, int len);
 int flash_write_svc(struct vault_service *svc, int idx);
+int flash_decrypt_read_sector(uint32_t sector, uint8_t *buf);
+int flash_encrypt_write_sector(uint32_t sector, uint8_t *buf);
 uint8_t flash_read_status(void);
 #endif
diff --git a/src/fsm.c b/src/fsm.c
index 148d93b..64332ef 100644
--- a/src/fsm.c
+++ b/src/fsm.c
@@ -24,6 +24,7 @@
 #include "bsp/board.h"
 #include "fsm.h"
 #include "ui.h"
+#include "cryptoengine.h"
 
 static enum vault_state vault_state = VAULT_OFF;
 static volatile int ui_event_pending = 0;
@@ -45,9 +46,15 @@ void fsm_set(enum vault_state st)
 
     if (st == vault_state)
         return;
+
     /* Change state */
     vault_state = st;
-    ui_event_cb(old_st, fsm_get());
+    ui_event_cb(old_st, vault_state);
+    if ((vault_state == VAULT_SETTINGS_MENU) && 
+            (old_st == VAULT_VERIFY_PASSPHRASE)){
+        /* Just logged in! */
+        disk_crypto_activate(1);
+    }
 }
 
 enum vault_state fsm_get(void)
diff --git a/src/msc_disk.c b/src/msc_disk.c
index 2888fe4..57824e4 100644
--- a/src/msc_disk.c
+++ b/src/msc_disk.c
@@ -25,6 +25,8 @@
 
 #include "bsp/board.h"
 #include "tusb.h"
+#include "flash.h"
+#include "fsm.h"
 
 #if CFG_TUD_MSC
 
@@ -32,6 +34,8 @@
 #define DISK_BLOCK_NUM 1024
 #define DISK_BLOCK_SIZE 512
 
+#define BLOCKS_PER_SECTOR (SPI_FLASH_SECTOR_SIZE / DISK_BLOCK_SIZE)
+
 extern const unsigned char disk0_img[];
 
 struct disk_lba {
@@ -39,11 +43,88 @@ struct disk_lba {
 };
 
 
+static const uint8_t last_cluster_init_info[4] = { 0xF8, 0xFF, 0xFF, 0xFF };
+
+static const uint8_t disk1_mbr[DISK_BLOCK_SIZE] = {
+    0xeb, 0x3c, 0x90, 0x6d, 0x6b, 0x66, 0x73, 0x2e, 0x66, 0x61, 0x74, 0x00,
+    0x02, 0x04, 0x01, 0x00, 0x02, 0x00, 0x02, 0x00, 0x04, 0xf8, 0x01, 0x00,
+    0x10, 0x00, 0x02, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+    0x80, 0x00, 0x29, 0x9d, 0x0b, 0x6e, 0xd0, 0x4e, 0x4f, 0x20, 0x4e, 0x41,
+    0x4d, 0x45, 0x20, 0x20, 0x20, 0x20, 0x46, 0x41, 0x54, 0x31, 0x32, 0x20,
+    0x20, 0x20, 0x0e, 0x1f, 0xbe, 0x5b, 0x7c, 0xac, 0x22, 0xc0, 0x74, 0x0b,
+    0x56, 0xb4, 0x0e, 0xbb, 0x07, 0x00, 0xcd, 0x10, 0x5e, 0xeb, 0xf0, 0x32,
+    0xe4, 0xcd, 0x16, 0xcd, 0x19, 0xeb, 0xfe, 0x54, 0x68, 0x69, 0x73, 0x20,
+    0x69, 0x73, 0x20, 0x6e, 0x6f, 0x74, 0x20, 0x61, 0x20, 0x62, 0x6f, 0x6f,
+    0x74, 0x61, 0x62, 0x6c, 0x65, 0x20, 0x64, 0x69, 0x73, 0x6b, 0x2e, 0x20,
+    0x20, 0x50, 0x6c, 0x65, 0x61, 0x73, 0x65, 0x20, 0x69, 0x6e, 0x73, 0x65,
+    0x72, 0x74, 0x20, 0x61, 0x20, 0x62, 0x6f, 0x6f, 0x74, 0x61, 0x62, 0x6c,
+    0x65, 0x20, 0x66, 0x6c, 0x6f, 0x70, 0x70, 0x79, 0x20, 0x61, 0x6e, 0x64,
+    0x0d, 0x0a, 0x70, 0x72, 0x65, 0x73, 0x73, 0x20, 0x61, 0x6e, 0x79, 0x20,
+    0x6b, 0x65, 0x79, 0x20, 0x74, 0x6f, 0x20, 0x74, 0x72, 0x79, 0x20, 0x61,
+    0x67, 0x61, 0x69, 0x6e, 0x20, 0x2e, 0x2e, 0x2e, 0x20, 0x0d, 0x0a, 0x00,
+    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x55, 0xaa
+};
+
+
 // whether host does safe-eject
 static bool ejected = false;
 
-static struct disk_lba lba_cache;
-static int lba_cached = -1;
+static struct disk_lba sector_cache[BLOCKS_PER_SECTOR];
+static int sector_cached = -1;
+
+static int sector_cache_load(uint32_t sector)
+{
+    int ret;
+    if (sector * SPI_FLASH_SECTOR_SIZE > DRIVE_FLASH_SIZE) {
+        sector_cached = -1;
+        return -1;
+    }
+    //ret = flash_decrypt_read_sector(sector, sector_cache[0].data);
+    ret = flash_read(sector * SPI_FLASH_SECTOR_SIZE + DRIVE_FLASH_OFFSET,
+            sector_cache[0].data, SPI_FLASH_SECTOR_SIZE);
+    if (ret < 0) {
+        sector_cached = -1;
+        return ret;
+    }
+    sector_cached = sector;
+}
+
+static void sector_cache_commit(void)
+{
+    int ret;
+    uint32_t sector = sector_cached;
+    if (sector * SPI_FLASH_SECTOR_SIZE > DRIVE_FLASH_SIZE)
+        return;
+    //flash_encrypt_write_sector(sector, sector_cache[0].data);
+    flash_write(sector * SPI_FLASH_SECTOR_SIZE + DRIVE_FLASH_OFFSET, 
+            sector_cache[0].data, SPI_FLASH_SECTOR_SIZE);
+}
 
 struct drive {
     const uint8_t lun;
@@ -55,7 +136,7 @@ struct drive {
 
 #define N_DRIVES 2
 
-static uint8_t drv_active[N_DRIVES] = {1, 0};
+static uint8_t drv_active[N_DRIVES] = {1, 1};
 
 
 const char global_vid[8] = "DLX";
@@ -64,14 +145,14 @@ const char vault_pid[16] = "MEP Vault";
 const char global_rev[4] = "0.1";
 
 static const struct drive Drives[N_DRIVES] = { 
-    { 0, "DLX", "MEP Tools", "0.1", 1024},
+    { 0, "DLX", "MEP Tools", "0.1", 256},
     { 1, "DLX", "MEP Vault", "0.1", 1024}
 };
 
 
 uint8_t tud_msc_get_maxlun_cb(void)
 {
-  return 1;
+  return 2;
 }
 
 // Invoked when received SCSI_CMD_INQUIRY
@@ -109,7 +190,7 @@ bool tud_msc_test_unit_ready_cb(uint8_t lun)
 void tud_msc_capacity_cb(uint8_t lun, uint32_t* block_count, uint16_t* block_size)
 {
     if (lun >= N_DRIVES) {
-        *block_count = DISK_BLOCK_NUM;
+        *block_count = 0;
         *block_size  = DISK_BLOCK_SIZE;
     } else {
         *block_count = Drives[lun].n_lba;
@@ -147,37 +228,37 @@ int32_t tud_msc_read10_cb(uint8_t lun, uint32_t lba, uint32_t offset, void* buff
 {
 
     const struct drive *drv;
-    uint8_t const *addr;
+    uint32_t sector, sec_off;
     if (lun >= N_DRIVES)
         return -1;
     drv = &Drives[lun];
 
     if ( lba >= drv->n_lba )
         return -1;
-
-
-
-
+    
     if (lun == 0) {
         struct disk_lba *msc_ro = (struct disk_lba *)disk0_img;
         memcpy(buffer, msc_ro[lba].data + offset, bufsize);
         asm volatile("DMB");
         return bufsize;
     } else {
-#if 0
-        if ( lba != lba_cached) {
-            if (lba == 0) {
-                memcpy(lba_cache.data, mbr.data, sizeof(struct disk_lba)); 
-                lba_cached = lba;
-                return 0;
-            }
-            memcpy(lba_cache.data, msc_ro[lba].data, sizeof(struct disk_lba)); 
-            lba_cached = lba;
+        if (fsm_get() < VAULT_MAIN_MENU) {
+            memset(buffer, 0, bufsize);
+            return bufsize;
+        }
+        if (lba == 0) {
+            memcpy(buffer, disk1_mbr + offset, bufsize);
+            return bufsize;
+        }
+        sector = lba / BLOCKS_PER_SECTOR;
+        sec_off = lba - (sector * BLOCKS_PER_SECTOR);
+         
+        if ( sector != sector_cached) {
+            sector_cache_load(sector);
             return 0;
         }
-        memcpy(buffer, lba_cache.data + offset, bufsize);
-#endif
-        memset(buffer, 0xA5, bufsize);
+        memcpy(buffer, sector_cache[sec_off].data + offset, bufsize);
+        asm volatile("DMB");
         return bufsize;
     }
 }
@@ -185,8 +266,8 @@ int32_t tud_msc_read10_cb(uint8_t lun, uint32_t lba, uint32_t offset, void* buff
 bool tud_msc_is_writable_cb (uint8_t lun)
 {
     if (lun != 1) {
-      /* Read only. */
-      return false;
+        /* Read only. */
+        return false;
     } else {
         return true;
     }
@@ -196,13 +277,26 @@ bool tud_msc_is_writable_cb (uint8_t lun)
 // Process data in buffer to disk's storage and return number of written bytes
 int32_t tud_msc_write10_cb(uint8_t lun, uint32_t lba, uint32_t offset, uint8_t* buffer, uint32_t bufsize)
 {
-  (void) lun;
-  // out of ramdisk
-  if ( lba >= DISK_BLOCK_NUM )
-      return -1;
-  /* Read only. Silently fail. */
-  (void) lba; (void) offset; (void) buffer;
-  return bufsize;
+    uint32_t sector, sec_off;
+    if (lun == 0)
+        return bufsize;
+    /* 
+    if (lba == 0)
+        return bufsize;
+    */
+    if ( lba >= Drives[lun].n_lba)
+        return -1;
+    if (fsm_get() < VAULT_MAIN_MENU)
+        return -1;
+    sector = lba / BLOCKS_PER_SECTOR;
+    sec_off = lba - (sector * BLOCKS_PER_SECTOR);
+    if ( sector != sector_cached) {
+        sector_cache_load(sector);
+    }
+    memcpy(sector_cache[sec_off].data + offset, buffer, bufsize);
+    sector_cache_commit();
+    asm volatile("DMB");
+    return bufsize;
 }
 
 // Callback invoked when received an SCSI command not in built-in list below
@@ -251,4 +345,39 @@ int32_t tud_msc_scsi_cb (uint8_t lun, uint8_t const scsi_cmd[16], void* buffer,
     return resplen;
 }
 
+void disk1_format(void)
+{
+    uint8_t block[DISK_BLOCK_SIZE];
+    int i, j;
+    if (cryptoengine_check_vault() < 0)
+        return;
+
+    memset(block, 0, DISK_BLOCK_SIZE);
+    for (i = 1; i < 5; i++) {
+        if (sector_cached != i) {
+            sector_cache_load(i);
+        }
+        for (j = 0; j < BLOCKS_PER_SECTOR; j++)
+            memcpy(sector_cache[j].data, block, DISK_BLOCK_SIZE);
+        sector_cache_commit();
+    }
+
+    /* Fill "last cluster" info */
+    sector_cache_load(0);
+    for (j = 0; j < BLOCKS_PER_SECTOR; j++)
+            memcpy(sector_cache[j].data, block, DISK_BLOCK_SIZE);
+    memcpy(sector_cache[0].data, disk1_mbr, DISK_BLOCK_SIZE);
+    memcpy(sector_cache[1].data, last_cluster_init_info, 4); 
+    memcpy(sector_cache[2].data, last_cluster_init_info, 4); 
+    sector_cache_commit();
+}
+
+void disk_crypto_activate(int status)
+{
+    drv_active[1] = status;
+    if (status) {
+
+    }
+}
+
 #endif
diff --git a/src/password_safe.c b/src/password_safe.c
index d5bc189..8195747 100644
--- a/src/password_safe.c
+++ b/src/password_safe.c
@@ -84,8 +84,6 @@ static uint Led[] = {
 #define SPI_MISO_PIN     4
 #define SPI_FLASH_CS_PIN 5
 
-
-
 uint16_t flash_info = 0xFFFF;
 
 #define AIRCR *(volatile uint32_t *)(0xE000ED0C)
@@ -166,6 +164,7 @@ void system_boot(void)
     spi_init(SPI, SPI_BAUDRATE);
     flash_info = flash_init(SPI);
     
+    
     printf("Initializing I2C Display\r\n");
     gpio_init(I2C_SDA_PIN);
     gpio_init(I2C_SCL_PIN);
@@ -311,8 +310,6 @@ int main(void) {
     int i;
     system_boot();
     printf("Loop started.\n");
-
-
     while (main_plug) {
         uint32_t ms_tick = board_millis();
         if ((ms_tick % 10) == 0) {
diff --git a/src/ui.c b/src/ui.c
index 884f483..be1d1cc 100644
--- a/src/ui.c
+++ b/src/ui.c
@@ -81,16 +81,15 @@ void ui_event_cb(enum vault_state old_st, enum vault_state st)
                 fsm_set(VAULT_VERIFY_FAILED);
                 return;
             }
-            gpio_put(GREEN_LED, 0);
             fsm_set(VAULT_MAIN_MENU);
-            return;
-            break;
+            return; /* State changed, avoid non-tail recursion */
         case VAULT_MAIN_MENU:
             display_text(0,"                 ");
             display_text(1,"                 ");
             display_text(2,"                 ");
             display_text(3,"                 ");
             display_text_inverse(0, 2, "[Main Menu]");
+            gpio_put(GREEN_LED, 0);
             break;
         case VAULT_SETTINGS_MENU:
             display_text(0,"                 ");