github.com/abemedia/go-don@v0.2.2-0.20240329015135-be88e32bb73b/encoding/protobuf/protobuf.go (about)

     1  // Package protobuf provides encoding and decoding of Protocol Buffers data.
     2  package protobuf
     3  
     4  import (
     5  	"reflect"
     6  	"sync"
     7  
     8  	"github.com/abemedia/go-don"
     9  	"github.com/abemedia/go-don/encoding"
    10  	"google.golang.org/protobuf/proto"
    11  )
    12  
    13  var (
    14  	messageType = reflect.TypeOf((*proto.Message)(nil)).Elem()
    15  	cache       sync.Map
    16  )
    17  
    18  func unmarshal(b []byte, v any) error {
    19  	typ := reflect.TypeOf(v)
    20  	fn, ok := cache.Load(typ)
    21  	if !ok {
    22  		if typ.Elem().Implements(messageType) {
    23  			fn = func(v any) any { return reflect.ValueOf(v).Elem().Interface() }
    24  		} else {
    25  			fn = func(v any) any { return v }
    26  		}
    27  		cache.Store(typ, fn)
    28  	}
    29  
    30  	m, ok := fn.(func(v any) any)(v).(proto.Message)
    31  	if !ok {
    32  		return don.ErrUnsupportedMediaType
    33  	}
    34  	return proto.Unmarshal(b, m)
    35  }
    36  
    37  func marshal(v any) ([]byte, error) {
    38  	m, ok := v.(proto.Message)
    39  	if !ok {
    40  		return nil, don.ErrNotAcceptable
    41  	}
    42  	return proto.Marshal(m)
    43  }
    44  
    45  func init() {
    46  	mediaType := "application/protobuf"
    47  	alias := "application/x-protobuf"
    48  
    49  	encoding.RegisterDecoder(unmarshal, mediaType, alias)
    50  	encoding.RegisterEncoder(marshal, mediaType, alias)
    51  }