Use generic trees for recovery

Author: Agusx1211

packages/wallet/primitives/src/extensions/recovery.ts

@@ -1,11 +1,14 @@
 import { Address, Bytes, Hash, Hex } from 'ox'
 import * as Payload from '../payload'
 import { getSignPayload } from 'ox/TypedData'
+import * as GenericTree from '../generic-tree'
 
 export const FLAG_RECOVERY_LEAF = 1
 export const FLAG_NODE = 3
 export const FLAG_BRANCH = 4
 
+const RECOVERY_LEAF_PREFIX = Bytes.fromString('Sequence recovery leaf:\n')
+
 /**
  * A leaf in the Recovery tree, storing:
  *  - signer who can queue a payload
@@ -19,23 +22,18 @@ export type RecoveryLeaf = {
   minTimestamp: bigint
 }
 
-/**
- * A node is just a 32-byte hash
- */
-export type NodeLeaf = Hex.Hex
-
 /**
  * A branch is a list of subtrees (≥2 in length).
  */
-export type Node = [Node, Node]
+export type Branch = [Tree, Tree]
 
 /**
  * The topology of a recovery tree can be either:
  * - A node (pair of subtrees)
  * - A node leaf (32-byte hash)
  * - A recovery leaf (signer with timing constraints)
  */
-export type Topology = Node | NodeLeaf | RecoveryLeaf
+export type Tree = Branch | GenericTree.Node | RecoveryLeaf
 
 /**
  * Type guard to check if a value is a RecoveryLeaf
@@ -44,25 +42,18 @@ export function isRecoveryLeaf(cand: any): cand is RecoveryLeaf {
   return typeof cand === 'object' && cand !== null && cand.type === 'leaf'
 }
 
-/**
- * Type guard to check if a value is a NodeLeaf (32-byte hash)
- */
-export function isNodeLeaf(cand: any): cand is NodeLeaf {
-  return typeof cand === 'string' && cand.length === 66 && cand.startsWith('0x')
-}
-
 /**
  * Type guard to check if a value is a Node (pair of subtrees)
  */
