-
Notifications
You must be signed in to change notification settings - Fork 1.8k
/
codec.go
120 lines (100 loc) · 3.19 KB
/
codec.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
// Based on https://github.com/trusch/grpc-proxy
// Copyright Michal Witkowski. Licensed under Apache2 license: https://github.com/trusch/grpc-proxy/blob/master/LICENSE.txt
package codec
import (
"fmt"
protoV1 "github.com/golang/protobuf/proto" //nolint:staticcheck
"google.golang.org/grpc/encoding"
"google.golang.org/protobuf/proto"
)
// Name is the name by which the proxy codec is registered in the encoding codec registry
// We have to say that we are the "proto" codec otherwise marshaling will fail.
const Name = "proto"
func init() {
Register()
}
// Register manually registers the codec.
func Register() {
encoding.RegisterCodec(codec())
}
// codec returns a proxying grpc.codec with the default protobuf codec as parent.
//
// See CodecWithParent.
func codec() encoding.Codec {
// since we have registered the default codec by importing it,
// we can fetch it from the registry and use it as our parent
// and overwrite the existing codec in the registry
return codecWithParent(&protoCodec{})
}
// CodecWithParent returns a proxying grpc.Codec with a user provided codec as parent.
//
// This codec is *crucial* to the functioning of the proxy. It allows the proxy server to be oblivious
// to the schema of the forwarded messages. It basically treats a gRPC message frame as raw bytes.
// However, if the server handler, or the client caller are not proxy-internal functions it will fall back
// to trying to decode the message using a fallback codec.
func codecWithParent(fallback encoding.Codec) encoding.Codec {
return &Proxy{parentCodec: fallback}
}
// Proxy satisfies the encoding.Codec interface.
type Proxy struct {
parentCodec encoding.Codec
}
// Frame holds the proxy transported data.
type Frame struct {
payload []byte
}
// ProtoMessage tags a frame as valid proto message.
func (f *Frame) ProtoMessage() {
// nop
}
// Raw returns the raw message.
// This is primarily useful for debugging.
func (f Frame) Raw() []byte {
return f.payload
}
// Marshal implements the encoding.Codec interface method.
func (p *Proxy) Marshal(v any) ([]byte, error) {
out, ok := v.(*Frame)
if !ok {
return p.parentCodec.Marshal(v)
}
return out.payload, nil
}
// Unmarshal implements the encoding.Codec interface method.
func (p *Proxy) Unmarshal(data []byte, v any) error {
dst, ok := v.(*Frame)
if !ok {
return p.parentCodec.Unmarshal(data, v)
}
dst.payload = data
return nil
}
// Name implements the encoding.Codec interface method.
func (*Proxy) Name() string {
return Name
}
// protoCodec is a Codec implementation with protobuf. It is the default rawCodec for gRPC.
type protoCodec struct{}
func (*protoCodec) Marshal(v any) ([]byte, error) {
switch x := v.(type) {
case proto.Message:
return proto.Marshal(x)
case protoV1.Message:
return protoV1.Marshal(x)
default:
return nil, fmt.Errorf("failed to marshal: message is %T, want proto.Message", x)
}
}
func (*protoCodec) Unmarshal(data []byte, v any) error {
switch x := v.(type) {
case proto.Message:
return proto.Unmarshal(data, x)
case protoV1.Message:
return protoV1.Unmarshal(data, x)
default:
return fmt.Errorf("failed to unmarshal: message is %T, want proto.Message", x)
}
}
func (*protoCodec) Name() string {
return "proxy>proto"
}