bootloader: more protobuf cleanup

micropython/bootloader/messages.c

@@ -12,7 +12,8 @@ void send_msg_Success(int iface)
     // response: Success message (id 2), payload len 0
     PB_CTX ctx;
     pb_start(&ctx, 2);
-    usb_hid_write_blocking(iface, pb_build(&ctx), 64, 1);
+    pb_end(&ctx);
+    usb_hid_write_blocking(iface, ctx.buf, ctx.pos, 1);
 }
 
 void send_msg_Failure(int iface)
@@ -21,8 +22,9 @@ void send_msg_Failure(int iface)
     //           - code = 99 (Failure_FirmwareError)
     PB_CTX ctx;
     pb_start(&ctx, 3);
-    pb_add_varint(&ctx, "\x08", 99);
-    usb_hid_write_blocking(iface, pb_build(&ctx), 64, 1);
+    pb_add_varint(&ctx, 1, 99);
+    pb_end(&ctx);
+    usb_hid_write_blocking(iface, ctx.buf, ctx.pos, 1);
 }
 
 void send_msg_Features(int iface, bool firmware_present)
@@ -36,13 +38,14 @@ void send_msg_Features(int iface, bool firmware_present)
     //           - firmware_present = True/False
     PB_CTX ctx;
     pb_start(&ctx, 17);
-    pb_add_string(&ctx, "\x0a", "trezor.io");
-    pb_add_varint(&ctx, "\x10", VERSION_MAJOR);
-    pb_add_varint(&ctx, "\x18", VERSION_MINOR);
-    pb_add_varint(&ctx, "\x20", VERSION_PATCH);
-    pb_add_bool(&ctx, "\x28", true);
-    pb_add_bool(&ctx, "\x90\x01", firmware_present);
-    usb_hid_write_blocking(iface, pb_build(&ctx), 64, 1);
+    pb_add_string(&ctx, 1, "trezor.io");
+    pb_add_varint(&ctx, 2, VERSION_MAJOR);
+    pb_add_varint(&ctx, 3, VERSION_MINOR);
+    pb_add_varint(&ctx, 4, VERSION_PATCH);
+    pb_add_bool(&ctx, 5, true);
+    pb_add_bool(&ctx, 18, firmware_present);
+    pb_end(&ctx);
+    usb_hid_write_blocking(iface, ctx.buf, ctx.pos, 1);
 }
 
 void send_msg_FirmwareRequest(int iface, uint32_t offset, uint32_t length)
@@ -52,7 +55,8 @@ void send_msg_FirmwareRequest(int iface, uint32_t offset, uint32_t length)
     //           - length = length
     PB_CTX ctx;
     pb_start(&ctx, 8);
-    pb_add_varint(&ctx, "\x08", offset);
-    pb_add_varint(&ctx, "\x10", length);
-    usb_hid_write_blocking(iface, pb_build(&ctx), 64, 1);
+    pb_add_varint(&ctx, 1, offset);
+    pb_add_varint(&ctx, 2, length);
+    pb_end(&ctx);
+    usb_hid_write_blocking(iface, ctx.buf, ctx.pos, 1);
 }

micropython/bootloader/protobuf.c

@@ -9,52 +9,65 @@ void pb_start(PB_CTX *ctx, uint16_t msg_id)
     ctx->buf[2] = '#';
     ctx->buf[3] = (msg_id >> 8) & 0xFF;
     ctx->buf[4] = msg_id & 0xFF;
-    ctx->size = 9;
+    ctx->pos = 9;
+    ctx->len = 0;
 }
 
-const uint8_t *pb_build(PB_CTX *ctx)
+void pb_end(PB_CTX *ctx)
 {
-    ctx->buf[5] = (ctx->size >> 24) & 0xFF;
-    ctx->buf[6] = (ctx->size >> 16) & 0xFF;
-    ctx->buf[7] = (ctx->size >> 8) & 0xFF;
-    ctx->buf[8] = ctx->size & 0xFF;
-    return ctx->buf;
+    ctx->buf[5] = (ctx->len >> 24) & 0xFF;
+    ctx->buf[6] = (ctx->len >> 16) & 0xFF;
+    ctx->buf[7] = (ctx->len >> 8) & 0xFF;
+    ctx->buf[8] = ctx->len & 0xFF;
+    // align to 64 bytes
+    ctx->pos += (-ctx->pos) & 63;
 }
 
