feat(cborlite): a light CBOR reader.

core/felt/cbor.go

@@ -2,12 +2,20 @@ package felt
 
 import (
 	"encoding/binary"
+	"errors"
 	"math"
 
+	"github.com/NethermindEth/juno/encoder/cborlite"
 	"github.com/consensys/gnark-crypto/ecc/stark-curve/fp"
 	"github.com/fxamacker/cbor/v2"
 )
 
+var (
+	_ cborlite.PrefixUnmarshaler = (*Felt)(nil)
+	_ cbor.Marshaler             = (*Felt)(nil)
+	_ cbor.Unmarshaler           = (*Felt)(nil)
+)
+
 // Fast, felt-specialized CBOR marshaling.
 func (z *Felt) MarshalCBOR() ([]byte, error) {
 	data := make([]byte, maxCBORFeltLen)
@@ -24,6 +32,14 @@ func (z *Felt) UnmarshalCBOR(data []byte) error {
 	return cbor.Unmarshal(data, (*fp.Element)(z))
 }
 
+func (z *Felt) UnmarshalCBORPrefix(data []byte) (int, error) {
+	consumed, ok := decodeLimbs(data, z)
+	if !ok {
+		return 0, errors.New("felt: not limb-encoded")
+	}
+	return consumed, nil
+}
+
 const (
 	// These derive from the CBOR spec
 	// Limb types are always unsigned int

core/felt/cbor_fastpath_test.go

@@ -12,6 +12,7 @@ import (
 	"github.com/NethermindEth/juno/core/felt"
 	"github.com/consensys/gnark-crypto/ecc/stark-curve/fp"
 	"github.com/fxamacker/cbor/v2"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
@@ -237,3 +238,66 @@ func FuzzCBORFastPathMarshalEquivalence(f *testing.F) {
 		requireMarshalEquivalent(t, &value)
 	})
 }
+
+func TestFeltUnmarshalCBORPrefix(t *testing.T) {
+	want := felt.FromUint64[felt.Felt](0xdeadbeef)
+	encoded, err := want.MarshalCBOR()
+	require.NoError(t, err)
+
+	t.Run("reads the felt and reports the count", func(t *testing.T) {
+		var got felt.Felt
+		consumed, err := got.UnmarshalCBORPrefix(encoded)
+
+		require.NoError(t, err)
+		assert.Equal(t, len(encoded), consumed)
+		assert.Equal(t, want, got)
+	})
+
+	t.Run("trailing bytes are allowed, and not counted", func(t *testing.T) {
+		var got felt.Felt
+		consumed, err := got.UnmarshalCBORPrefix(append(append([]byte{}, encoded...), 0xff, 0xff))
+
+		require.NoError(t, err)
+		assert.Equal(t, len(encoded), consumed, "counting the trailer would misread the next field")
+		assert.Equal(t, want, got)
+	})
+
+	t.Run("rejects what is not a limb-encoded felt", func(t *testing.T) {
+		for name, data := range map[string][]byte{
+			"empty":                  {},
+			"truncated":              encoded[:len(encoded)-1],
+			"not an array":           {0x18, 0x2a},
+			"array of the wrong len": {0x83, 0x01, 0x02, 0x03},
+			"null":                   {0xf6},
+			"a limb that is signed":  {0x84, 0x20, 0x01, 0x02, 0x03},
+		} {
+			t.Run(name, func(t *testing.T) {
+				got := felt.FromUint64[felt.Felt](7)
+				consumed, err := got.UnmarshalCBORPrefix(data)
+
+				assert.Error(t, err)
+				assert.Zero(t, consumed)
+				assert.Equal(t, felt.FromUint64[felt.Felt](7), got,
+					"a rejected input must not touch the destination")
+			})
+		}
+	})
+}
+
+// TestTrailingBytesLeaveTheReceiverAlone covers what the fastpath table cannot: its
+// fixtures decode to the zero felt, so a receiver written by a rejected parse looks
+// exactly like one that was never touched.
+func TestTrailingBytesLeaveTheReceiverAlone(t *testing.T) {
+	var value felt.Felt
+	value.SetUint64(0xdeadbeef)
+
+	encoded, err := value.MarshalCBOR()
+	require.NoError(t, err)
+
+	withTrailing := append(append([]byte{}, encoded...), 0x00)
+
+	var got felt.Felt
+	err = got.UnmarshalCBOR(withTrailing)
+	require.Error(t, err, "trailing bytes must not decode")
+	require.True(t, got.IsZero(), "a failed decode must not write the receiver")
+}

core/felt/slice.go

@@ -2,13 +2,21 @@ package felt
 
 import (
 	"encoding/binary"
+	"errors"
 	"math"
 
+	"github.com/NethermindEth/juno/encoder/cborlite"
 	"github.com/fxamacker/cbor/v2"
 )
 
 type Slice[F FeltLike] []F
 
+var (
+	_ cborlite.PrefixUnmarshaler = (*Slice[Felt])(nil)
+	_ cbor.Marshaler             = (*Slice[Felt])(nil)
+	_ cbor.Unmarshaler           = (*Slice[Felt])(nil)
+)
+
 const (
 	// maxCBORArrayHeaderLen: 1 major/info byte + 4 length bytes (uint32)
 	maxCBORArrayHeaderLen = 1 + 4
@@ -33,36 +41,55 @@ func (s Slice[F]) MarshalCBOR() ([]byte, error) {
 }
 
 func (s *Slice[F]) UnmarshalCBOR(data []byte) error {
+	var out Slice[F]
+	consumed, ok := decodeSlicePrefix(data, &out)
+	if ok && consumed == len(data) {
+		*s = out
+		return nil
+	}
+	return s.unmarshalCBORGeneric(data)
+}
+
+func (s *Slice[F]) UnmarshalCBORPrefix(data []byte) (int, error) {
+	consumed, ok := decodeSlicePrefix(data, s)
+	if !ok {
+		return 0, errors.New("felt: not a limb-encoded array")
+	}
+	return consumed, nil
+}
+
+func decodeSlicePrefix[F FeltLike](data []byte, out *Slice[F]) (int, bool) {
+	if consumed, isNull := cborlite.ReadNull(data); isNull {
+		*out = nil
+		return consumed, true
+	}
+
 	size, offset, ok := decodeCBORArrayHeader(data)
 	if !ok {
-		return s.unmarshalGeneric(data)
+		return 0, false
 	}
 
 	// Checking if size was corrupted to avoid allocating a malicious amount of data
 	maxPossibleFelts := (len(data) - offset) / minCBORFeltLen
 	if size < 0 || size > maxPossibleFelts {
-		return s.unmarshalGeneric(data)
+		return 0, false
 	}
 
 	buffer := make([]F, size)
 	for i := range buffer {
 		consumed, ok := decodeLimbs(data[offset:], &buffer[i])
 		if !ok {
-			return s.unmarshalGeneric(data)
+			return 0, false
 		}
 		offset += consumed
 	}
 
-	if offset != len(data) {
-		return s.unmarshalGeneric(data)
-	}
-
-	*s = buffer
-	return nil
+	*out = buffer
+	return offset, true
 }
 
-// unmarshalGeneric handles any shape the fast path does not recognise.
-func (s *Slice[F]) unmarshalGeneric(data []byte) error {
+// unmarshalCBORGeneric handles any shape the fast path does not recognise.
+func (s *Slice[F]) unmarshalCBORGeneric(data []byte) error {
 	var buffer []F
 	if err := cbor.Unmarshal(data, &buffer); err != nil {
 		return err

core/felt/slice_test.go

@@ -10,6 +10,7 @@ import (
 
 	"github.com/NethermindEth/juno/core/felt"
 	"github.com/fxamacker/cbor/v2"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
@@ -370,3 +371,73 @@ func FuzzSliceDecodeJSONEquivalence(fz *testing.F) {
 		requireSliceDecodeJSONEquivalent(t, data)
 	})
 }
+
+func TestFeltSliceUnmarshalCBORPrefix(t *testing.T) {
+	want := felt.Slice[felt.Felt]{
+		felt.FromUint64[felt.Felt](1),
+		felt.FromUint64[felt.Felt](0xffffffffffffffff),
+	}
+	encoded, err := want.MarshalCBOR()
+	require.NoError(t, err)
+
+	t.Run("reads the slice and reports the count", func(t *testing.T) {
+		var got felt.Slice[felt.Felt]
+		consumed, err := got.UnmarshalCBORPrefix(encoded)
+
+		require.NoError(t, err)
+		assert.Equal(t, len(encoded), consumed)
+		assert.Equal(t, want, got)
+	})
+
+	t.Run("trailing bytes are allowed, and not counted", func(t *testing.T) {
+		var got felt.Slice[felt.Felt]
+		consumed, err := got.UnmarshalCBORPrefix(append(append([]byte{}, encoded...), 0xff))
+
+		require.NoError(t, err)
+		assert.Equal(t, len(encoded), consumed)
+	})
+
+	t.Run("an empty slice is empty, not nil", func(t *testing.T) {
+		empty, err := felt.Slice[felt.Felt]{}.MarshalCBOR()
+		require.NoError(t, err)
+
+		var got felt.Slice[felt.Felt]
+		consumed, err := got.UnmarshalCBORPrefix(empty)
+
+		require.NoError(t, err)
+		assert.Equal(t, len(empty), consumed)
+		assert.NotNil(t, got)
+		assert.Empty(t, got)
+	})
+
+	// A nil slice goes on the wire as null, so declining it would send every value
+	// holding one back to the generic decoder over an empty field.
+	t.Run("null reads as a nil slice", func(t *testing.T) {
+		got := felt.Slice[felt.Felt]{felt.FromUint64[felt.Felt](7)}
+		consumed, err := got.UnmarshalCBORPrefix([]byte{0xf6})
+
+		require.NoError(t, err)
+		assert.Equal(t, 1, consumed)
+		assert.Nil(t, got)
+	})
+
+	t.Run("rejects what is not a felt array", func(t *testing.T) {
+		for name, data := range map[string][]byte{
+			"empty":                         {},
+			"truncated":                     encoded[:len(encoded)-1],
+			"not an array":                  {0x18, 0x2a},
+			"an element that is not a felt": {0x82, 0x01, 0x02},
+			// A count far past the buffer would size an allocation if it were trusted.
+			"count past the buffer": {0x98, 0xff, 0x01},
+		} {
+			t.Run(name, func(t *testing.T) {
+				var got felt.Slice[felt.Felt]
+				consumed, err := got.UnmarshalCBORPrefix(data)
+
+				assert.Error(t, err)
+				assert.Zero(t, consumed)
+				assert.Nil(t, got)
+			})
+		}
+	})
+}

encoder/cborlite/cache.go

@@ -0,0 +1,67 @@
+package cborlite
+
+import (
+	"reflect"
+	"sync"
+)
+
+const maxNestedPlans = 32
+
+type reader func(target reflect.Value, data []byte) (consumed int, err error)
+
+// builtReader is what a build produced, including an error.
+type builtReader struct {
+	reader
+	err error
+}
+
+// cacheKey identifies a reader by type and strictness.
+type cacheKey struct {
+	valueType reflect.Type
+	strict    bool
+}
+
+var (
+	// Finished readers
+	readers sync.Map // cacheKey -> builtReader
+
+	// buildMutex serialises building.
+	buildMutex sync.Mutex
+	buildDepth int
+
+	// Finished plans.
+	plans = map[cacheKey]*plan{}
+)
+
+// cachedReader builds a reader and cache it so it only builds once.
+func cachedReader(valueType reflect.Type, strict bool) (reader, error) {
+	key := cacheKey{valueType: valueType, strict: strict}
+
+	cached, ok := readers.Load(key)
+	if ok {
+		built := cached.(builtReader)
+		return built.reader, built.err
+	}
+
+	buildMutex.Lock()
+	defer buildMutex.Unlock()
+
+	// Another goroutine may have finished while this one waited.
+	cached, ok = readers.Load(key)
+	if ok {
+		built := cached.(builtReader)
+		return built.reader, built.err
+	}
+
+	built, err := buildReader(valueType, strict)
+	readers.Store(key, builtReader{reader: built, err: err})
+	return built, err
+}
+
+func buildReader(valueType reflect.Type, strict bool) (reader, error) {
+	// If the type reads itself (e.g. implements UnmarshalCBORPrefix), it takes priority
+	if read, ok := specialTypeReader(valueType); ok {
+		return read, nil
+	}
+	return kindReader(valueType, strict)
+}