From 5c227a4a47d2a9fa01e64a5c6cdbcde4a3f7abda Mon Sep 17 00:00:00 2001
From: Pavol Rusnak <stick@gk2.sk>
Date: Sat, 1 Apr 2017 22:30:10 +0200
Subject: [PATCH] loader: compute pubkey rather than use precomputed one,
 prepare for vendor keys

---
 micropython/bootloader/main.c |  2 +-
 micropython/loader/main.c     |  7 ++--
 micropython/trezorhal/image.c | 66 ++++++++++++++++++++---------------
 micropython/trezorhal/image.h |  2 +-
 4 files changed, 42 insertions(+), 35 deletions(-)

diff --git a/micropython/bootloader/main.c b/micropython/bootloader/main.c
index b611e378e9..ddaf2e3e81 100644
--- a/micropython/bootloader/main.c
+++ b/micropython/bootloader/main.c
@@ -116,7 +116,7 @@ bool copy_sdcard(void)
 void check_and_jump(void)
 {
     BOOTLOADER_PRINTLN("checking loader");
-    if (image_check_signature((const uint8_t *)LOADER_START)) {
+    if (image_check_signature((const uint8_t *)LOADER_START, NULL)) {
         BOOTLOADER_PRINTLN("valid loader image");
         // TODO: remove debug wait
         BOOTLOADER_PRINTLN("waiting 1 second");
diff --git a/micropython/loader/main.c b/micropython/loader/main.c
index 5c7e3a2d92..b1bc35b40a 100644
--- a/micropython/loader/main.c
+++ b/micropython/loader/main.c
@@ -47,13 +47,10 @@ void check_and_jump(void)
         return;
     }
 
-    // TODO: use keys from vendor header in image_check_signature
-    if (image_check_signature((const uint8_t *)(FIRMWARE_START + vhdr.hdrlen))) {
+    if (image_check_signature((const uint8_t *)(FIRMWARE_START + vhdr.hdrlen), &vhdr)) {
         LOADER_PRINTLN("valid firmware image");
-        // TODO: remove debug wait
         display_vendor(vhdr.vimg, (const char *)vhdr.vstr, vhdr.vstr_len);
-        HAL_Delay(1000);
-        // end
+        HAL_Delay(1000); // TODO: remove?
         LOADER_PRINTLN("JUMP!");
         jump_to(FIRMWARE_START + vhdr.hdrlen + HEADER_SIZE);
     } else {
diff --git a/micropython/trezorhal/image.c b/micropython/trezorhal/image.c
index 4b6a3c5aad..06991d1633 100644
--- a/micropython/trezorhal/image.c
+++ b/micropython/trezorhal/image.c
@@ -38,37 +38,47 @@ bool image_parse_header(const uint8_t *data, image_header *header)
     return true;
 }
 
-#define KEYMASK(A, B, C) ((1 << (A - 1)) | (1 << (B - 1)) | (1 << (C - 1)))
+static const uint8_t * const SATOSHILABS_PUBKEYS[] = {
+    (const uint8_t *)"AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA",
+    (const uint8_t *)"AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA",
+    (const uint8_t *)"AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA",
+    (const uint8_t *)"AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA",
+    (const uint8_t *)"AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA",
+};
 
-static const uint8_t *get_pubkey(uint8_t index)
+static const uint8_t *compute_pubkey(const vendor_header *vhdr, uint8_t sigmask)
 {
-    switch (index) {
-        case KEYMASK(1, 2, 3):
-            return (const uint8_t *)"AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA";
-        case KEYMASK(1, 2, 4):
-            return (const uint8_t *)"AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA";
-        case KEYMASK(1, 2, 5):
-            return (const uint8_t *)"AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA";
-        case KEYMASK(1, 3, 4):
-            return (const uint8_t *)"AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA";
-        case KEYMASK(1, 3, 5):
-            return (const uint8_t *)"AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA";
-        case KEYMASK(1, 4, 5):
-            return (const uint8_t *)"AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA";
-        case KEYMASK(2, 3, 4):
-            return (const uint8_t *)"AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA";
-        case KEYMASK(2, 3, 5):
-            return (const uint8_t *)"AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA";
-        case KEYMASK(2, 4, 5):
-            return (const uint8_t *)"AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA";
-        case KEYMASK(3, 4, 5):
-            return (const uint8_t *)"AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA";
-        default:
-            return NULL;
+    uint8_t vsig_m;
+    uint8_t vsig_n;
+    const uint8_t * const *vpub;
+
+    if (vhdr) {
+        vsig_m = vhdr->vsig_n;
+        vsig_n = vhdr->vsig_n;
+        vpub = vhdr->vpub;
+    } else {
+        vsig_m = 3;
+        vsig_n = 5;
+        vpub = SATOSHILABS_PUBKEYS;
     }
+
+    if (!vsig_m || !vsig_n) return NULL;
+    if (vsig_m > vsig_n) return NULL;
+
+    // discard bits higher than vsig_n
+    sigmask &= ((1 << vsig_n) - 1);
+
+    // remove if number of set bits in sigmask is not equal to vsig_m
+    if (__builtin_popcount(sigmask) != vsig_m) return NULL;
+
+    // TODO: add keys from vpub according to sigmask
+    (void)vpub;
+    (void)sigmask;
+
+    return NULL;
 }
 
-bool image_check_signature(const uint8_t *data)
+bool image_check_signature(const uint8_t *data, const vendor_header *vhdr)
 {
     image_header hdr;
     if (!image_parse_header(data, &hdr)) {
@@ -85,7 +95,7 @@ bool image_check_signature(const uint8_t *data)
     blake2s_Update(&ctx, data + HEADER_SIZE, hdr.codelen);
     blake2s_Final(&ctx, hash, BLAKE2S_DIGEST_LENGTH);
 
-    const uint8_t *pub = get_pubkey(hdr.sigmask);
+    const uint8_t *pub = compute_pubkey(vhdr, hdr.sigmask);
 
     // TODO: remove debug skip of unsigned
     if (!pub) return true;
@@ -154,7 +164,7 @@ bool vendor_check_signature(const uint8_t *data)
     }
     blake2s_Final(&ctx, hash, BLAKE2S_DIGEST_LENGTH);
 
-    const uint8_t *pub = get_pubkey(hdr.sigmask);
+    const uint8_t *pub = compute_pubkey(NULL, hdr.sigmask);
 
     // TODO: remove debug skip of unsigned
     if (!pub) return true;
diff --git a/micropython/trezorhal/image.h b/micropython/trezorhal/image.h
index 8c4bd3b02f..f0e0c6412f 100644
--- a/micropython/trezorhal/image.h
+++ b/micropython/trezorhal/image.h
@@ -34,7 +34,7 @@ typedef struct {
 
 bool image_parse_header(const uint8_t *data, image_header *header);
 
-bool image_check_signature(const uint8_t *data);
+bool image_check_signature(const uint8_t *data, const vendor_header *vhdr);
 
 bool vendor_parse_header(const uint8_t *data, vendor_header *header);