khenaidoo | f333355 | 2021-12-15 16:52:31 -0500 | [diff] [blame] | 1 | package dynamic |
| 2 | |
| 3 | // Binary serialization and de-serialization for dynamic messages |
| 4 | |
| 5 | import ( |
| 6 | "fmt" |
| 7 | "github.com/golang/protobuf/proto" |
| 8 | "github.com/jhump/protoreflect/codec" |
| 9 | "io" |
| 10 | ) |
| 11 | |
| 12 | // defaultDeterminism, if true, will mean that calls to Marshal will produce |
| 13 | // deterministic output. This is used to make the output of proto.Marshal(...) |
| 14 | // deterministic (since there is no way to have that convey determinism intent). |
| 15 | // **This is only used from tests.** |
| 16 | var defaultDeterminism = false |
| 17 | |
| 18 | // Marshal serializes this message to bytes, returning an error if the operation |
| 19 | // fails. The resulting bytes are in the standard protocol buffer binary format. |
| 20 | func (m *Message) Marshal() ([]byte, error) { |
| 21 | var b codec.Buffer |
| 22 | b.SetDeterministic(defaultDeterminism) |
| 23 | if err := m.marshal(&b); err != nil { |
| 24 | return nil, err |
| 25 | } |
| 26 | return b.Bytes(), nil |
| 27 | } |
| 28 | |
| 29 | // MarshalAppend behaves exactly the same as Marshal, except instead of allocating a |
| 30 | // new byte slice to marshal into, it uses the provided byte slice. The backing array |
| 31 | // for the returned byte slice *may* be the same as the one that was passed in, but |
| 32 | // it's not guaranteed as a new backing array will automatically be allocated if |
| 33 | // more bytes need to be written than the provided buffer has capacity for. |
| 34 | func (m *Message) MarshalAppend(b []byte) ([]byte, error) { |
| 35 | codedBuf := codec.NewBuffer(b) |
| 36 | codedBuf.SetDeterministic(defaultDeterminism) |
| 37 | if err := m.marshal(codedBuf); err != nil { |
| 38 | return nil, err |
| 39 | } |
| 40 | return codedBuf.Bytes(), nil |
| 41 | } |
| 42 | |
| 43 | // MarshalDeterministic serializes this message to bytes in a deterministic way, |
| 44 | // returning an error if the operation fails. This differs from Marshal in that |
| 45 | // map keys will be sorted before serializing to bytes. The protobuf spec does |
| 46 | // not define ordering for map entries, so Marshal will use standard Go map |
| 47 | // iteration order (which will be random). But for cases where determinism is |
| 48 | // more important than performance, use this method instead. |
| 49 | func (m *Message) MarshalDeterministic() ([]byte, error) { |
| 50 | var b codec.Buffer |
| 51 | b.SetDeterministic(true) |
| 52 | if err := m.marshal(&b); err != nil { |
| 53 | return nil, err |
| 54 | } |
| 55 | return b.Bytes(), nil |
| 56 | } |
| 57 | |
| 58 | // MarshalAppendDeterministic behaves exactly the same as MarshalDeterministic, |
| 59 | // except instead of allocating a new byte slice to marshal into, it uses the |
| 60 | // provided byte slice. The backing array for the returned byte slice *may* be |
| 61 | // the same as the one that was passed in, but it's not guaranteed as a new |
| 62 | // backing array will automatically be allocated if more bytes need to be written |
| 63 | // than the provided buffer has capacity for. |
| 64 | func (m *Message) MarshalAppendDeterministic(b []byte) ([]byte, error) { |
| 65 | codedBuf := codec.NewBuffer(b) |
| 66 | codedBuf.SetDeterministic(true) |
| 67 | if err := m.marshal(codedBuf); err != nil { |
| 68 | return nil, err |
| 69 | } |
| 70 | return codedBuf.Bytes(), nil |
| 71 | } |
| 72 | |
| 73 | func (m *Message) marshal(b *codec.Buffer) error { |
| 74 | if err := m.marshalKnownFields(b); err != nil { |
| 75 | return err |
| 76 | } |
| 77 | return m.marshalUnknownFields(b) |
| 78 | } |
| 79 | |
| 80 | func (m *Message) marshalKnownFields(b *codec.Buffer) error { |
| 81 | for _, tag := range m.knownFieldTags() { |
| 82 | itag := int32(tag) |
| 83 | val := m.values[itag] |
| 84 | fd := m.FindFieldDescriptor(itag) |
| 85 | if fd == nil { |
| 86 | panic(fmt.Sprintf("Couldn't find field for tag %d", itag)) |
| 87 | } |
| 88 | if err := b.EncodeFieldValue(fd, val); err != nil { |
| 89 | return err |
| 90 | } |
| 91 | } |
| 92 | return nil |
| 93 | } |
| 94 | |
| 95 | func (m *Message) marshalUnknownFields(b *codec.Buffer) error { |
| 96 | for _, tag := range m.unknownFieldTags() { |
| 97 | itag := int32(tag) |
| 98 | sl := m.unknownFields[itag] |
| 99 | for _, u := range sl { |
| 100 | if err := b.EncodeTagAndWireType(itag, u.Encoding); err != nil { |
| 101 | return err |
| 102 | } |
| 103 | switch u.Encoding { |
| 104 | case proto.WireBytes: |
| 105 | if err := b.EncodeRawBytes(u.Contents); err != nil { |
| 106 | return err |
| 107 | } |
| 108 | case proto.WireStartGroup: |
| 109 | _, _ = b.Write(u.Contents) |
| 110 | if err := b.EncodeTagAndWireType(itag, proto.WireEndGroup); err != nil { |
| 111 | return err |
| 112 | } |
| 113 | case proto.WireFixed32: |
| 114 | if err := b.EncodeFixed32(u.Value); err != nil { |
| 115 | return err |
| 116 | } |
| 117 | case proto.WireFixed64: |
| 118 | if err := b.EncodeFixed64(u.Value); err != nil { |
| 119 | return err |
| 120 | } |
| 121 | case proto.WireVarint: |
| 122 | if err := b.EncodeVarint(u.Value); err != nil { |
| 123 | return err |
| 124 | } |
| 125 | default: |
| 126 | return codec.ErrBadWireType |
| 127 | } |
| 128 | } |
| 129 | } |
| 130 | return nil |
| 131 | } |
| 132 | |
| 133 | // Unmarshal de-serializes the message that is present in the given bytes into |
| 134 | // this message. It first resets the current message. It returns an error if the |
| 135 | // given bytes do not contain a valid encoding of this message type. |
| 136 | func (m *Message) Unmarshal(b []byte) error { |
| 137 | m.Reset() |
| 138 | if err := m.UnmarshalMerge(b); err != nil { |
| 139 | return err |
| 140 | } |
| 141 | return m.Validate() |
| 142 | } |
| 143 | |
| 144 | // UnmarshalMerge de-serializes the message that is present in the given bytes |
| 145 | // into this message. Unlike Unmarshal, it does not first reset the message, |
| 146 | // instead merging the data in the given bytes into the existing data in this |
| 147 | // message. |
| 148 | func (m *Message) UnmarshalMerge(b []byte) error { |
| 149 | return m.unmarshal(codec.NewBuffer(b), false) |
| 150 | } |
| 151 | |
| 152 | func (m *Message) unmarshal(buf *codec.Buffer, isGroup bool) error { |
| 153 | for !buf.EOF() { |
| 154 | fd, val, err := buf.DecodeFieldValue(m.FindFieldDescriptor, m.mf) |
| 155 | if err != nil { |
| 156 | if err == codec.ErrWireTypeEndGroup { |
| 157 | if isGroup { |
| 158 | // finished parsing group |
| 159 | return nil |
| 160 | } |
| 161 | return codec.ErrBadWireType |
| 162 | } |
| 163 | return err |
| 164 | } |
| 165 | |
| 166 | if fd == nil { |
| 167 | if m.unknownFields == nil { |
| 168 | m.unknownFields = map[int32][]UnknownField{} |
| 169 | } |
| 170 | uv := val.(codec.UnknownField) |
| 171 | u := UnknownField{ |
| 172 | Encoding: uv.Encoding, |
| 173 | Value: uv.Value, |
| 174 | Contents: uv.Contents, |
| 175 | } |
| 176 | m.unknownFields[uv.Tag] = append(m.unknownFields[uv.Tag], u) |
| 177 | } else if err := mergeField(m, fd, val); err != nil { |
| 178 | return err |
| 179 | } |
| 180 | } |
| 181 | if isGroup { |
| 182 | return io.ErrUnexpectedEOF |
| 183 | } |
| 184 | return nil |
| 185 | } |