rlp: add back Iterator.Count, with fixes

rlp/iterator.go

@@ -69,6 +69,15 @@ func (it *Iterator) Value() []byte {
 	return it.next
 }
 
+// Count returns the remaining number of items.
+// Note this is O(n) and the result may be incorrect if the list data is invalid.
+// The returned count is always an upper bound on the remaining items
+// that will be visited by the iterator.
+func (it *Iterator) Count() int {
+	count, _ := CountValues(it.data)
+	return count
+}
+
 // Offset returns the offset of the current value into the list data.
 func (it *Iterator) Offset() int {
 	return it.offset - len(it.next)

rlp/iterator_test.go

@@ -17,6 +17,7 @@
 package rlp
 
 import (
+	"io"
 	"testing"
 
 	"github.com/ethereum/go-ethereum/common/hexutil"
@@ -54,6 +55,9 @@ func TestIterator(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
+	if c := txit.Count(); c != 2 {
+		t.Fatal("wrong Count:", c)
+	}
 	var i = 0
 	for txit.Next() {
 		if txit.err != nil {
@@ -65,3 +69,65 @@ func TestIterator(t *testing.T) {
 		t.Errorf("count wrong, expected %d got %d", i, exp)
 	}
 }
+
+func TestIteratorErrors(t *testing.T) {
+	tests := []struct {
+		input     []byte
+		wantCount int // expected Count before iterating
+		wantErr   error
+	}{
+		// Second item string header claims 3 bytes content, but only 2 remain.
+		{unhex("C4 01 83AABB"), 2, ErrValueTooLarge},
+		// Second item truncated: B9 requires 2 size bytes, none available.
+		{unhex("C2 01 B9"), 2, io.ErrUnexpectedEOF},
+		// 0x05 should be encoded directly, not as 81 05.
+		{unhex("C3 01 8105"), 2, ErrCanonSize},
+		// Long-form string header B8 used for 1-byte content (< 56).
+		{unhex("C4 01 B801AA"), 2, ErrCanonSize},
+		// Long-form list header F8 used for 1-byte content (< 56).
+		{unhex("C4 01 F80101"), 2, ErrCanonSize},
+	}
+	for _, tt := range tests {
+		it, err := NewListIterator(tt.input)
+		if err != nil {
+			t.Fatal("NewListIterator error:", err)
+		}
+		if c := it.Count(); c != tt.wantCount {
+			t.Fatalf("%x: Count = %d, want %d", tt.input, c, tt.wantCount)
+		}
+		n := 0
+		for it.Next() {
+			if it.Err() != nil {
+				break
+			}
+			n++
+		}
+		if wantN := tt.wantCount - 1; n != wantN {
+			t.Fatalf("%x: got %d valid items, want %d", tt.input, n, wantN)
+		}
+		if it.Err() != tt.wantErr {
+			t.Fatalf("%x: got error %v, want %v", tt.input, it.Err(), tt.wantErr)
+		}
+		if it.Next() {
+			t.Fatalf("%x: Next returned true after error", tt.input)
+		}
+	}
+}
+
+func FuzzIteratorCount(f *testing.F) {
+	examples := [][]byte{unhex("010203"), unhex("018142"), unhex("01830202")}
+	for _, e := range examples {
+		f.Add(e)
+	}
+	f.Fuzz(func(t *testing.T, in []byte) {
+		it := newIterator(in, 0)
+		count := it.Count()
+		i := 0
+		for it.Next() {
+			i++
+		}
+		if i != count {
+			t.Fatalf("%x: count %d not equal to %d iterations", in, count, i)
+		}
+	})
+}

rlp/raw.go

@@ -285,7 +285,7 @@ func CountValues(b []byte) (int, error) {
 	for ; len(b) > 0; i++ {
 		_, tagsize, size, err := readKind(b)
 		if err != nil {
-			return 0, err
+			return i + 1, err
 		}
 		b = b[tagsize+size:]
 	}

rlp/raw_test.go

@@ -288,9 +288,9 @@ func TestCountValues(t *testing.T) {
 		{"820101 820202 8403030303 04", 4, nil},
 
 		// size errors
-		{"8142", 0, ErrCanonSize},
-		{"01 01 8142", 0, ErrCanonSize},
-		{"02 84020202", 0, ErrValueTooLarge},
+		{"8142", 1, ErrCanonSize},
+		{"01 01 8142", 3, ErrCanonSize},
+		{"02 84020202", 2, ErrValueTooLarge},
 
 		{
 			input: "A12000BF49F440A1CD0527E4D06E2765654C0F56452257516D793A9B8D604DCFDF2AB853F851808D10000000000000000000000000A056E81F171BCC55A6FF8345E692C0F86E5B48E01B996CADC001622FB5E363B421A0C5D2460186F7233C927E7DB2DCC703C0E500B653CA82273B7BFAD8045D85A470",

trie/node.go

@@ -161,11 +161,14 @@ func decodeNodeUnsafe(hash, buf []byte) (node, error) {
 	if err != nil {
 		return nil, fmt.Errorf("decode error: %v", err)
 	}
-	switch c, _ := rlp.CountValues(elems); c {
-	case 2:
+	c, err := rlp.CountValues(elems)
+	switch {
+	case err != nil:
+		return nil, fmt.Errorf("invalid node list: %v", err)
+	case c == 2:
 		n, err := decodeShort(hash, elems)
 		return n, wrapError(err, "short")
-	case 17:
+	case c == 17:
 		n, err := decodeFull(hash, elems)
 		return n, wrapError(err, "full")
 	default: