blob: 524f1bcbe0c8b37894c43c9d8a9e9d2dadc83abd [file] [log] [blame]
khenaidooa46458b2021-12-15 16:50:44 -05001package codec
2
3import (
4 "github.com/golang/protobuf/proto"
5)
6
7// EncodeVarint writes a varint-encoded integer to the Buffer.
8// This is the format for the
9// int32, int64, uint32, uint64, bool, and enum
10// protocol buffer types.
11func (cb *Buffer) EncodeVarint(x uint64) error {
12 for x >= 1<<7 {
13 cb.buf = append(cb.buf, uint8(x&0x7f|0x80))
14 x >>= 7
15 }
16 cb.buf = append(cb.buf, uint8(x))
17 return nil
18}
19
20// EncodeTagAndWireType encodes the given field tag and wire type to the
21// buffer. This combines the two values and then writes them as a varint.
22func (cb *Buffer) EncodeTagAndWireType(tag int32, wireType int8) error {
23 v := uint64((int64(tag) << 3) | int64(wireType))
24 return cb.EncodeVarint(v)
25}
26
27// EncodeFixed64 writes a 64-bit integer to the Buffer.
28// This is the format for the
29// fixed64, sfixed64, and double protocol buffer types.
30func (cb *Buffer) EncodeFixed64(x uint64) error {
31 cb.buf = append(cb.buf,
32 uint8(x),
33 uint8(x>>8),
34 uint8(x>>16),
35 uint8(x>>24),
36 uint8(x>>32),
37 uint8(x>>40),
38 uint8(x>>48),
39 uint8(x>>56))
40 return nil
41}
42
43// EncodeFixed32 writes a 32-bit integer to the Buffer.
44// This is the format for the
45// fixed32, sfixed32, and float protocol buffer types.
46func (cb *Buffer) EncodeFixed32(x uint64) error {
47 cb.buf = append(cb.buf,
48 uint8(x),
49 uint8(x>>8),
50 uint8(x>>16),
51 uint8(x>>24))
52 return nil
53}
54
55// EncodeRawBytes writes a count-delimited byte buffer to the Buffer.
56// This is the format used for the bytes protocol buffer
57// type and for embedded messages.
58func (cb *Buffer) EncodeRawBytes(b []byte) error {
59 if err := cb.EncodeVarint(uint64(len(b))); err != nil {
60 return err
61 }
62 cb.buf = append(cb.buf, b...)
63 return nil
64}
65
66// EncodeMessage writes the given message to the buffer.
67func (cb *Buffer) EncodeMessage(pm proto.Message) error {
68 bytes, err := marshalMessage(cb.buf, pm, cb.deterministic)
69 if err != nil {
70 return err
71 }
72 cb.buf = bytes
73 return nil
74}
75
76// EncodeDelimitedMessage writes the given message to the buffer with a
77// varint-encoded length prefix (the delimiter).
78func (cb *Buffer) EncodeDelimitedMessage(pm proto.Message) error {
79 bytes, err := marshalMessage(cb.tmp, pm, cb.deterministic)
80 if err != nil {
81 return err
82 }
83 // save truncated buffer if it was grown (so we can re-use it and
84 // curtail future allocations)
85 if cap(bytes) > cap(cb.tmp) {
86 cb.tmp = bytes[:0]
87 }
88 return cb.EncodeRawBytes(bytes)
89}
90
91func marshalMessage(b []byte, pm proto.Message, deterministic bool) ([]byte, error) {
92 // We try to use the most efficient way to marshal to existing slice.
93
94 if deterministic {
95 // see if the message has custom deterministic methods, preferring an
96 // "append" method over one that must always re-allocate
97 madm, ok := pm.(interface {
98 MarshalAppendDeterministic(b []byte) ([]byte, error)
99 })
100 if ok {
101 return madm.MarshalAppendDeterministic(b)
102 }
103
104 mdm, ok := pm.(interface {
105 MarshalDeterministic() ([]byte, error)
106 })
107 if ok {
108 bytes, err := mdm.MarshalDeterministic()
109 if err != nil {
110 return nil, err
111 }
112 if len(b) == 0 {
113 return bytes, nil
114 }
115 return append(b, bytes...), nil
116 }
117
118 var buf proto.Buffer
119 buf.SetDeterministic(true)
120 if err := buf.Marshal(pm); err != nil {
121 return nil, err
122 }
123 bytes := buf.Bytes()
124 if len(b) == 0 {
125 return bytes, nil
126 }
127 return append(b, bytes...), nil
128 }
129
130 mam, ok := pm.(interface {
131 // see if we can append the message, vs. having to re-allocate
132 MarshalAppend(b []byte) ([]byte, error)
133 })
134 if ok {
135 return mam.MarshalAppend(b)
136 }
137
138 // lowest common denominator
139 bytes, err := proto.Marshal(pm)
140 if err != nil {
141 return nil, err
142 }
143 if len(b) == 0 {
144 return bytes, nil
145 }
146 return append(b, bytes...), nil
147}