forked from prysmaticlabs/prysm
-
Notifications
You must be signed in to change notification settings - Fork 0
/
ssz.go
163 lines (147 loc) · 4.41 KB
/
ssz.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
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
package encoder
import (
"bytes"
"fmt"
"io"
"github.com/gogo/protobuf/proto"
"github.com/golang/snappy"
"github.com/prysmaticlabs/go-ssz"
"github.com/sirupsen/logrus"
)
var _ = NetworkEncoding(&SszNetworkEncoder{})
// MaxChunkSize allowed for decoding messages.
const MaxChunkSize = uint64(1 << 20) // 1Mb
// SszNetworkEncoder supports p2p networking encoding using SimpleSerialize
// with snappy compression (if enabled).
type SszNetworkEncoder struct {
UseSnappyCompression bool
}
func (e SszNetworkEncoder) doEncode(msg interface{}) ([]byte, error) {
return ssz.Marshal(msg)
}
// Encode the proto message to the io.Writer.
func (e SszNetworkEncoder) Encode(w io.Writer, msg interface{}) (int, error) {
if msg == nil {
return 0, nil
}
b, err := e.doEncode(msg)
if err != nil {
return 0, err
}
if e.UseSnappyCompression {
return writeSnappyBuffer(w, b)
}
return w.Write(b)
}
// EncodeWithLength the proto message to the io.Writer. This encoding prefixes the byte slice with a protobuf varint
// to indicate the size of the message.
func (e SszNetworkEncoder) EncodeWithLength(w io.Writer, msg interface{}) (int, error) {
if msg == nil {
return 0, nil
}
b, err := e.doEncode(msg)
if err != nil {
return 0, err
}
// write varint first
_, err = w.Write(proto.EncodeVarint(uint64(len(b))))
if err != nil {
return 0, err
}
if e.UseSnappyCompression {
return writeSnappyBuffer(w, b)
}
return w.Write(b)
}
// EncodeWithMaxLength the proto message to the io.Writer. This encoding prefixes the byte slice with a protobuf varint
// to indicate the size of the message. This checks that the encoded message isn't larger than the provided max limit.
func (e SszNetworkEncoder) EncodeWithMaxLength(w io.Writer, msg interface{}, maxSize uint64) (int, error) {
if msg == nil {
return 0, nil
}
b, err := e.doEncode(msg)
if err != nil {
return 0, err
}
if uint64(len(b)) > maxSize {
return 0, fmt.Errorf("size of encoded message is %d which is larger than the provided max limit of %d", len(b), maxSize)
}
// write varint first
_, err = w.Write(proto.EncodeVarint(uint64(len(b))))
if err != nil {
return 0, err
}
if e.UseSnappyCompression {
return writeSnappyBuffer(w, b)
}
return w.Write(b)
}
func (e SszNetworkEncoder) doDecode(b []byte, to interface{}) error {
return ssz.Unmarshal(b, to)
}
// Decode the bytes to the protobuf message provided.
func (e SszNetworkEncoder) Decode(b []byte, to interface{}) error {
if e.UseSnappyCompression {
newBuffer := bytes.NewBuffer(b)
r := snappy.NewReader(newBuffer)
newObj := make([]byte, len(b))
numOfBytes, err := r.Read(newObj)
if err != nil {
return err
}
return e.doDecode(newObj[:numOfBytes], to)
}
return e.doDecode(b, to)
}
// DecodeWithLength the bytes from io.Reader to the protobuf message provided.
func (e SszNetworkEncoder) DecodeWithLength(r io.Reader, to interface{}) error {
return e.DecodeWithMaxLength(r, to, MaxChunkSize)
}
// DecodeWithMaxLength the bytes from io.Reader to the protobuf message provided.
// This checks that the decoded message isn't larger than the provided max limit.
func (e SszNetworkEncoder) DecodeWithMaxLength(r io.Reader, to interface{}, maxSize uint64) error {
if maxSize > MaxChunkSize {
return fmt.Errorf("maxSize %d exceeds max chunk size %d", maxSize, MaxChunkSize)
}
msgLen, err := readVarint(r)
if err != nil {
return err
}
if e.UseSnappyCompression {
r = snappy.NewReader(r)
}
if msgLen > maxSize {
return fmt.Errorf("size of decoded message is %d which is larger than the provided max limit of %d", msgLen, maxSize)
}
b := make([]byte, e.MaxLength(int(msgLen)))
numOfBytes, err := r.Read(b)
if err != nil {
return err
}
return e.doDecode(b[:numOfBytes], to)
}
// ProtocolSuffix returns the appropriate suffix for protocol IDs.
func (e SszNetworkEncoder) ProtocolSuffix() string {
if e.UseSnappyCompression {
return "/ssz_snappy"
}
return "/ssz"
}
// MaxLength specifies the maximum possible length of an encoded
// chunk of data.
func (e SszNetworkEncoder) MaxLength(length int) int {
if e.UseSnappyCompression {
return snappy.MaxEncodedLen(length)
}
return length
}
// Writes a bytes value through a snappy buffered writer.
func writeSnappyBuffer(w io.Writer, b []byte) (int, error) {
bufWriter := snappy.NewBufferedWriter(w)
defer func() {
if err := bufWriter.Close(); err != nil {
logrus.WithError(err).Error("Failed to close snappy buffered writer")
}
}()
return bufWriter.Write(b)
}