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.
21 package proto
22 23 import (
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.
33 const Name = "proto"
34 35 func init() {
36 encoding.RegisterCodecV2(&codecV2{})
37 }
38 39 // codec is a CodecV2 implementation with protobuf. It is the default codec for
40 // gRPC.
41 type codecV2 struct{}
42 43 func (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 // Important: if we remove this Size call then we cannot use
50 // UseCachedSize in MarshalOptions below.
51 size := proto.Size(vv)
52 53 // MarshalOptions with UseCachedSize allows reusing the result from the
54 // previous Size call. This is safe here because:
55 //
56 // 1. We just computed the size.
57 // 2. We assume the message is not being mutated concurrently.
58 //
59 // Important: If the proto.Size call above is removed, using UseCachedSize
60 // becomes unsafe and may lead to incorrect marshaling.
61 //
62 // For more details, see the doc of UseCachedSize:
63 // https://pkg.go.dev/google.golang.org/protobuf/proto#MarshalOptions
64 marshalOptions := proto.MarshalOptions{UseCachedSize: true}
65 66 if mem.IsBelowBufferPoolingThreshold(size) {
67 buf, err := marshalOptions.Marshal(vv)
68 if err != nil {
69 return nil, err
70 }
71 data = append(data, mem.SliceBuffer(buf))
72 } else {
73 pool := mem.DefaultBufferPool()
74 buf := pool.Get(size)
75 if _, err := marshalOptions.MarshalAppend((*buf)[:0], vv); err != nil {
76 pool.Put(buf)
77 return nil, err
78 }
79 data = append(data, mem.NewBuffer(buf, pool))
80 }
81 82 return data, nil
83 }
84 85 func (c *codecV2) Unmarshal(data mem.BufferSlice, v any) (err error) {
86 vv := messageV2Of(v)
87 if vv == nil {
88 return fmt.Errorf("failed to unmarshal, message is %T, want proto.Message", v)
89 }
90 91 buf := data.MaterializeToBuffer(mem.DefaultBufferPool())
92 defer buf.Free()
93 // TODO: Upgrade proto.Unmarshal to support mem.BufferSlice. Right now, it's not
94 // really possible without a major overhaul of the proto package, but the
95 // vtprotobuf library may be able to support this.
96 return proto.Unmarshal(buf.ReadOnlyData(), vv)
97 }
98 99 func messageV2Of(v any) proto.Message {
100 switch v := v.(type) {
101 case protoadapt.MessageV1:
102 return protoadapt.MessageV2Of(v)
103 case protoadapt.MessageV2:
104 return v
105 }
106 107 return nil
108 }
109 110 func (c *codecV2) Name() string {
111 return Name
112 }
113