Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

UnmarshalFrom support in codec.Codec #3335

Merged
merged 11 commits into from
Sep 4, 2024
1 change: 1 addition & 0 deletions codec/codec.go
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@ var (
type Codec interface {
MarshalInto(interface{}, *wrappers.Packer) error
Unmarshal([]byte, interface{}) error
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Would it make sense to remove Unmarshal entirely and have the codec manager do the exhaustion check?

This interface feels weird. (Not really having to do with your PR... but the existing code feels weird)

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If we'd prefer to leave that as a potential followup - I think that's reasonable too.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think that should be a follow up instead of being part of this PR

UnmarshalFrom(*wrappers.Packer, interface{}) error

// Returns the size, in bytes, of [value] when it's marshaled
Size(value interface{}) (int, error)
Expand Down
26 changes: 26 additions & 0 deletions codec/codectest/codectest.go
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,8 @@ import (

"github.com/stretchr/testify/require"

"github.com/ava-labs/avalanchego/utils/wrappers"

codecpkg "github.com/ava-labs/avalanchego/codec"
)

Expand Down Expand Up @@ -75,6 +77,7 @@ var (
{"Slice Length Overflow", TestSliceLengthOverflow},
{"Map", TestMap},
{"Can Marshal Large Slices", TestCanMarshalLargeSlices},
{"Implements UnmarshalFrom", TestImplementsUnmarshalFrom},
}

MultipleTagsTests = []NamedTest{
Expand Down Expand Up @@ -1153,3 +1156,26 @@ func FuzzStructUnmarshal(codec codecpkg.GeneralCodec, f *testing.F) {
require.Len(bytes, size)
})
}

func TestImplementsUnmarshalFrom(t testing.TB, codec codecpkg.GeneralCodec) {
require := require.New(t)

p := wrappers.Packer{MaxSize: 1024}
p.PackFixedBytes([]byte{0, 1, 2}) // pack 3 extra bytes prefix

mySlice := []bool{true, false, true, true}

require.NoError(codec.MarshalInto(mySlice, &p))

p.PackFixedBytes([]byte{7, 7, 7}) // pack 3 extra bytes suffix

bytesLen, err := codec.Size(mySlice)
require.NoError(err)
require.Equal(3+bytesLen+3, p.Offset)

p = wrappers.Packer{Bytes: p.Bytes, MaxSize: p.MaxSize, Offset: 3}

var sliceUnmarshaled []bool
require.NoError(codec.UnmarshalFrom(&p, &sliceUnmarshaled))
require.Equal(mySlice, sliceUnmarshaled)
aaronbuchwald marked this conversation as resolved.
Show resolved Hide resolved
}
24 changes: 15 additions & 9 deletions codec/reflectcodec/type_codec.go
Original file line number Diff line number Diff line change
Expand Up @@ -499,18 +499,10 @@ func (c *genericCodec) marshal(
// Unmarshal unmarshals [bytes] into [dest], where [dest] must be a pointer or
// interface
func (c *genericCodec) Unmarshal(bytes []byte, dest interface{}) error {
if dest == nil {
return codec.ErrUnmarshalNil
}

p := wrappers.Packer{
Bytes: bytes,
}
destPtr := reflect.ValueOf(dest)
if destPtr.Kind() != reflect.Ptr {
return errNeedPointer
}
if err := c.unmarshal(&p, destPtr.Elem(), nil /*=typeStack*/); err != nil {
if err := c.UnmarshalFrom(&p, dest); err != nil {
return err
}
if p.Offset != len(bytes) {
Expand All @@ -523,6 +515,20 @@ func (c *genericCodec) Unmarshal(bytes []byte, dest interface{}) error {
return nil
}

// UnmarshalFrom unmarshals [p.Bytes] into [dest], where [dest] must be a pointer or
// interface
func (c *genericCodec) UnmarshalFrom(p *wrappers.Packer, dest interface{}) error {
containerman17 marked this conversation as resolved.
Show resolved Hide resolved
if dest == nil {
return codec.ErrUnmarshalNil
}

destPtr := reflect.ValueOf(dest)
if destPtr.Kind() != reflect.Ptr {
return errNeedPointer
}
return c.unmarshal(p, destPtr.Elem(), nil /*=typeStack*/)
}

// Unmarshal from p.Bytes into [value]. [value] must be addressable.
//
// c.lock should be held for the duration of this function
Expand Down
Loading