-export function isNode(cand: any): cand is Node {
-  return Array.isArray(cand) && cand.length === 2 && isTopology(cand[0]) && isTopology(cand[1])
+export function isBranch(cand: any): cand is Branch {
+  return Array.isArray(cand) && cand.length === 2 && isTree(cand[0]) && isTree(cand[1])
 }
 
 /**
  * Type guard to check if a value is a Topology
  */
-export function isTopology(cand: any): cand is Topology {
-  return isRecoveryLeaf(cand) || isNodeLeaf(cand) || isNode(cand)
+export function isTree(cand: any): cand is Tree {
+  return isRecoveryLeaf(cand) || GenericTree.isNode(cand) || isBranch(cand)
 }
 
 /**
@@ -79,24 +70,8 @@ export const DOMAIN_VERSION = '1'
  * For node leaves, it returns the hash directly.
  * For nodes, it hashes the concatenation of the hashes of both subtrees.
  */
-export function hashConfiguration(topology: Topology): Hex.Hex {
-  if (isRecoveryLeaf(topology)) {
-    return Hash.keccak256(
-      Bytes.concat(
-        Bytes.fromString('Sequence recovery leaf:\n'),
-        Bytes.fromHex(topology.signer, { size: 20 }),
-        Bytes.padLeft(Bytes.fromNumber(topology.requiredDeltaTime), 32),
-        Bytes.padLeft(Bytes.fromNumber(topology.minTimestamp), 32),
-      ),
-      { as: 'Hex' },
-    )
-  } else if (isNodeLeaf(topology)) {
-    return topology
-  } else if (isNode(topology)) {
-    return Hash.keccak256(Hex.concat(hashConfiguration(topology[0]), hashConfiguration(topology[1])), { as: 'Hex' })
-  } else {
-    throw new Error('Invalid topology')
-  }
+export function hashConfiguration(topology: Tree): Hex.Hex {
+  return GenericTree.hash(toGenericTree(topology))
 }
 
 /**
@@ -107,13 +82,13 @@ export function hashConfiguration(topology: Topology): Hex.Hex {
  * - leaves: Array of RecoveryLeaf nodes
  * - isComplete: boolean indicating if all leaves are present (no node references)
  */
-export function getRecoveryLeaves(topology: Topology): { leaves: RecoveryLeaf[]; isComplete: boolean } {
+export function getRecoveryLeaves(topology: Tree): { leaves: RecoveryLeaf[]; isComplete: boolean } {
   const isComplete = true
   if (isRecoveryLeaf(topology)) {
     return { leaves: [topology], isComplete }
-  } else if (isNodeLeaf(topology)) {
+  } else if (GenericTree.isNode(topology)) {
     return { leaves: [], isComplete: false }
-  } else if (isNode(topology)) {
+  } else if (isBranch(topology)) {
     const left = getRecoveryLeaves(topology[0])
     const right = getRecoveryLeaves(topology[1])
     return { leaves: [...left.leaves, ...right.leaves], isComplete: left.isComplete && right.isComplete }
@@ -129,7 +104,7 @@ export function getRecoveryLeaves(topology: Topology): { leaves: RecoveryLeaf[];
  * @returns The decoded Topology object
  * @throws Error if the encoding is invalid
  */
-export function decodeTopology(encoded: Bytes.Bytes): Topology {
+export function decodeTopology(encoded: Bytes.Bytes): Tree {
   const { nodes, leftover } = parseBranch(encoded)
   if (leftover.length > 0) {
     throw new Error('Leftover bytes in branch')
@@ -146,12 +121,12 @@ export function decodeTopology(encoded: Bytes.Bytes): Topology {
  * - leftover: Any remaining unparsed bytes
  * @throws Error if the encoding is invalid
  */
-export function parseBranch(encoded: Bytes.Bytes): { nodes: Topology[]; leftover: Bytes.Bytes } {
+export function parseBranch(encoded: Bytes.Bytes): { nodes: Tree[]; leftover: Bytes.Bytes } {
   if (encoded.length === 0) {
     throw new Error('Empty branch')
   }
 
-  const nodes: Topology[] = []
+  const nodes: Tree[] = []
   let index = 0
 
   while (index < encoded.length) {
@@ -208,7 +183,7 @@ export function parseBranch(encoded: Bytes.Bytes): { nodes: Topology[]; leftover
  * @param signer - The signer address to keep
  * @returns The trimmed topology
  */
-export function trimTopology(topology: Topology, signer: Address.Address): Topology {
+export function trimTopology(topology: Tree, signer: Address.Address): Tree {
   if (isRecoveryLeaf(topology)) {
     if (topology.signer === signer) {
       return topology
@@ -217,20 +192,20 @@ export function trimTopology(topology: Topology, signer: Address.Address): Topol
     }
   }
 
-  if (isNodeLeaf(topology)) {
+  if (GenericTree.isNode(topology)) {
     return topology
   }
 
-  if (isNode(topology)) {
+  if (isBranch(topology)) {
     const left = trimTopology(topology[0], signer)
     const right = trimTopology(topology[1], signer)
 
     // If both are hashes, we can just return the hash of the node
-    if (isNodeLeaf(left) && isNodeLeaf(right)) {
+    if (GenericTree.isNode(left) && GenericTree.isNode(right)) {
       return hashConfiguration(topology)
     }
 
-    return [left, right] as Node
+    return [left, right] as Branch
   }
 
   throw new Error('Invalid topology')
@@ -243,11 +218,11 @@ export function trimTopology(topology: Topology, signer: Address.Address): Topol
  * @returns The binary encoded topology
  * @throws Error if the topology is invalid
  */
-export function encodeTopology(topology: Topology): Bytes.Bytes {
-  if (isNode(topology)) {
+export function encodeTopology(topology: Tree): Bytes.Bytes {
+  if (isBranch(topology)) {
     const encoded0 = encodeTopology(topology[0]!)
     const encoded1 = encodeTopology(topology[1]!)
-    const isBranching = isNode(topology[1]!)
+    const isBranching = isBranch(topology[1]!)
 
     if (isBranching) {
       // max 3 bytes for the size
@@ -263,7 +238,7 @@ export function encodeTopology(topology: Topology): Bytes.Bytes {
     }
   }
 
-  if (isNodeLeaf(topology)) {
+  if (GenericTree.isNode(topology)) {
     const flag = Bytes.fromNumber(FLAG_NODE)
     const nodeHash = Bytes.fromHex(topology, { size: 32 })
     return Bytes.concat(flag, nodeHash)
@@ -296,7 +271,7 @@ export function encodeTopology(topology: Topology): Bytes.Bytes {
  * @returns A binary tree structure
  * @throws Error if the nodes array is empty
  */
-function foldNodes(nodes: Topology[]): Topology {
+function foldNodes(nodes: Tree[]): Tree {
   if (nodes.length === 0) {
     throw new Error('Empty signature tree')
   }
@@ -305,9 +280,9 @@ function foldNodes(nodes: Topology[]): Topology {
     return nodes[0]!
   }
 
-  let tree: Topology = nodes[0]!
+  let tree: Tree = nodes[0]!
   for (let i = 1; i < nodes.length; i++) {
-    tree = [tree, nodes[i]!] as Topology
+    tree = [tree, nodes[i]!] as Tree
   }
   return tree
 }
@@ -321,7 +296,7 @@ function foldNodes(nodes: Topology[]): Topology {
  * @returns A topology tree structure
  * @throws Error if the leaves array is empty
  */
-export function fromRecoveryLeaves(leaves: RecoveryLeaf[]): Topology {
+export function fromRecoveryLeaves(leaves: RecoveryLeaf[]): Tree {
   if (leaves.length === 0) {
     throw new Error('Cannot build a tree with zero leaves')
   }
@@ -333,7 +308,7 @@ export function fromRecoveryLeaves(leaves: RecoveryLeaf[]): Topology {
   const mid = Math.floor(leaves.length / 2)
   const left = fromRecoveryLeaves(leaves.slice(0, mid))
   const right = fromRecoveryLeaves(leaves.slice(mid))
-  return [left, right] as Node
+  return [left, right] as Branch
 }
 
 /**
@@ -389,3 +364,74 @@ export function hashRecoveryPayload(
   const structHash = Bytes.fromHex(getSignPayload(Payload.toTyped(wallet, noChainId ? 0n : chainId, payload)))
   return Hash.keccak256(Bytes.concat(Bytes.fromString('\x19\x01'), Hex.toBytes(ds), structHash), { as: 'Hex' })
 }
+
+/**
+ * Convert a RecoveryTree topology to a generic tree format
+ *
+ * @param topology - The recovery tree topology to convert
+ * @returns A generic tree that produces the same root hash
+ */
+export function toGenericTree(topology: Tree): GenericTree.Tree {
+  if (isRecoveryLeaf(topology)) {
+    // Convert recovery leaf to generic leaf
+    return {
+      type: 'leaf',
+      value: Bytes.concat(
+        RECOVERY_LEAF_PREFIX,
+        Bytes.fromHex(topology.signer, { size: 20 }),
+        Bytes.padLeft(Bytes.fromNumber(topology.requiredDeltaTime), 32),
+        Bytes.padLeft(Bytes.fromNumber(topology.minTimestamp), 32),
+      ),
+    }
+  } else if (GenericTree.isNode(topology)) {
+    // Node leaves are already in the correct format
+    return topology
+  } else if (isBranch(topology)) {
+    // Convert node to branch
+    return [toGenericTree(topology[0]), toGenericTree(topology[1])]
+  } else {
+    throw new Error('Invalid topology')
+  }
+}
+
+/**
+ * Convert a generic tree back to a RecoveryTree topology
+ *
+ * @param tree - The generic tree to convert
+ * @returns A recovery tree topology that produces the same root hash
+ */
+export function fromGenericTree(tree: GenericTree.Tree): Tree {
+  if (GenericTree.isLeaf(tree)) {
+    // Convert generic leaf back to recovery leaf
+    const bytes = tree.value
+    if (
+      bytes.length !== RECOVERY_LEAF_PREFIX.length + 84 ||
+      !Bytes.isEqual(bytes.slice(0, RECOVERY_LEAF_PREFIX.length), RECOVERY_LEAF_PREFIX)
+    ) {
+      throw new Error('Invalid recovery leaf format')
+    }
+
+    const offset = RECOVERY_LEAF_PREFIX.length
+    const signer = Address.from(Hex.fromBytes(bytes.slice(offset, offset + 20)))
+    const requiredDeltaTime = Bytes.toBigInt(bytes.slice(offset + 20, offset + 52))
+    const minTimestamp = Bytes.toBigInt(bytes.slice(offset + 52, offset + 84))
+
+    return {
+      type: 'leaf',
+      signer,
+      requiredDeltaTime,
+      minTimestamp,
+    }
+  } else if (GenericTree.isNode(tree)) {
+    // Nodes are already in the correct format
+    return tree
+  } else if (GenericTree.isBranch(tree)) {
+    // Convert branch back to node
+    if (tree.length !== 2) {
+      throw new Error('Recovery tree only supports binary branches')
+    }
+    return [fromGenericTree(tree[0]), fromGenericTree(tree[1])] as Branch
+  } else {
+    throw new Error('Invalid tree format')
+  }
+}