Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
62 changes: 46 additions & 16 deletions codec/grpc/grpc_codec.go
Original file line number Diff line number Diff line change
@@ -1,33 +1,63 @@
package grpc

import "fmt"
import (
"google.golang.org/grpc/encoding"
"google.golang.org/grpc/mem"

// Guarantee that the built-in proto is called registered before this one
// so that it can be replaced.
_ "google.golang.org/grpc/encoding/proto"
)

// Name is the name registered for the proto compressor.
const Name = "proto"

type Codec struct{}

type vtprotoMessage interface {
MarshalVT() ([]byte, error)
MarshalToSizedBufferVT(data []byte) (int, error)
UnmarshalVT([]byte) error
SizeVT() int
}

type Codec struct {
fallback encoding.CodecV2
}

func (Codec) Marshal(v interface{}) ([]byte, error) {
vt, ok := v.(vtprotoMessage)
if !ok {
return nil, fmt.Errorf("failed to marshal, message is %T (missing vtprotobuf helpers)", v)
func (Codec) Name() string { return Name }

func (c *Codec) Marshal(v any) (mem.BufferSlice, error) {
if m, ok := v.(vtprotoMessage); ok {
size := m.SizeVT()
if mem.IsBelowBufferPoolingThreshold(size) {
buf := make([]byte, 0, size)

Choose a reason for hiding this comment

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

There is a bug on this line IMO and it should go like either make([]byte, size) or mem.BufferSlice{mem.SliceBuffer(buf[:size])} on the line 35. Because the original slice reference will have its len still set to 0 effectively returning 0 length byte array from the method.

Copy link
Author

Choose a reason for hiding this comment

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

Yes! Apologies -- I have this fixed in my PR for SpiceDB but didn't carry over the changes here.

if _, err := m.MarshalToSizedBufferVT(buf[:size]); err != nil {
return nil, err
}
return mem.BufferSlice{mem.SliceBuffer(buf)}, nil
}
pool := mem.DefaultBufferPool()
buf := pool.Get(size)
if _, err := m.MarshalToSizedBufferVT((*buf)[:size]); err != nil {
pool.Put(buf)
return nil, err
}
return mem.BufferSlice{mem.NewBuffer(buf, pool)}, nil
}
return vt.MarshalVT()

return c.fallback.Marshal(v)
}

func (Codec) Unmarshal(data []byte, v interface{}) error {
vt, ok := v.(vtprotoMessage)
if !ok {
return fmt.Errorf("failed to unmarshal, message is %T (missing vtprotobuf helpers)", v)
func (c *Codec) Unmarshal(data mem.BufferSlice, v any) error {
if m, ok := v.(vtprotoMessage); ok {
buf := data.MaterializeToBuffer(mem.DefaultBufferPool())
defer buf.Free()
return m.UnmarshalVT(buf.ReadOnlyData())
}
return vt.UnmarshalVT(data)

return c.fallback.Unmarshal(data, v)
}

func (Codec) Name() string {
return Name
func init() {
encoding.RegisterCodecV2(&Codec{
fallback: encoding.GetCodecV2("proto"),
Copy link

Choose a reason for hiding this comment

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

We can use name from proto so we will avoid unused import https://pkg.go.dev/google.golang.org/grpc/encoding/proto#pkg-constants

Suggested change
fallback: encoding.GetCodecV2("proto"),
fallback: encoding.GetCodecV2(proto.Name),

})
}
16 changes: 9 additions & 7 deletions go.mod
Original file line number Diff line number Diff line change
@@ -1,20 +1,22 @@
module github.com/planetscale/vtprotobuf

go 1.20
go 1.21

toolchain go1.23.0

require (
github.com/stretchr/testify v1.8.4
google.golang.org/grpc v1.58.2
google.golang.org/protobuf v1.31.0
google.golang.org/grpc v1.67.0-dev.0.20240821211139-9ab8b62505c8
google.golang.org/protobuf v1.34.2
)

require (
github.com/davecgh/go-spew v1.1.1 // indirect
github.com/golang/protobuf v1.5.3 // indirect
github.com/pmezard/go-difflib v1.0.0 // indirect
golang.org/x/net v0.14.0 // indirect
golang.org/x/sys v0.11.0 // indirect
golang.org/x/text v0.12.0 // indirect
google.golang.org/genproto/googleapis/rpc v0.0.0-20230822172742-b8732ec3820d // indirect
golang.org/x/net v0.28.0 // indirect
golang.org/x/sys v0.24.0 // indirect
golang.org/x/text v0.17.0 // indirect
google.golang.org/genproto/googleapis/rpc v0.0.0-20240814211410-ddb44dafa142 // indirect
gopkg.in/yaml.v3 v3.0.1 // indirect
)
12 changes: 12 additions & 0 deletions go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -11,19 +11,31 @@ github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcU
github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo=
golang.org/x/net v0.14.0 h1:BONx9s002vGdD9umnlX1Po8vOZmrgH34qlHcD1MfK14=
golang.org/x/net v0.14.0/go.mod h1:PpSgVXXLK0OxS0F31C1/tv6XNguvCrnXIDrFMspZIUI=
golang.org/x/net v0.28.0 h1:a9JDOJc5GMUJ0+UDqmLT86WiEy7iWyIhz8gz8E4e5hE=
golang.org/x/net v0.28.0/go.mod h1:yqtgsTWOOnlGLG9GFRrK3++bGOUEkNBoHZc8MEDWPNg=
golang.org/x/sys v0.11.0 h1:eG7RXZHdqOJ1i+0lgLgCpSXAp6M3LYlAo6osgSi0xOM=
golang.org/x/sys v0.11.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.24.0 h1:Twjiwq9dn6R1fQcyiK+wQyHWfaz/BJB+YIpzU/Cv3Xg=
golang.org/x/sys v0.24.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
golang.org/x/text v0.12.0 h1:k+n5B8goJNdU7hSvEtMUz3d1Q6D/XW4COJSJR6fN0mc=
golang.org/x/text v0.12.0/go.mod h1:TvPlkZtksWOMsz7fbANvkp4WM8x/WCo/om8BMLbz+aE=
golang.org/x/text v0.17.0 h1:XtiM5bkSOt+ewxlOE/aE/AKEHibwj/6gvWMl9Rsh0Qc=
golang.org/x/text v0.17.0/go.mod h1:BuEKDfySbSR4drPmRPG/7iBdf8hvFMuRexcpahXilzY=
golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
google.golang.org/genproto/googleapis/rpc v0.0.0-20230822172742-b8732ec3820d h1:uvYuEyMHKNt+lT4K3bN6fGswmK8qSvcreM3BwjDh+y4=
google.golang.org/genproto/googleapis/rpc v0.0.0-20230822172742-b8732ec3820d/go.mod h1:+Bk1OCOj40wS2hwAMA+aCW9ypzm63QTBBHp6lQ3p+9M=
google.golang.org/genproto/googleapis/rpc v0.0.0-20240814211410-ddb44dafa142 h1:e7S5W7MGGLaSu8j3YjdezkZ+m1/Nm0uRVRMEMGk26Xs=
google.golang.org/genproto/googleapis/rpc v0.0.0-20240814211410-ddb44dafa142/go.mod h1:UqMtugtsSgubUsoxbuAoiCXvqvErP7Gf0so0mK9tHxU=
google.golang.org/grpc v1.58.2 h1:SXUpjxeVF3FKrTYQI4f4KvbGD5u2xccdYdurwowix5I=
google.golang.org/grpc v1.58.2/go.mod h1:tgX3ZQDlNJGU96V6yHh1T/JeoBQ2TXdr43YbYSsCJk0=
google.golang.org/grpc v1.67.0-dev.0.20240821211139-9ab8b62505c8 h1:b61arqKMyICRrCxna/qIyjYGWlfem3U2iLpF7LMnDaU=
google.golang.org/grpc v1.67.0-dev.0.20240821211139-9ab8b62505c8/go.mod h1:1gLDyUQU7CTLJI90u3nXZ9ekeghjeM7pTDZlqFNg2AA=
google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp09yW+WbY/TyQbw=
google.golang.org/protobuf v1.26.0/go.mod h1:9q0QmTI4eRPtz6boOQmLYwt+qCgq0jsYwAQnmE0givc=
google.golang.org/protobuf v1.31.0 h1:g0LDEJHgrBl9N9r17Ru3sqWhkIx2NB67okBHPwC7hs8=
google.golang.org/protobuf v1.31.0/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqwMG9pJV4I=
google.golang.org/protobuf v1.34.2 h1:6xV6lTsCfpGD21XK49h7MhtcApnLqkfYgPcdHftf6hg=
google.golang.org/protobuf v1.34.2/go.mod h1:qYOHts0dSfpeUzUFpOMr/WGzszTmLH+DiWniOlNbLDw=
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
Expand Down