From 9f4c4b37738710934686403f5e7e060aa9fb2940 Mon Sep 17 00:00:00 2001
From: Alexis Grojean <alexis.grojean@ledger.fr>
Date: Mon, 20 Nov 2023 14:15:17 +0100
Subject: [PATCH] Small improvements in APDU handling

---
 src/main.rs             | 30 +++++++++++-------------------
 tests/test_error_cmd.py | 18 +++++++++---------
 2 files changed, 20 insertions(+), 28 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 2966b8d8..f9db80a9 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -44,6 +44,7 @@ use handlers::{
 
 ledger_device_sdk::set_panic!(ledger_device_sdk::exiting_panic);
 
+use ledger_device_sdk::testing::debug_print;
 // CLA (APDU class byte) for all APDUs.
 const CLA: u8 = 0xe0;
 // P2 for last APDU to receive.
@@ -132,7 +133,7 @@ extern "C" fn sample_main() {
         match ui_menu_main(&mut comm) {
             Event::Command(ins) => match handle_apdu(&mut comm, ins.into(), &mut tx_ctx) {
                 Ok(()) => comm.reply_ok(),
-                Err(sw) => comm.reply(Reply::from(sw)),
+                Err(sw) => comm.reply(sw),
             },
             _ => (),
         }
@@ -140,10 +141,14 @@ extern "C" fn sample_main() {
 }
 
 fn handle_apdu(comm: &mut Comm, ins: Ins, ctx: &mut TxContext) -> Result<(), AppSW> {
-    if comm.rx == 0 {
+    // Reject any APDU that does not have a minimum length of 4 bytes
+    // The APDU must have at least 5 bytes: CLA, INS, P1, P2, Lc
+    if comm.rx < 4 {
         return Err(AppSW::WrongDataLength);
     }
 
+    let data = comm.get_data().map_err(|_| AppSW::WrongDataLength)?;
+
     let apdu_metadata = comm.get_apdu_metadata();
 
     if apdu_metadata.cla != CLA {
@@ -167,14 +172,8 @@ fn handle_apdu(comm: &mut Comm, ins: Ins, ctx: &mut TxContext) -> Result<(), App
             if apdu_metadata.p1 > 1 || apdu_metadata.p2 != 0 {
                 return Err(AppSW::WrongP1P2);
             }
-
-            match comm.get_data() {
-                Ok(data) => {
-                    if data.len() == 0 {
-                        return Err(AppSW::WrongDataLength);
-                    }
-                }
-                Err(_) => return Err(AppSW::WrongDataLength),
+            if data.len() == 0 {
+                return Err(AppSW::WrongDataLength);
             }
 
             return handler_get_public_key(comm, apdu_metadata.p1 == 1);
@@ -186,16 +185,9 @@ fn handle_apdu(comm: &mut Comm, ins: Ins, ctx: &mut TxContext) -> Result<(), App
             {
                 return Err(AppSW::WrongP1P2);
             }
-
-            match comm.get_data() {
-                Ok(data) => {
-                    if data.len() == 0 {
-                        return Err(AppSW::WrongDataLength);
-                    }
-                }
-                Err(_) => return Err(AppSW::WrongDataLength),
+            if data.len() == 0 {
+                return Err(AppSW::WrongDataLength);
             }
-
             return handler_sign_tx(
                 comm,
                 apdu_metadata.p1,
diff --git a/tests/test_error_cmd.py b/tests/test_error_cmd.py
index 277f2f8a..2e45e9dd 100644
--- a/tests/test_error_cmd.py
+++ b/tests/test_error_cmd.py
@@ -34,15 +34,15 @@ def test_wrong_p1p2(backend):
     assert e.value.status == Errors.SW_WRONG_P1P2
 
 # Ensure the app returns an error when a bad data length is used
-# def test_wrong_data_length(backend):
-#     # APDUs must be at least 5 bytes: CLA, INS, P1, P2, Lc.
-#     with pytest.raises(ExceptionRAPDU) as e:
-#         backend.exchange_raw(b"E0030000")
-#     assert e.value.status == Errors.SW_WRONG_DATA_LENGTH
-#     # APDUs advertises a too long length
-#     with pytest.raises(ExceptionRAPDU) as e:
-#         backend.exchange_raw(b"E003000005")
-#     assert e.value.status == Errors.SW_WRONG_DATA_LENGTH
+def test_wrong_data_length(backend):
+    # APDUs must be at least 4 bytes: CLA, INS, P1, P2.
+    with pytest.raises(ExceptionRAPDU) as e:
+        backend.exchange_raw(bytes.fromhex("E00300"))
+    assert e.value.status == Errors.SW_WRONG_DATA_LENGTH
+    # APDUs advertises a too long length
+    with pytest.raises(ExceptionRAPDU) as e:
+        backend.exchange_raw(bytes.fromhex("E003000005"))
+    assert e.value.status == Errors.SW_WRONG_DATA_LENGTH
 
 
 # Ensure there is no state confusion when trying wrong APDU sequences