github.com/authzed/spicedb@v1.32.1-0.20240520085336-ebda56537386/pkg/proto/dispatch/v1/01_codec.go (about)

     1  // This file registers a gRPC codec that replaces the default gRPC proto codec
     2  // with one that attempts to use protobuf codecs in the following order:
     3  // - vtprotobuf
     4  // - google.golang.org/encoding/proto
     5  
     6  package dispatchv1
     7  
     8  import (
     9  	"fmt"
    10  
    11  	"google.golang.org/grpc/encoding"
    12  	"google.golang.org/protobuf/proto"
    13  
    14  	// Guarantee that the built-in proto is called registered before this one
    15  	// so that it can be replaced.
    16  	_ "google.golang.org/grpc/encoding/proto"
    17  )
    18  
    19  // Name is the name registered for the proto compressor.
    20  const Name = "proto"
    21  
    22  type vtprotoMessage interface {
    23  	MarshalVT() ([]byte, error)
    24  	UnmarshalVT([]byte) error
    25  }
    26  
    27  type vtprotoCodec struct{}
    28  
    29  func (vtprotoCodec) Name() string { return Name }
    30  
    31  func (vtprotoCodec) Marshal(v any) ([]byte, error) {
    32  	if m, ok := v.(vtprotoMessage); ok {
    33  		return m.MarshalVT()
    34  	}
    35  
    36  	if m, ok := v.(proto.Message); ok {
    37  		return proto.Marshal(m)
    38  	}
    39  
    40  	return nil, fmt.Errorf("failed to marshal, message is %T, want proto.Message", v)
    41  }
    42  
    43  func (vtprotoCodec) Unmarshal(data []byte, v any) error {
    44  	if m, ok := v.(vtprotoMessage); ok {
    45  		return m.UnmarshalVT(data)
    46  	}
    47  
    48  	if m, ok := v.(proto.Message); ok {
    49  		return proto.Unmarshal(data, m)
    50  	}
    51  
    52  	return fmt.Errorf("failed to unmarshal, message is %T, want proto.Message", v)
    53  }
    54  
    55  func init() {
    56  	encoding.RegisterCodec(vtprotoCodec{})
    57  }