From 3f294f958288e5f124c32a5563754d2892ce3117 Mon Sep 17 00:00:00 2001
From: Lucas PASCAL <lucas.pascal@ledger.com>
Date: Fri, 16 Sep 2022 14:43:51 +0200
Subject: [PATCH] Minor fixes after review

---
 src/apdu_constants.h                       |  3 +++
 src/main.c                                 |  6 ++++--
 src_features/signMessage/cmd_signMessage.c | 10 ++++++----
 3 files changed, 13 insertions(+), 6 deletions(-)

diff --git a/src/apdu_constants.h b/src/apdu_constants.h
index ce1c96a..219b7d0 100644
--- a/src/apdu_constants.h
+++ b/src/apdu_constants.h
@@ -58,6 +58,9 @@
 #define OFFSET_LC    4
 #define OFFSET_CDATA 5
 
+#define ERR_APDU_EMPTY         0x6982
+#define ERR_APDU_SIZE_MISMATCH 0x6983
+
 void handleGetPublicKey(uint8_t p1,
                         uint8_t p2,
                         const uint8_t *dataBuffer,
diff --git a/src/main.c b/src/main.c
index fa3db8e..2a32e22 100644
--- a/src/main.c
+++ b/src/main.c
@@ -770,9 +770,11 @@ void app_main(void) {
                 // no apdu received, well, reset the session, and reset the
                 // bootloader configuration
                 if (rx == 0) {
-                    THROW(0x6982);
+                    THROW(ERR_APDU_EMPTY);
+                }
+                if (rx > OFFSET_LC && rx != (G_io_apdu_buffer[OFFSET_LC] + 5)) {
+                    THROW(ERR_APDU_SIZE_MISMATCH);
                 }
-                PRINTF("New APDU received:\n%.*H\n", rx, G_io_apdu_buffer);
 
                 handleApdu(&flags, &tx);
             }
diff --git a/src_features/signMessage/cmd_signMessage.c b/src_features/signMessage/cmd_signMessage.c
index 4c257b7..4b4008f 100644
--- a/src_features/signMessage/cmd_signMessage.c
+++ b/src_features/signMessage/cmd_signMessage.c
@@ -21,7 +21,8 @@ static const char SIGN_MAGIC[] =
  * @param[in] sw status word
  */
 static void apdu_reply(uint16_t sw) {
-    *(uint16_t *) G_io_apdu_buffer = __builtin_bswap16(sw);
+    G_io_apdu_buffer[0] = (sw >> 8) & 0xff;
+    G_io_apdu_buffer[1] = sw & 0xff;
     io_exchange(CHANNEL_APDU | IO_RETURN_AFTER_TX, 2);
 }
 
@@ -134,7 +135,7 @@ static const uint8_t *first_apdu_data(const uint8_t *data, uint16_t *length) {
  * @param[in] length the data length
  * @return whether it was successful or not
  */
-static bool feed_hash(const uint8_t *const data, uint8_t length) {
+static bool feed_hash(const uint8_t *const data, const uint8_t length) {
     if (length > tmpCtx.messageSigningContext.remainingLength) {
         PRINTF("Error: Length mismatch ! (%u > %u)!\n",
                length,
@@ -217,11 +218,12 @@ bool handleSignPersonalMessage(uint8_t p1,
                                const uint8_t *const payload,
                                uint8_t length) {
     const uint8_t *data = payload;
+    uint16_t u16_length = length;
 
     (void) p2;
     processed_size = 0;
     if (p1 == P1_FIRST) {
-        if ((data = first_apdu_data(data, (uint16_t *) &length)) == NULL) {
+        if ((data = first_apdu_data(data, &u16_length)) == NULL) {
             return false;
         }
         processed_size = data - payload;
@@ -231,7 +233,7 @@ bool handleSignPersonalMessage(uint8_t p1,
         return false;
     }
 
-    if (!feed_hash(data, length)) {
+    if (!feed_hash(data, u16_length)) {
         return false;
     }