reject trees with inconsistent key heights, probably

src/atmst/mst/node.py

@@ -31,14 +31,17 @@ class MSTNode:
 
 	# NB: __init__ is auto-generated by dataclass decorator
 
-	# these checks should never fail, and could be skipped for performance
+	# these checks could be elided for "trusted" trees
 	def __post_init__(self) -> None:
 		# TODO: maybe check that they're tuples here?
 		# implicitly, the length of self.subtrees must be at least 1
 		if len(self.subtrees) != len(self.keys) + 1:
 			raise ValueError("Invalid subtree count")
 		if len(self.keys) != len(self.vals):
 			raise ValueError("Mismatched keys/vals lengths")
+		key_heights = [MSTNode.key_height(k) for k in self.keys]
+		if any(k!=key_heights[0] for k in key_heights):
+			raise ValueError("Inconsistent key heights")
 
 	@classmethod
 	def empty_root(cls) -> "Self":

src/atmst/mst/node_walker.py

@@ -38,8 +38,9 @@ class StackFrame:
 	ns: NodeStore
 	stack: List[StackFrame]
 
-	def __init__(self, ns: NodeStore, root_cid: Optional[CID], lpath: Optional[str]=PATH_MIN, rpath: Optional[str]=PATH_MAX) -> None:
+	def __init__(self, ns: NodeStore, root_cid: Optional[CID], lpath: Optional[str]=PATH_MIN, rpath: Optional[str]=PATH_MAX, trusted: Optional[bool]=False) -> None:
 		self.ns = ns
+		self.trusted = trusted
 		self.stack = [self.StackFrame(
 			node=MSTNode.empty_root() if root_cid is None else self.ns.get_node(root_cid),
 			lpath=lpath,
@@ -48,7 +49,7 @@ def __init__(self, ns: NodeStore, root_cid: Optional[CID], lpath: Optional[str]=
 		)]
 
 	def subtree_walker(self) -> "Self":
-		return NodeWalker(self.ns, self.subtree, self.lpath, self.rpath)
+		return NodeWalker(self.ns, self.subtree, self.lpath, self.rpath, self.trusted)
 
 	@property
 	def frame(self) -> StackFrame:
@@ -98,16 +99,22 @@ def right_or_up(self) -> None:
 
 	def right(self) -> None:
 		if not self.can_go_right:
-			raise Exception("cursor is already at rightmost position in node")
+			raise ValueError("cursor is already at rightmost position in node")
 		self.frame.idx += 1
 
 	def down(self) -> None:
 		subtree = self.frame.node.subtrees[self.frame.idx]
 		if subtree is None:
-			raise Exception("oi, you can't recurse here mate (subtree is None)")
+			raise ValueError("oi, you can't recurse here mate (subtree is None)")
+
+		subtree_node = self.ns.get_node(subtree)
+
+		if not self.trusted: # if we "trust" the source we can elide this check
+			if subtree_node.height != self.height - 1:
+				raise ValueError("inconsistent subtree height")
 
 		self.stack.append(self.StackFrame(
-			node=self.ns.get_node(subtree),
+			node=subtree_node,
 			lpath=self.lpath,
 			rpath=self.rpath,
 			idx=0