1/*
2 *
3 * Copyright 2024 gRPC authors.
4 *
5 * Licensed under the Apache License, Version 2.0 (the "License");
6 * you may not use this file except in compliance with the License.
7 * You may obtain a copy of the License at
8 *
9 * http://www.apache.org/licenses/LICENSE-2.0
10 *
11 * Unless required by applicable law or agreed to in writing, software
12 * distributed under the License is distributed on an "AS IS" BASIS,
13 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14 * See the License for the specific language governing permissions and
15 * limitations under the License.
16 *
17 */
18
19// Package proto defines the protobuf codec. Importing this package will
20// register the codec.
21package proto
22
23import (
24 "fmt"
25
26 "google.golang.org/grpc/encoding"
27 "google.golang.org/grpc/mem"
28 "google.golang.org/protobuf/proto"
29 "google.golang.org/protobuf/protoadapt"
30)
31
32// Name is the name registered for the proto compressor.
33const Name = "proto"
34
35func init() {
36 encoding.RegisterCodecV2(&codecV2{})
37}
38
39// codec is a CodecV2 implementation with protobuf. It is the default codec for
40// gRPC.
41type codecV2 struct{}
42
43func (c *codecV2) Marshal(v any) (data mem.BufferSlice, err error) {
44 vv := messageV2Of(v)
45 if vv == nil {
46 return nil, fmt.Errorf("proto: failed to marshal, message is %T, want proto.Message", v)
47 }
48
49 size := proto.Size(vv)
50 if mem.IsBelowBufferPoolingThreshold(size) {
51 buf, err := proto.Marshal(vv)
52 if err != nil {
53 return nil, err
54 }
55 data = append(data, mem.SliceBuffer(buf))
56 } else {
57 pool := mem.DefaultBufferPool()
58 buf := pool.Get(size)
59 if _, err := (proto.MarshalOptions{}).MarshalAppend((*buf)[:0], vv); err != nil {
60 pool.Put(buf)
61 return nil, err
62 }
63 data = append(data, mem.NewBuffer(buf, pool))
64 }
65
66 return data, nil
67}
68
69func (c *codecV2) Unmarshal(data mem.BufferSlice, v any) (err error) {
70 vv := messageV2Of(v)
71 if vv == nil {
72 return fmt.Errorf("failed to unmarshal, message is %T, want proto.Message", v)
73 }
74
75 buf := data.MaterializeToBuffer(mem.DefaultBufferPool())
76 defer buf.Free()
77 // TODO: Upgrade proto.Unmarshal to support mem.BufferSlice. Right now, it's not
78 // really possible without a major overhaul of the proto package, but the
79 // vtprotobuf library may be able to support this.
80 return proto.Unmarshal(buf.ReadOnlyData(), vv)
81}
82
83func messageV2Of(v any) proto.Message {
84 switch v := v.(type) {
85 case protoadapt.MessageV1:
86 return protoadapt.MessageV2Of(v)
87 case protoadapt.MessageV2:
88 return v
89 }
90
91 return nil
92}
93
94func (c *codecV2) Name() string {
95 return Name
96}