Fix proof generation and verification code

proof_ipa.go

@@ -387,7 +387,10 @@ type stemInfo struct {
 // PreStateTreeFromProof builds a stateless prestate tree from the proof.
 func PreStateTreeFromProof(proof *Proof, rootC *Point) (VerkleNode, error) { // skipcq: GO-R1005
 	if len(proof.Keys) != len(proof.PreValues) {
-		return nil, fmt.Errorf("incompatible number of keys and values: %d != %d", len(proof.Keys), len(proof.PreValues))
+		return nil, fmt.Errorf("incompatible number of keys and pre-values: %d != %d", len(proof.Keys), len(proof.PreValues))
+	}
+	if len(proof.Keys) != len(proof.PostValues) {
+		return nil, fmt.Errorf("incompatible number of keys and post-values: %d != %d", len(proof.Keys), len(proof.PostValues))
 	}
 	stems := make([][]byte, 0, len(proof.Keys))
 	for _, k := range proof.Keys {
@@ -404,6 +407,11 @@ func PreStateTreeFromProof(proof *Proof, rootC *Point) (VerkleNode, error) { //
 		poas  = proof.PoaStems
 	)
 
+	// The proof of absence stems must be sorted. If that isn't the case, the proof is invalid.
+	if !sort.IsSorted(bytesSlice(proof.PoaStems)) {
+		return nil, fmt.Errorf("proof of absence stems are not sorted")
+	}
+
 	// assign one or more stem to each stem info
 	for _, es := range proof.ExtStatus {
 		depth := es >> 3
@@ -414,25 +422,57 @@ func PreStateTreeFromProof(proof *Proof, rootC *Point) (VerkleNode, error) { //
 		}
 		switch si.stemType {
 		case extStatusAbsentEmpty:
+			// All keys that are part of a proof of absence, must contain empty
+			// prestate values. If that isn't the case, the proof is invalid.
+			for i, k := range proof.Keys { // TODO: DoS risk, use map or binary search.
+				if bytes.HasPrefix(k, path) {
+					if proof.PreValues[i] != nil {
+						return nil, fmt.Errorf("proof of absence (empty) stem %x has a value", si.stem)
+					}
+				}
+			}
 		case extStatusAbsentOther:
 			si.stem = poas[0]
 			poas = poas[1:]
+			// All keys that are part of a proof of absence, must contain empty
+			// prestate values. If that isn't the case, the proof is invalid.
+			for i, k := range proof.Keys { // TODO: DoS risk, use map or binary search.
+				if bytes.HasPrefix(k, si.stem) {
+					if proof.PreValues[i] != nil {
+						return nil, fmt.Errorf("proof of absence (other) stem %x has a value", si.stem)
+					}
+				}
+			}
 		default:
 			// the first stem could be missing (e.g. the second stem in the
 			// group is the one that is present. Compare each key to the first
 			// stem, along the length of the path only.
 			stemPath := stems[stemIndex][:len(path)]
 			si.values = map[byte][]byte{}
-			for i, k := range proof.Keys {
+			for i, k := range proof.Keys { // TODO: DoS risk, use map or binary search.
 				if bytes.Equal(k[:len(path)], stemPath) && proof.PreValues[i] != nil {
 					si.values[k[31]] = proof.PreValues[i]
 					si.has_c1 = si.has_c1 || (k[31] < 128)
 					si.has_c2 = si.has_c2 || (k[31] >= 128)
 					// This key has values, its stem is the one that
 					// is present.
-					si.stem = k[:31]
+					if si.stem == nil {
+						si.stem = k[:31]
+						continue
+					}
+					// Any other key with values must have the same
+					// same previously detected stem. If that isn't the case,
+					// the proof is invalid.
+					if !bytes.Equal(si.stem, k[:31]) {
+						return nil, fmt.Errorf("multiple keys with values found for stem %x", k[:31])
+					}
 				}
 			}
+			// For a proof of presence, we must always have detected a stem.
+			// If that isn't the case, the proof is invalid.
+			if si.stem == nil {
+				return nil, fmt.Errorf("no stem found for path %x", path)
+			}
 		}
 		info[string(path)] = si
 		paths = append(paths, path)
@@ -513,3 +553,9 @@ func PostStateTreeFromStateDiff(preroot VerkleNode, statediff StateDiff) (Verkle
 
 	return postroot, nil
 }
+
+type bytesSlice [][]byte
+
+func (x bytesSlice) Len() int           { return len(x) }
+func (x bytesSlice) Less(i, j int) bool { return bytes.Compare(x[i], x[j]) < 0 }
+func (x bytesSlice) Swap(i, j int)      { x[i], x[j] = x[j], x[i] }

tree.go

@@ -434,7 +434,7 @@ func (n *InternalNode) InsertValuesAtStem(stem []byte, values [][]byte, resolver
 // commitments that have not been assigned a node. It returns
 // the same list, save the commitments that were consumed
 // during this call.
-func (n *InternalNode) CreatePath(path []byte, stemInfo stemInfo, comms []*Point, values [][]byte) ([]*Point, error) {
+func (n *InternalNode) CreatePath(path []byte, stemInfo stemInfo, comms []*Point, values [][]byte) ([]*Point, error) { // skipcq: GO-R1005
 	if len(path) == 0 {
 		return comms, errors.New("invalid path")
 	}
@@ -448,6 +448,12 @@ func (n *InternalNode) CreatePath(path []byte, stemInfo stemInfo, comms []*Point
 			// unknown node.
 			n.children[path[0]] = Empty{}
 		case extStatusAbsentOther:
+			if len(comms) == 0 {
+				return comms, fmt.Errorf("missing commitment for stem %x", stemInfo.stem)
+			}
+			if len(stemInfo.stem) != StemSize {
+				return comms, fmt.Errorf("invalid stem size %d", len(stemInfo.stem))
+			}
 			// insert poa stem
 			newchild := &LeafNode{
 				commitment: comms[0],
@@ -459,6 +465,12 @@ func (n *InternalNode) CreatePath(path []byte, stemInfo stemInfo, comms []*Point
 			n.children[path[0]] = newchild
 			comms = comms[1:]
 		case extStatusPresent:
+			if len(comms) == 0 {
+				return comms, fmt.Errorf("missing commitment for stem %x", stemInfo.stem)
+			}
+			if len(stemInfo.stem) != StemSize {
+				return comms, fmt.Errorf("invalid stem size %d", len(stemInfo.stem))
+			}
 			// insert stem
 			newchild := &LeafNode{
 				commitment: comms[0],
@@ -469,12 +481,18 @@ func (n *InternalNode) CreatePath(path []byte, stemInfo stemInfo, comms []*Point
 			n.children[path[0]] = newchild
 			comms = comms[1:]
 			if stemInfo.has_c1 {
+				if len(comms) == 0 {
+					return comms, fmt.Errorf("missing commitment for stem %x", stemInfo.stem)
+				}
 				newchild.c1 = comms[0]
 				comms = comms[1:]
 			} else {
 				newchild.c1 = new(Point)
 			}
 			if stemInfo.has_c2 {
+				if len(comms) == 0 {
+					return comms, fmt.Errorf("missing commitment for stem %x", stemInfo.stem)
+				}
 				newchild.c2 = comms[0]
 				comms = comms[1:]
 			} else {
@@ -483,6 +501,8 @@ func (n *InternalNode) CreatePath(path []byte, stemInfo stemInfo, comms []*Point
 			for b, value := range stemInfo.values {
 				newchild.values[b] = value
 			}
+		default:
+			return comms, fmt.Errorf("invalid stem type %d", stemInfo.stemType)
 		}
 		return comms, nil
 	}