Serde Header

pkg/sr/serde.go

@@ -1,19 +1,14 @@
 package sr
 
 import (
+	"bytes"
 	"encoding/binary"
 	"errors"
-	"io"
 	"reflect"
 	"sync"
 	"sync/atomic"
 )
 
-// The wire format for encoded types is 0, then big endian uint32 of the ID,
-// then the encoded message.
-//
-// https://docs.confluent.io/platform/current/schema-registry/serdes-develop/index.html#wire-format
-
 var (
 	// ErrNotRegistered is returned from Serde when attempting to encode a
 	// value or decode an ID that has not been registered, or when using
@@ -80,8 +75,9 @@ type tserde struct {
 	gen          func() any
 	typeof       reflect.Type
 
-	index    []int          // for encoding, an optional index we use
-	subindex map[int]tserde // for decoding, we look up sub-indices in the payload
+	index         []int          // for encoding, an optional index we use
+	subindex      map[int]tserde // for decoding, we look up sub-indices in the payload
+	subindexDepth int            // for decoding, we need to know the maximum depth of the subindex map
 }
 
 // Serde encodes and decodes values according to the schema registry wire
@@ -101,6 +97,7 @@ type Serde struct {
 	mu    sync.Mutex
 
 	defaults []SerdeOpt
+	header   SerdeHeader
 }
 
 var (
@@ -134,6 +131,20 @@ func (s *Serde) SetDefaults(opts ...SerdeOpt) {
 	s.defaults = opts
 }
 
+// SetHeader configures which header should be used when encoding and decoding
+// values. If the header is set to nil it falls back to confluentHeader.
+func (s *Serde) SetHeader(header SerdeHeader) {
+	s.header = header
+}
+
+// Header returns the configured header.
+func (s *Serde) Header() SerdeHeader {
+	if s.header == nil {
+		return defaultSerdeHeader
+	}
+	return s.header
+}
+
 // Register registers a schema ID and the value it corresponds to, as well as
 // the encoding or decoding functions. You need to register functions depending
 // on whether you are only encoding, only decoding, or both.
@@ -172,7 +183,16 @@ func (s *Serde) Register(id int, v any, opts ...SerdeOpt) {
 	// tree to find the end node we are initializing.
 	k := id
 	at := m[k]
-	for _, idx := range t.index {
+	max := func(i, j int) int {
+		if i > j {
+			return i
+		}
+		return j
+	}
+	for i, idx := range t.index {
+		at.subindexDepth = max(at.subindexDepth, len(t.index)-i)
+		m[k] = at
+
 		m = at.subindex
 		k = idx
 		at = m[k]
@@ -186,15 +206,16 @@ func (s *Serde) Register(id int, v any, opts ...SerdeOpt) {
 
 	// Now, we initialize the end node.
 	t = tserde{
-		id:           uint32(id),
-		exists:       true,
-		encode:       t.encode,
-		appendEncode: t.appendEncode,
-		decode:       t.decode,
-		gen:          t.gen,
-		typeof:       typeof,
-		index:        t.index,
-		subindex:     at.subindex,
+		id:            uint32(id),
+		exists:        true,
+		encode:        t.encode,
+		appendEncode:  t.appendEncode,
+		decode:        t.decode,
+		gen:           t.gen,
+		typeof:        typeof,
+		index:         t.index,
+		subindex:      at.subindex,
+		subindexDepth: at.subindexDepth,
 	}
 	m[k] = t
 }
@@ -219,8 +240,8 @@ func tserdeMapClone(m map[int]tserde, at int, index []int) map[int]tserde {
 	return dup
 }
 
-// Encode encodes a value according to the schema registry wire format and
-// returns it. If EncodeFn was not used, this returns ErrNotRegistered.
+// Encode encodes a value and prepends the header according to the configured
+// SerdeHeader. If EncodeFn was not used, this returns ErrNotRegistered.
 func (s *Serde) Encode(v any) ([]byte, error) {
 	return s.AppendEncode(nil, v)
 }
@@ -234,23 +255,9 @@ func (s *Serde) AppendEncode(b []byte, v any) ([]byte, error) {
 		return b, ErrNotRegistered
 	}
 
-	b = append(b,
-		0,
-		byte(t.id>>24),
-		byte(t.id>>16),
-		byte(t.id>>8),
-		byte(t.id>>0),
-	)
-
-	if len(t.index) > 0 {
-		if len(t.index) == 1 && t.index[0] == 0 {
-			b = append(b, 0) // first-index shortcut (one type in the protobuf)
-		} else {
-			b = binary.AppendVarint(b, int64(len(t.index)))
-			for _, idx := range t.index {
-				b = binary.AppendVarint(b, int64(idx))
-			}
-		}
+	b, err := s.Header().AppendEncode(b, int(t.id), t.index)
+	if err != nil {
+		return nil, err
 	}
 
 	if t.appendEncode != nil {
@@ -316,44 +323,119 @@ func (s *Serde) DecodeNew(b []byte) (any, error) {
 }
 
 func (s *Serde) decodeFind(b []byte) ([]byte, tserde, error) {
-	if len(b) < 5 || b[0] != 0 {
-		return nil, tserde{}, ErrBadHeader
+	h := s.Header()
+
+	id, b, err := h.DecodeID(b)
+	if err != nil {
+		return nil, tserde{}, err
 	}
-	id := binary.BigEndian.Uint32(b[1:5])
-	b = b[5:]
 
-	t := s.loadIDs()[int(id)]
+	t := s.loadIDs()[id]
 	if len(t.subindex) > 0 {
-		r := bReader{b}
-		br := io.ByteReader(&r)
-		l, err := binary.ReadVarint(br)
-		if l == 0 { // length 0 is a shortcut for length 1, index 0
-			t = t.subindex[0]
-		}
-		for err == nil && t.subindex != nil && l > 0 {
-			var idx int64
-			idx, err = binary.ReadVarint(br)
-			t = t.subindex[int(idx)]
-			l--
-		}
+		var index []int
+		index, b, err = h.DecodeIndex(b, t.subindexDepth)
 		if err != nil {
-			return nil, t, err
+			return nil, tserde{}, err
+		}
+		for _, idx := range index {
+			if t.subindex == nil {
+				return nil, tserde{}, ErrNotRegistered
+			}
+			t = t.subindex[idx]
 		}
-		b = r.b
 	}
 	if !t.exists {
-		return nil, t, ErrNotRegistered
+		return nil, tserde{}, ErrNotRegistered
 	}
 	return b, t, nil
 }
 
-type bReader struct{ b []byte }
+// SerdeHeader encodes and decodes a message header.
+type SerdeHeader interface {
+	// AppendEncode encodes a schema ID and optional index to b, returning the
+	// updated slice or an error.
+	AppendEncode(b []byte, id int, index []int) ([]byte, error)
+	// DecodeID decodes an ID from in, returning the ID and the remaining bytes,
+	// or an error.
+	DecodeID(in []byte) (id int, out []byte, err error)
+	// DecodeIndex decodes at most maxLength of a schema index from in,
+	// returning the index and remaining bytes, or an error.
+	DecodeIndex(in []byte, maxLength int) (index []int, out []byte, err error)
+}
+
+var defaultSerdeHeader = new(confluentHeader)
+
+// confluentHeader is a SerdeHeader that produces the Confluent wire format. It
+// starts with 0, then big endian uint32 of the ID, then index (only protobuf),
+// then the encoded message.
+//
+// https://docs.confluent.io/platform/current/schema-registry/serdes-develop/index.html#wire-format
+type confluentHeader struct{}
+
+// AppendEncode appends an encoded header to b according to the Confluent wire
+// format and returns it. Error is always nil.
+func (*confluentHeader) AppendEncode(b []byte, id int, index []int) ([]byte, error) {
+	b = append(
+		b,
+		0,
+		byte(id>>24),
+		byte(id>>16),
+		byte(id>>8),
+		byte(id>>0),
+	)
+
+	if len(index) > 0 {
+		if len(index) == 1 && index[0] == 0 {
+			b = append(b, 0) // first-index shortcut (one type in the protobuf)
+		} else {
+			b = binary.AppendVarint(b, int64(len(index)))
+			for _, idx := range index {
+				b = binary.AppendVarint(b, int64(idx))
+			}
+		}
+	}
 
-func (b *bReader) ReadByte() (byte, error) {
-	if len(b.b) > 0 {
-		r := b.b[0]
-		b.b = b.b[1:]
-		return r, nil
+	return b, nil
+}
+
+// DecodeID strips and decodes the schema ID from b. It returns the ID alongside
+// the unread bytes. If the header does not contain the magic byte or b contains
+// less than 5 bytes it returns ErrBadHeader.
+func (*confluentHeader) DecodeID(b []byte) (int, []byte, error) {
+	if len(b) < 5 || b[0] != 0 {
+		return 0, nil, ErrBadHeader
+	}
+	id := binary.BigEndian.Uint32(b[1:5])
+	return int(id), b[5:], nil
+}
+
+// DecodeIndex strips and decodes indices from b. It returns the index slice
+// alongside the unread bytes. It expects b to be the output of DecodeID (schema
+// ID should already be stripped away). If maxLength is greater than 0 and the
+// encoded data contains more indices than maxLength the function returns
+// ErrNotRegistered.
+func (*confluentHeader) DecodeIndex(b []byte, maxLength int) ([]int, []byte, error) {
+	buf := bytes.NewBuffer(b)
+	l, err := binary.ReadVarint(buf)
+	if err != nil {
+		return nil, nil, err
+	}
+	if l == 0 { // length 0 is a shortcut for length 1, index 0
+		return []int{0}, buf.Bytes(), nil
+	}
+	if l < 0 { // index length can't be negative
+		return nil, nil, ErrBadHeader
+	}
+	if maxLength > 0 && int(l) > maxLength { // index count is greater than expected
+		return nil, nil, ErrNotRegistered
+	}
+	index := make([]int, l)
+	for i := range index {
+		idx, err := binary.ReadVarint(buf)
+		if err != nil {
+			return nil, nil, err
+		}
+		index[i] = int(idx)
 	}
-	return 0, io.EOF
+	return index, buf.Bytes(), nil
 }

pkg/sr/serde_test.go

@@ -135,7 +135,86 @@ func TestSerde(t *testing.T) {
 	if _, err := serde.DecodeNew([]byte{0, 0, 0, 0, 3}); err != io.EOF {
 		t.Errorf("got %v != exp io.EOF", err)
 	}
+	if _, err := serde.DecodeNew([]byte{0, 0, 0, 0, 3, 8}); err != ErrNotRegistered {
+		t.Errorf("got %v != exp ErrNotRegistered", err)
+	}
 	if _, err := serde.DecodeNew([]byte{0, 0, 0, 0, 99}); err != ErrNotRegistered {
 		t.Errorf("got %v != exp ErrNotRegistered", err)
 	}
 }
+
+func TestConfluentHeader(t *testing.T) {
+	var h confluentHeader
+
+	for i, test := range []struct {
+		id     int
+		index  []int
+		expEnc []byte
+	}{
+		{id: 1, index: nil, expEnc: []byte{0, 0, 0, 0, 1}},
+		{id: 256, index: nil, expEnc: []byte{0, 0, 0, 1, 0}},
+		{id: 2, index: []int{0}, expEnc: []byte{0, 0, 0, 0, 2, 0}},
+		{id: 3, index: []int{1}, expEnc: []byte{0, 0, 0, 0, 3, 2, 2}},
+		{id: 4, index: []int{1, 2, 3}, expEnc: []byte{0, 0, 0, 0, 4, 6, 2, 4, 6}},
+	} {
+		b, err := h.AppendEncode(nil, test.id, test.index)
+		if err != nil {
+			t.Errorf("#%d AppendEncode: got unexpected err %v", i, err)
+			continue
+		}
+		if !bytes.Equal(b, test.expEnc) {
+			t.Errorf("#%d: AppendEncode(%v) != exp(%v)", i, b, test.expEnc)
+			continue
+		}
+
+		if b2, _ := h.AppendEncode([]byte("foo"), test.id, test.index); !bytes.Equal(b2, append([]byte("foo"), b...)) {
+			t.Errorf("#%d got AppendEncode(%v) != AppendEncode(foo%v)", i, b2, b)
+		}
+
+		id, b2, err := h.DecodeID(b)
+		if err != nil {
+			t.Errorf("#%d DecodeID: got unexpected err %v", i, err)
+			continue
+		}
+		if id != test.id {
+			t.Errorf("#%d: DecodeID: id(%v) != exp(%v)", i, id, test.id)
+			continue
+		}
+		if test.index == nil && len(b2) != 0 {
+			t.Errorf("#%d: DecodeID: bytes(%v) != exp([])", i, b2)
+			continue
+		}
+
+		if test.index != nil {
+			index, b3, err := h.DecodeIndex(b2, len(test.index))
+			if err != nil {
+				t.Errorf("#%d DecodeIndex: got unexpected err %v", i, err)
+				continue
+			}
+			if !reflect.DeepEqual(index, test.index) {
+				t.Errorf("#%d: DecodeIndex: index(%v) != exp(%v)", i, index, test.index)
+				continue
+			}
+			if len(b3) != 0 {
+				t.Errorf("#%d: DecodeIndex: bytes(%v) != exp([])", i, b3)
+				continue
+			}
+		}
+	}
+
+	if _, _, err := h.DecodeID([]byte{1, 0, 0, 0, 0, 1}); err != ErrBadHeader {
+		t.Errorf("got %v != exp ErrBadHeader", err)
+	}
+	if _, _, err := h.DecodeID([]byte{0, 0, 0, 0}); err != ErrBadHeader {
+		t.Errorf("got %v != exp ErrBadHeader", err)
+	}
+	if _, _, err := h.DecodeIndex([]byte{2}, 1); err != io.EOF {
+		t.Errorf("got %v != exp io.EOF", err)
+	}
+	if _, _, err := h.DecodeIndex([]byte{6, 2, 4, 6}, 2); err != ErrNotRegistered {
+		t.Errorf("got %v != exp ErrNotRegistered", err)
+	}
+	if _, _, err := h.DecodeIndex([]byte{1}, 2); err != ErrBadHeader {
+		t.Errorf("got %v != exp ErrBadHeader", err)
+	}
+}