From 84bd8dae01bf3a2e363390361ce52ab77f23e360 Mon Sep 17 00:00:00 2001
From: Salvatore Ingala <6681844+bigspider@users.noreply.github.com>
Date: Thu, 2 Nov 2023 16:38:03 +0100
Subject: [PATCH] Increase maximum depth of the supported taptrees on non-NanoS
 devices; generalize corresponding test so that it adapts to the different
 limits

---
 src/common/wallet.h                        |  7 ++++++-
 src/handler/lib/check_merkle_tree_sorted.h |  5 +++--
 tests/test_e2e_tapscripts.py               | 13 ++++++++++---
 3 files changed, 19 insertions(+), 6 deletions(-)

diff --git a/src/common/wallet.h b/src/common/wallet.h
index b440eac5..c43b57d0 100644
--- a/src/common/wallet.h
+++ b/src/common/wallet.h
@@ -73,7 +73,12 @@
     MAX(MAX_WALLET_POLICY_SERIALIZED_LENGTH_V1, MAX_WALLET_POLICY_SERIALIZED_LENGTH_V2)
 
 // maximum depth of a taproot tree that we support
-#define MAX_TAPTREE_POLICY_DEPTH 4
+// (here depth 1 means only the root of the taptree)
+#ifdef TARGET_NANOS
+#define MAX_TAPTREE_POLICY_DEPTH 5
+#else
+#define MAX_TAPTREE_POLICY_DEPTH 9
+#endif
 
 typedef struct {
     uint32_t master_key_derivation[MAX_BIP32_PATH_STEPS];
diff --git a/src/handler/lib/check_merkle_tree_sorted.h b/src/handler/lib/check_merkle_tree_sorted.h
index 2cb8993e..15426d68 100644
--- a/src/handler/lib/check_merkle_tree_sorted.h
+++ b/src/handler/lib/check_merkle_tree_sorted.h
@@ -2,11 +2,12 @@
 
 #include "../../boilerplate/dispatcher.h"
 #include "../../common/merkle.h"
+#include "../../common/wallet.h"
 
 // this flow aborts if any element is larger than this size
-// This is enough for a PSBT with the control block of a taptree of depth 5 tapbranches
 // TODO: we might remove this limitation altogether with a more careful implementation.
-#define MAX_CHECK_MERKLE_TREE_SORTED_PREIMAGE_SIZE 162
+// Here we make sure that we have enough space for control block of a taptree of the maximum supported depth
+#define MAX_CHECK_MERKLE_TREE_SORTED_PREIMAGE_SIZE (34 + 32*(MAX_TAPTREE_POLICY_DEPTH - 1))
 
 typedef void (*merkle_tree_elements_callback_t)(struct dispatcher_context_s *,
                                                 void *,
diff --git a/tests/test_e2e_tapscripts.py b/tests/test_e2e_tapscripts.py
index da2aa17c..31997efd 100644
--- a/tests/test_e2e_tapscripts.py
+++ b/tests/test_e2e_tapscripts.py
@@ -236,21 +236,28 @@ def test_e2e_tapscript_multi_a_2of2(rpc, rpc_test_wallet, client: Client, specul
                  rpc, rpc_test_wallet, client, speculos_globals, comm)
 
 
-def test_e2e_tapscript_depth4(rpc, rpc_test_wallet, client: Client, speculos_globals: SpeculosGlobals, comm: Union[TransportClient, SpeculosClient]):
+def test_e2e_tapscript_maxdepth(rpc, rpc_test_wallet, client: Client, speculos_globals: SpeculosGlobals, comm: Union[TransportClient, SpeculosClient], model: str):
     # A taproot tree with maximum supported depth, where the internal key is in the deepest script
 
+    MAX_TAPTREE_POLICY_DEPTH = 4 if model == "nanos" else 9
+
+    # Make the most unbalanced tree where each script is a simple pk()
+    parts = [f"pk(@{i}/**)" for i in range(1, MAX_TAPTREE_POLICY_DEPTH)]
+    descriptor_template = "tr(@0/**,{" + ',{'.join(parts) + f",pk(@{MAX_TAPTREE_POLICY_DEPTH}/**)" + "}" * (MAX_TAPTREE_POLICY_DEPTH - 1) + ")"
+
     keys_info = []
-    for _ in range(4):
+    for _ in range(MAX_TAPTREE_POLICY_DEPTH):
         _, core_xpub_orig = create_new_wallet()
         keys_info.append(core_xpub_orig)
 
+    # the last (deepest) script is the only one we sign with the ledger key
     path = "499'/1'/0'"
     internal_xpub = get_internal_xpub(speculos_globals.seed, path)
     keys_info.append(f"[{speculos_globals.master_key_fingerprint.hex()}/{path}]{internal_xpub}")
 
     wallet_policy = WalletPolicy(
         name="Tapscriptception",
-        descriptor_template="tr(@0/**,{pk(@1/**),{pk(@2/**),{pk(@3/**),pk(@4/**)}}})",
+        descriptor_template=descriptor_template,
         keys_info=keys_info)
 
     run_test_e2e(wallet_policy, [], rpc, rpc_test_wallet, client, speculos_globals, comm)