Merge pull request #573 from LedgerHQ/misc-fixes

### src/boilerplate/io_ext.c
@@ -112,6 +112,13 @@ void app_ticker_event_callback(void) {
 }
 
 void ioe_add_to_response(const void *rdata, size_t rdata_len) {
+    if (G_output_len > IO_APDU_BUFFER_SIZE - 2) {
+        // No room left for the status word
+        G_output_len = IO_APDU_BUFFER_SIZE;
+        write_u16_be(G_io_apdu_buffer, IO_APDU_BUFFER_SIZE - 2, SW_WRONG_RESPONSE_LENGTH);
+        return;
+    }
+
     size_t remaining = (IO_APDU_BUFFER_SIZE - 2) - G_output_len;
 
     if (rdata_len > remaining) {
@@ -126,7 +133,7 @@ void ioe_add_to_response(const void *rdata, size_t rdata_len) {
 }
 
 void ioe_finalize_response(uint16_t sw) {
-    if (G_output_len >= IO_APDU_BUFFER_SIZE - 2) {
+    if (G_output_len > IO_APDU_BUFFER_SIZE - 2) {
         G_output_len = IO_APDU_BUFFER_SIZE;
         write_u16_be(G_io_apdu_buffer, IO_APDU_BUFFER_SIZE - 2, SW_WRONG_RESPONSE_LENGTH);
     } else {

### src/common/wallet.c
@@ -1796,6 +1796,13 @@ static int parse_script(buffer_t *in_buf,
                 return WITH_ERROR(-1, "Invalid k and/or n");
             }
 
+            // In legacy P2SH, the redeem script must not exceed the 520-byte consensus limit
+            if ((context_flags & CONTEXT_WITHIN_SH) != 0 &&
+                (context_flags & CONTEXT_WITHIN_WSH) == 0 &&
+                node->n > MAX_PUBKEYS_PER_P2SH_MULTISIG) {
+                return WITH_ERROR(-1, "Too many keys for multisig in legacy P2SH");
+            }
+
             if (token == TOKEN_SORTEDMULTI || token == TOKEN_SORTEDMULTI_A) {
                 node->base.flags.is_miniscript = 0;
             } else if (token == TOKEN_MULTI) {

### src/common/wallet.h
@@ -18,6 +18,12 @@
 // bitcoin-core supports up to 20, but we limit to 16 as bigger pushes require special handling.
 #define MAX_PUBKEYS_PER_MULTISIG 16
 
+// The maximum number of keys for multi/sortedmulti in legacy P2SH: the redeem script
+// (3 + 34*n bytes) must not exceed the 520-byte consensus limit for stack elements.
+#define MAX_PUBKEYS_PER_P2SH_MULTISIG 15
+_Static_assert(3 + 34 * MAX_PUBKEYS_PER_P2SH_MULTISIG <= 520,
+               "P2SH multisig redeem script exceeds 520 bytes");
+
 // The maximum number of keys supported in a musig() key expression
 // It is basically unlimited in theory, but we need to set a practical limit.
 // The implementation of MuSig2 requires quite a few large arrays (for example, the pubnonces are

### src/error_codes.h
@@ -18,6 +18,12 @@
 // No key in the wallet policy was recognized as internal.
 #define EC_REGISTER_WALLET_POLICY_HAS_NO_INTERNAL_KEY 0x0002
 
+// A musig() key expression in the wallet policy has more than one internal key.
+#define EC_REGISTER_WALLET_MUSIG_WITH_MULTIPLE_INTERNAL_KEYS 0x0003
+
+// The wallet policy has more key expressions with internal keys than supported for signing.
+#define EC_REGISTER_WALLET_TOO_MANY_INTERNAL_KEY_EXPRESSIONS 0x0004
+
 /**
  * SIGN_PSBT
  */

### src/handler/register_wallet.c
@@ -21,6 +21,7 @@
 /* SDK headers */
 #include "bip32.h"
 #include "cx.h"
+#include "ledger_assert.h"
 #include "os.h"
 #include "read.h"
 #include "write.h"
@@ -45,6 +46,10 @@
 
 static bool is_policy_acceptable(const policy_node_t *policy);
 static bool is_policy_name_acceptable(const char *name, size_t name_len);
+static bool check_policy_is_signable(
+    dispatcher_context_t *dc,
+    const policy_node_t *policy,
+    const key_type_e keys_type[static MAX_N_KEYS_IN_WALLET_POLICY]);
 
 static const uint8_t BIP0341_NUMS_PUBKEY[] = {0x02, 0x50, 0x92, 0x9b, 0x74, 0xc1, 0xa0, 0x49, 0x54,
                                               0xb7, 0x8b, 0x4b, 0x60, 0x35, 0xe9, 0x7a, 0x5e, 0x07,
@@ -169,6 +174,11 @@ __attribute__((noinline)) static void confirm_and_register_wallet(
         return;
     }
 
+    // Make sure not to register policies that we would later be unable to sign for
+    if (!check_policy_is_signable(dc, policy, keys_type)) {
+        return;
+    }
+
     // Try to compute the cleartext spending-path lines. If the descriptor
     // doesn't classify (DC_OTHER), has non-canonical derivations, or its
     // confusion score exceeds the threshold, the cleartext block is skipped
@@ -334,3 +344,65 @@ static bool is_policy_name_acceptable(const char *name, size_t name_len) {
 
     return true;
 }
+
+/**
+ * Performs the checks needed to make sure that the device will later be able to sign for the
+ * policy, beyond its validity. Rejects the policy if:
+ * - a musig() key expression has more than one internal key: MuSig2 signing produces at most one
+ *   partial signature per key expression, so such a musig() would be tricky to spend (and it is
+ *   not useful anyway);
+ * - there are more than MAX_INTERNAL_KEY_EXPRESSIONS key expressions with internal keys; signing
+ *   only supports up to that number,
+ *
+ * Returns true on success, false if at least one condition is violated.
+ * It will send the error status before returning.
+ */
+static bool check_policy_is_signable(
+    dispatcher_context_t *dc,
+    const policy_node_t *policy,
+    const key_type_e keys_type[static MAX_N_KEYS_IN_WALLET_POLICY]) {
+    int n_key_expressions = get_keyexpr_by_index(policy, 0, NULL, NULL);
+    LEDGER_ASSERT(n_key_expressions >= 0, "Unexpected error retrieving key expression");
+
+    size_t n_internal_key_expressions = 0;
+    for (int i = 0; i < n_key_expressions; i++) {
+        policy_node_keyexpr_t *key_expr;
+        int ret = get_keyexpr_by_index(policy, i, NULL, &key_expr);
+        LEDGER_ASSERT(ret >= 0, "Unexpected error retrieving key expression");
+
+        int n_internal_keys = 0;
+        if (key_expr->type == KEY_EXPRESSION_NORMAL) {
+            if (keys_type[key_expr->k.key_index] == PUBKEY_TYPE_INTERNAL) {
+                ++n_internal_keys;
+            }
+        } else if (key_expr->type == KEY_EXPRESSION_MUSIG) {
+            const musig_aggr_key_info_t *musig_info = key_expr->m.musig_info;
+            for (int j = 0; j < musig_info->n; j++) {
+                if (keys_type[musig_info->key_indexes[j]] == PUBKEY_TYPE_INTERNAL) {
+                    ++n_internal_keys;
+                }
+            }
+            if (n_internal_keys > 1) {
+                PRINTF("A musig() key expression has more than one internal key\n");
+                SEND_SW_EC(dc,
+                           SW_NOT_SUPPORTED,
+                           EC_REGISTER_WALLET_MUSIG_WITH_MULTIPLE_INTERNAL_KEYS);
+                return false;
+            }
+        } else {
+            LEDGER_ASSERT(false, "Unexpected key expression type");
+        }
+
+        if (n_internal_keys > 0) {
+            ++n_internal_key_expressions;
+        }
+    }
+
+    if (n_internal_key_expressions > MAX_INTERNAL_KEY_EXPRESSIONS) {
+        PRINTF("Too many internal key expressions. The maximum supported is %d\n",
+               MAX_INTERNAL_KEY_EXPRESSIONS);
+        SEND_SW_EC(dc, SW_NOT_SUPPORTED, EC_REGISTER_WALLET_TOO_MANY_INTERNAL_KEY_EXPRESSIONS);
+        return false;
+    }
+    return true;
+}

### tests/test_register_wallet.py
@@ -227,6 +227,80 @@ def test_register_wallet_unsupported_policy(navigator: Navigator, firmware: Firm
     assert len(e.value.data) == 0
 
 
+def test_register_wallet_musig_multiple_internal_keys(client: RaggerClient):
+    # The device can only produce one partial signature per musig() key expression, therefore
+    # musig() key expressions with more than one internal key are not supported.
+
+    # defined in error_codes.h
+    EC_REGISTER_WALLET_MUSIG_WITH_MULTIPLE_INTERNAL_KEYS = 0x0003
+
+    keys_info = [
+        # @0 and @1 are both internal
+        "[f5acc2fd/48'/1'/0'/2']tpubDFAqEGNyad35aBCKUAXbQGDjdVhNueno5ZZVEn3sQbW5ci457gLR7HyTmHBg93oourBssgUxuWz1jX5uhc1qaqFo9VsybY1J5FuedLfm4dK",
+        "[f5acc2fd/48'/1'/0'/1']tpubDFAqEGNyad35YgH8zxvxFZqNUoPtr5mDojs7wzbXQBHTZ4xHeVXG6w2HvsKvjBpaRpTmjYDjdPg5w2c6Wvu8QBkyMDrmBWdCyqkDM7reSsY",
+        "tpubDE7NQymr4AFtewpAsWtnreyq9ghkzQBXpCZjWLFVRAvnbf7vya2eMTvT2fPapNqL8SuVvLQdbUbMfWLVDCZKnsEBqp6UK93QEzL8Ck23AwF",
+    ]
+
+    for descriptor_template in [
+        "tr(musig(@0,@1)/**,pk(@2/**))",
+        "tr(musig(@2,@0,@1)/**)",
+        "tr(@0/**,pk(musig(@1,@2,@0)/**))",
+        "tr(@2/**,multi_a(1,@0/<2;3>/*,musig(@0,@1)/**))",
+    ]:
+        with pytest.raises(ExceptionRAPDU) as e:
+            client.register_wallet(WalletPolicy(
+                name="Musig with 2 internal keys",
+                descriptor_template=descriptor_template,
+                keys_info=keys_info,
+            ))
+
+        assert DeviceException.exc.get(e.value.status) == NotSupportedError
+        assert len(e.value.data) == 2
+        error_code = int.from_bytes(e.value.data, 'big')
+        assert error_code == EC_REGISTER_WALLET_MUSIG_WITH_MULTIPLE_INTERNAL_KEYS
+
+
+def test_register_wallet_too_many_internal_key_expressions(client: RaggerClient):
+    # The same internal key can appear in many key expressions (with disjoint derivations), but
+    # signing supports at most MAX_INTERNAL_KEY_EXPRESSIONS (8) key expressions with internal keys,
+    # counting musig() key expressions that contain an internal key. Registration must reject
+    # policies with more of them.
+
+    # defined in error_codes.h
+    EC_REGISTER_WALLET_TOO_MANY_INTERNAL_KEY_EXPRESSIONS = 0x0004
+
+    internal_key = "[f5acc2fd/48'/1'/0'/2']tpubDFAqEGNyad35aBCKUAXbQGDjdVhNueno5ZZVEn3sQbW5ci457gLR7HyTmHBg93oourBssgUxuWz1jX5uhc1qaqFo9VsybY1J5FuedLfm4dK"
+    external_key = "tpubDE7NQymr4AFtewpAsWtnreyq9ghkzQBXpCZjWLFVRAvnbf7vya2eMTvT2fPapNqL8SuVvLQdbUbMfWLVDCZKnsEBqp6UK93QEzL8Ck23AwF"
+
+    # @0/<2;3>/*, @0/<4;5>/*, ..., @0/<16;17>/*
+    keys = [f"@0/<{2*i};{2*i+1}>/*" for i in range(1, 9)]
+    leaves = [f"pk({k})" for k in keys]
+    tree = (f"{{{{{{{leaves[0]},{leaves[1]}}},{{{leaves[2]},{leaves[3]}}}}},"
+            f"{{{{{leaves[4]},{leaves[5]}}},{{{leaves[6]},{leaves[7]}}}}}}}")
+
+    for descriptor_template, keys_info in [
+        # key path + 8 leaves
+        (f"tr(@0/<0;1>/*,{tree})", [internal_key]),
+        # 9 key expressions in a single leaf
+        (f"tr(@1/**,multi_a(1,@0/<0;1>/*,{','.join(keys)}))", [internal_key, external_key]),
+        # a musig() with an internal key counts as an internal key expression
+        (f"tr(musig(@0,@1)/<0;1>/*,multi_a(1,{','.join(keys)}))", [internal_key, external_key]),
+        # same in segwit v0
+        (f"wsh(multi(1,@0/<0;1>/*,{','.join(keys)}))", [internal_key]),
+    ]:
+        with pytest.raises(ExceptionRAPDU) as e:
+            client.register_wallet(WalletPolicy(
+                name="Too many internal keys",
+                descriptor_template=descriptor_template,
+                keys_info=keys_info,
+            ))
+
+        assert DeviceException.exc.get(e.value.status) == NotSupportedError
+        assert len(e.value.data) == 2
+        error_code = int.from_bytes(e.value.data, 'big')
+        assert error_code == EC_REGISTER_WALLET_TOO_MANY_INTERNAL_KEY_EXPRESSIONS
+
+
 def test_register_miniscript_long_policy(navigator: Navigator, firmware: Firmware, client:
                                          RaggerClient, test_name: str, speculos_globals):
     # This test makes sure that policies longer than 256 bytes work as expected on all devices

### unit-tests/test_wallet.c
@@ -523,6 +523,33 @@ static void test_failures(void **state) {
                                  "@17/**))",
                                  out,
                                  sizeof(out)));
+
+    // legacy P2SH multisig is limited to 15 keys, as the redeem script can't exceed 520 bytes
+    assert_true(0 > parse_policy("sh(multi(2,@0/**,@1/**,@2/**,@3/**,@4/**,@5/**,@6/**,@7/**,"
+                                 "@8/**,@9/**,@10/**,@11/**,@12/**,@13/**,@14/**,@15/**))",
+                                 out,
+                                 sizeof(out)));
+    assert_true(0 > parse_policy("sh(sortedmulti(2,@0/**,@1/**,@2/**,@3/**,@4/**,@5/**,@6/**,@7/**,"
+                                 "@8/**,@9/**,@10/**,@11/**,@12/**,@13/**,@14/**,@15/**))",
+                                 out,
+                                 sizeof(out)));
+    assert_true(0 <= parse_policy("sh(multi(2,@0/**,@1/**,@2/**,@3/**,@4/**,@5/**,@6/**,@7/**,"
+                                  "@8/**,@9/**,@10/**,@11/**,@12/**,@13/**,@14/**))",
+                                  out,
+                                  sizeof(out)));
+    assert_true(0 <= parse_policy("sh(sortedmulti(2,@0/**,@1/**,@2/**,@3/**,@4/**,@5/**,@6/**,"
+                                  "@7/**,@8/**,@9/**,@10/**,@11/**,@12/**,@13/**,@14/**))",
+                                  out,
+                                  sizeof(out)));
+    // no such limit for segwit
+    assert_true(0 <= parse_policy("wsh(multi(2,@0/**,@1/**,@2/**,@3/**,@4/**,@5/**,@6/**,@7/**,"
+                                  "@8/**,@9/**,@10/**,@11/**,@12/**,@13/**,@14/**,@15/**))",
+                                  out,
+                                  sizeof(out)));
+    assert_true(0 <= parse_policy("sh(wsh(sortedmulti(2,@0/**,@1/**,@2/**,@3/**,@4/**,@5/**,@6/**,"
+                                  "@7/**,@8/**,@9/**,@10/**,@11/**,@12/**,@13/**,@14/**,@15/**)))",
+                                  out,
+                                  sizeof(out)));
 }
 
 enum TestMode {