From 214ec2b9aef6d662f3b5cd66b7bdd51f91040b00 Mon Sep 17 00:00:00 2001
From: Pol Henarejos <pol.henarejos@cttc.es>
Date: Fri, 2 Sep 2022 18:02:31 +0200
Subject: [PATCH] Adding support for U2F_MSG

---
 src/apdu.c        |  2 +-
 src/apdu.h        |  2 +-
 src/usb/hid/hid.c | 95 +++++++++++++++++++++++++++++++++++++++++++----
 3 files changed, 90 insertions(+), 9 deletions(-)

diff --git a/src/apdu.c b/src/apdu.c
index 91ca9ca..2f73961 100644
--- a/src/apdu.c
+++ b/src/apdu.c
@@ -1,5 +1,5 @@
 /*
- * This file is part of the Pico HSM distribution (https://github.com/polhenarejos/pico-hsm).
+ * This file is part of the Pico HSM SDK distribution (https://github.com/polhenarejos/pico-hsm-sdk).
  * Copyright (c) 2022 Pol Henarejos.
  *
  * This program is free software: you can redistribute it and/or modify
diff --git a/src/apdu.h b/src/apdu.h
index 5fb9eb5..d8533c8 100644
--- a/src/apdu.h
+++ b/src/apdu.h
@@ -1,5 +1,5 @@
 /*
- * This file is part of the Pico HSM distribution (https://github.com/polhenarejos/pico-hsm).
+ * This file is part of the Pico HSM SDK distribution (https://github.com/polhenarejos/pico-hsm-sdk).
  * Copyright (c) 2022 Pol Henarejos.
  *
  * This program is free software: you can redistribute it and/or modify
diff --git a/src/usb/hid/hid.c b/src/usb/hid/hid.c
index 07b1a76..911e528 100644
--- a/src/usb/hid/hid.c
+++ b/src/usb/hid/hid.c
@@ -22,8 +22,18 @@
 #include "apdu.h"
 #include "usb.h"
 
+#define U2F_MAX_PACKET_SIZE (64 - 7 + 128 * (64 - 5))
+
 static bool mounted = false;
 
+typedef struct msg_packet {
+    uint16_t len;
+    uint16_t current_len;
+    uint8_t data[U2F_MAX_PACKET_SIZE];
+} __packed msg_packet_t;
+
+msg_packet_t msg_packet;
+
 void tud_mount_cb()
 {
     mounted = true;
@@ -91,7 +101,7 @@ void tud_hid_set_report_cb(uint8_t itf, uint8_t report_id, hid_report_type_t rep
 
 void hid_write_offset(uint16_t size, uint16_t offset) {
     if (*usb_get_tx() != 0x81)
-        DEBUG_PAYLOAD(usb_get_tx()+offset,size+10);
+        DEBUG_PAYLOAD(usb_get_tx()+offset, size);
     usb_write_offset(size, offset);
 }
 
@@ -99,13 +109,39 @@ void hid_write(uint16_t size) {
     hid_write_offset(size, 0);
 }
 
+void u2f_error(uint8_t error) {
+    u2f_resp->cid = u2f_req->cid;
+    u2f_resp->init.cmd = U2FHID_ERROR;
+    u2f_resp->init.bcntl = 1;
+    u2f_resp->init.data[0] = ERR_INVALID_CMD;
+    hid_write(64);
+}
+
+uint8_t last_cmd = 0;
+
 int driver_process_usb_packet(uint16_t read) {
-    if (read >= 10)
+    int apdu_sent = 0;
+    if (read >= 5)
     {
-        if (FRAME_TYPE(u2f_req) == TYPE_INIT) {
+        DEBUG_PAYLOAD(usb_get_rx(),64);
+        memset(u2f_resp, 0, sizeof(U2FHID_FRAME));
+        if (FRAME_TYPE(u2f_req) == TYPE_INIT)
+        {
             printf("command %x\n", FRAME_CMD(u2f_req));
             printf("len %d\n", MSG_LEN(u2f_req));
-            DEBUG_PAYLOAD(u2f_req->init.data, MSG_LEN(u2f_req));
+            msg_packet.len = 0;
+            if (MSG_LEN(u2f_req) > 64 - 7)
+            {
+                msg_packet.current_len = 0;
+                msg_packet.len = MSG_LEN(u2f_req);
+                memcpy(msg_packet.data + msg_packet.current_len, u2f_req->init.data, 64-7);
+                msg_packet.current_len += 64 - 7;
+            }
+            last_cmd = u2f_req->init.cmd;
+        }
+        else {
+            memcpy(msg_packet.data + msg_packet.current_len, u2f_req->cont.data, MIN(64 - 5, msg_packet.len - msg_packet.current_len));
+            msg_packet.current_len += MIN(64 - 5, msg_packet.len - msg_packet.current_len);
         }
         if (u2f_req->init.cmd == U2FHID_INIT) {
             U2FHID_INIT_REQ *req = (U2FHID_INIT_REQ *)u2f_req->init.data;
@@ -122,20 +158,65 @@ int driver_process_usb_packet(uint16_t read) {
             u2f_resp->init.bcntl = 17;
             u2f_resp->init.bcnth = 0;
             hid_write(64);
+            card_start();
             DEBUG_PAYLOAD((uint8_t *)u2f_resp, u2f_resp->init.bcntl+7);
         }
+        else if (u2f_req->init.cmd == U2FHID_WINK) {
+            if (MSG_LEN(u2f_req) != 0) {
+                u2f_error(ERR_INVALID_LEN);
+            }
+            memcpy(u2f_resp, u2f_req, sizeof(U2FHID_FRAME));
+            hid_write(64);
+        }
+        else if ((u2f_req->init.cmd == U2FHID_MSG && msg_packet.len == 0) || (msg_packet.len == msg_packet.current_len && msg_packet.len > 0)) {
+            if (msg_packet.current_len == msg_packet.len)
+                apdu_sent = apdu_process(msg_packet.data, msg_packet.len);
+            else
+                apdu_sent = apdu_process(u2f_req->init.data, MSG_LEN(u2f_req));
+             DEBUG_PAYLOAD(apdu.data, (int)apdu.nc);
+        }
+        else {
+            u2f_error(ERR_INVALID_CMD);
+        }
         // echo back anything we received from host
         //tud_hid_report(0, buffer, bufsize);
         printf("END\n");
         usb_clear_rx();
     }
-    return 0;
+    return apdu_sent;
 }
 
 void driver_exec_timeout() {
 
 }
 
-void driver_exec_finished(size_t size_next) {
-
+uint8_t *driver_prepare_response() {
+    u2f_resp = (U2FHID_FRAME *)usb_get_tx();
+    apdu.rdata = u2f_resp->init.data;
+    return u2f_resp->init.data;
+}
+
+void driver_exec_finished(size_t size_next) {
+    driver_exec_finished_cont(size_next, 7);
+}
+
+void driver_exec_finished_cont(size_t size_next, size_t offset) {
+    offset -= 7;
+    u2f_resp = (U2FHID_FRAME *)(usb_get_tx() + offset);
+    u2f_resp->cid = u2f_req->cid;
+    u2f_resp->init.cmd = last_cmd;
+    u2f_resp->init.bcnth = size_next >> 8;
+    u2f_resp->init.bcntl = size_next & 0xff;
+    hid_write_offset(64, offset);
+    size_next -= MIN(64-7, size_next);
+    u2f_resp += 64;
+    uint8_t seq = 0;
+    while (size_next > 0)
+    {
+        u2f_resp->cid = u2f_req->cid;
+        u2f_resp->cont.seq = seq++;
+        hid_write_offset(64, (uint8_t *)u2f_resp-(usb_get_tx()+offset));
+        size_next -= MIN(64 - 5, size_next);
+        u2f_resp += 64;
+    }
 }