-static void pb_add_id(PB_CTX *ctx, const char *id)
+inline static void pb_append(PB_CTX *ctx, uint8_t b)
 {
-    size_t len = strlen(id);
-    memcpy(ctx->buf + ctx->size, id, len);
-    ctx->buf[ctx->size] += len;
+    ctx->buf[ctx->pos] = b;
+    ctx->pos++;
+    if (ctx->pos % 64 == 0) {
+        ctx->buf[ctx->pos] = '?';
+        ctx->pos++;
+    }
+    ctx->len++;
 }
 
-void pb_add_bool(PB_CTX *ctx, const char *id, bool val)
+static void pb_varint(PB_CTX *ctx, uint32_t val)
 {
-    pb_add_id(ctx, id);
-    ctx->buf[ctx->size] = val;
-    ctx->size++;
-}
-
-void pb_add_string(PB_CTX *ctx, const char *id, const char *val)
-{
-    pb_add_varint(ctx, id, strlen(val));
-    size_t len = strlen(val);
-    memcpy(ctx->buf + ctx->size, val, len);
-    ctx->buf[ctx->size] += len;
-}
-
-void pb_add_varint(PB_CTX *ctx, const char *id, uint32_t val)
-{
-    pb_add_id(ctx, id);
     for (;;) {
         if (val < 0x80) {
-            ctx->buf[ctx->size] = val;
-            ctx->size++;
+            pb_append(ctx, val & 0x7F);
             break;
         } else {
-            ctx->buf[ctx->size] = (val & 0x7F) | 0x80;
-            ctx->size++;
+            pb_append(ctx, (val & 0x7F) | 0x80);
             val >>= 7;
         }
     }
 }
+
+void pb_add_bool(PB_CTX *ctx, uint32_t field_number, bool val)
+{
+    field_number = (field_number << 3) | 0;
+    pb_varint(ctx, field_number);
+    pb_append(ctx, val);
+}
+
+void pb_add_string(PB_CTX *ctx, uint32_t field_number, const char *val)
+{
+    field_number = (field_number << 3) | 2;
+    pb_varint(ctx, field_number);
+    size_t len = strlen(val);
+    pb_varint(ctx, len);
+    for (size_t i = 0; i < len; i++) {
+        pb_append(ctx, val[i]);
+    }
+}
+
+void pb_add_varint(PB_CTX *ctx, uint32_t field_number, uint32_t val)
+{
+    field_number = (field_number << 3) | 0;
+    pb_varint(ctx, field_number);
+    pb_varint(ctx, val);
+}

micropython/bootloader/protobuf.h

@@ -5,15 +5,16 @@
 #include <stdbool.h>
 
 typedef struct {
-    uint8_t buf[64];
-    uint32_t size;
+    uint8_t buf[128];
+    uint32_t pos;
+    uint32_t len;
 } PB_CTX;
 
 void pb_start(PB_CTX *ctx, uint16_t msg_id);
-const uint8_t *pb_build(PB_CTX *ctx);
+void pb_end(PB_CTX *ctx);
 
-void pb_add_bool(PB_CTX *ctx, const char *id, bool val);
-void pb_add_string(PB_CTX *ctx, const char *id, const char *val);
-void pb_add_varint(PB_CTX *ctx, const char *id, uint32_t val);
+void pb_add_bool(PB_CTX *ctx, uint32_t field_number, bool val);
+void pb_add_string(PB_CTX *ctx, uint32_t field_number, const char *val);
+void pb_add_varint(PB_CTX *ctx, uint32_t field_number, uint32_t val);
 
 #endif