diff --git a/net/edf/register.go b/net/edf/register.go index 2006bbe3..8eeab258 100644 --- a/net/edf/register.go +++ b/net/edf/register.go @@ -60,7 +60,12 @@ func RegisterTypeOf(v any) error { fenc := func(value reflect.Value, b *lib.Buffer, _ *stateEncode) error { v := value.Interface().(Marshaler) - buf := b.Extend(4) + // Record the offset for the length prefix instead of using the + // slice returned by Extend. If MarshalEDF triggers buffer + // reallocation, the Extend slice becomes stale but the offset + // remains valid against the new backing array. + lenPrefixOffset := b.Len() + b.Extend(4) l := b.Len() if err := v.MarshalEDF(b); err != nil { return err @@ -70,7 +75,7 @@ func RegisterTypeOf(v any) error { if int64(lenBinary) > int64(math.MaxUint32-1) { return ErrBinaryTooLong } - binary.BigEndian.PutUint32(buf, uint32(lenBinary)) + binary.BigEndian.PutUint32(b.B[lenPrefixOffset:], uint32(lenBinary)) return nil } encoders.Store(tov, regEncoder(name, fenc)) diff --git a/net/edf/register_marshaler_test.go b/net/edf/register_marshaler_test.go new file mode 100644 index 00000000..d8334e11 --- /dev/null +++ b/net/edf/register_marshaler_test.go @@ -0,0 +1,77 @@ +package edf + +// Verifies that Encode/Decode roundtrip works correctly when MarshalEDF +// writes enough data to trigger lib.Buffer reallocation. A small initial +// buffer capacity (64 bytes) guarantees that MarshalEDF's Write call +// triggers append-reallocation for any non-trivial payload. + +import ( + "bytes" + "io" + "testing" + + "ergo.services/ergo/gen" + "ergo.services/ergo/lib" +) + +type largePayload struct { + Data []byte +} + +func (lp largePayload) MarshalEDF(w io.Writer) error { + _, err := w.Write(lp.Data) + return err +} + +func (lp *largePayload) UnmarshalEDF(data []byte) error { + lp.Data = make([]byte, len(data)) + copy(lp.Data, data) + return nil +} + +func TestMarshalerBufferReallocation(t *testing.T) { + if err := RegisterTypeOf(largePayload{}); err != nil && err != gen.ErrTaken { + t.Fatal(err) + } + + // Small payload that fits without reallocation, and a large payload + // that forces the buffer to grow. + sizes := []struct { + name string + n int + }{ + {"small", 100}, + {"large", lib.DefaultBufferLength * 2}, + } + + for _, tc := range sizes { + t.Run(tc.name, func(t *testing.T) { + data := make([]byte, tc.n) + for i := range data { + data[i] = byte(i % 251) + } + original := largePayload{Data: data} + + // Small capacity forces MarshalEDF to trigger buffer reallocation. + buf := &lib.Buffer{B: make([]byte, 0, 64)} + + if err := Encode(original, buf, Options{}); err != nil { + t.Fatalf("Encode: %v", err) + } + + decoded, _, err := Decode(buf.B, Options{}) + if err != nil { + t.Fatalf("Decode: %v (encoded %d bytes)", err, buf.Len()) + } + + got, ok := decoded.(largePayload) + if !ok { + t.Fatalf("expected largePayload, got %T", decoded) + } + if !bytes.Equal(got.Data, original.Data) { + t.Fatalf("roundtrip mismatch: got %d bytes, want %d bytes", + len(got.Data), len(original.Data)) + } + }) + } +}