@@ -63,6 +63,8 @@ const (
63
63
// The maximum write buffer size. This *must* be multiple of
64
64
// altsRecordDefaultLength.
65
65
altsWriteBufferMaxSize = 512 * 1024 // 512KiB
66
+ // The initial buffer used to read from the network.
67
+ altsReadBufferInitialSize = 32 * 1024 // 32KiB
66
68
)
67
69
68
70
var (
@@ -83,7 +85,7 @@ type conn struct {
83
85
net.Conn
84
86
crypto ALTSRecordCrypto
85
87
// buf holds data that has been read from the connection and decrypted,
86
- // but has not yet been returned by Read.
88
+ // but has not yet been returned by Read. It is a sub-slice of protected.
87
89
buf []byte
88
90
payloadLengthLimit int
89
91
// protected holds data read from the network but have not yet been
@@ -111,21 +113,13 @@ func NewConn(c net.Conn, side core.Side, recordProtocol string, key []byte, prot
111
113
}
112
114
overhead := MsgLenFieldSize + msgTypeFieldSize + crypto .EncryptionOverhead ()
113
115
payloadLengthLimit := altsRecordDefaultLength - overhead
114
- var protectedBuf []byte
115
- if protected == nil {
116
- // We pre-allocate protected to be of size
117
- // 2*altsRecordDefaultLength-1 during initialization. We only
118
- // read from the network into protected when protected does not
119
- // contain a complete frame, which is at most
120
- // altsRecordDefaultLength-1 (bytes). And we read at most
121
- // altsRecordDefaultLength (bytes) data into protected at one
122
- // time. Therefore, 2*altsRecordDefaultLength-1 is large enough
123
- // to buffer data read from the network.
124
- protectedBuf = make ([]byte , 0 , 2 * altsRecordDefaultLength - 1 )
125
- } else {
126
- protectedBuf = make ([]byte , len (protected ))
127
- copy (protectedBuf , protected )
128
- }
116
+ // We pre-allocate protected to be of size 32KB during initialization.
117
+ // We increase the size of the buffer by the required amount if it can't
118
+ // hold a complete encrypted record.
119
+ protectedBuf := make ([]byte , max (altsReadBufferInitialSize , len (protected )))
120
+ // Copy additional data from hanshaker service.
121
+ copy (protectedBuf , protected )
122
+ protectedBuf = protectedBuf [:len (protected )]
129
123
130
124
altsConn := & conn {
131
125
Conn : c ,
@@ -162,11 +156,21 @@ func (p *conn) Read(b []byte) (n int, err error) {
162
156
// Check whether a complete frame has been received yet.
163
157
for len (framedMsg ) == 0 {
164
158
if len (p .protected ) == cap (p .protected ) {
165
- tmp := make ([]byte , len (p .protected ), cap (p .protected )+ altsRecordDefaultLength )
166
- copy (tmp , p .protected )
167
- p .protected = tmp
159
+ // We can parse the length header to know exactly how large
160
+ // the buffer needs to be to hold the entire frame.
161
+ length , didParse := parseMessageLength (p .protected )
162
+ if ! didParse {
163
+ // The protected buffer is initialized with a capacity of
164
+ // larger than 4B. It should always hold the message length
165
+ // header.
166
+ panic (fmt .Sprintf ("protected buffer length shorter than expected: %d vs %d" , len (p .protected ), MsgLenFieldSize ))
167
+ }
168
+ oldProtectedBuf := p .protected
169
+ p .protected = make ([]byte , int (length )+ MsgLenFieldSize )
170
+ copy (p .protected , oldProtectedBuf )
171
+ p .protected = p .protected [:len (oldProtectedBuf )]
168
172
}
169
- n , err = p .Conn .Read (p .protected [len (p .protected ):min ( cap (p .protected ), len ( p . protected ) + altsRecordDefaultLength )])
173
+ n , err = p .Conn .Read (p .protected [len (p .protected ):cap (p .protected )])
170
174
if err != nil {
171
175
return 0 , err
172
176
}
@@ -185,6 +189,15 @@ func (p *conn) Read(b []byte) (n int, err error) {
185
189
}
186
190
ciphertext := msg [msgTypeFieldSize :]
187
191
192
+ // Decrypt directly into the buffer, avoiding a copy from p.buf if
193
+ // possible.
194
+ if len (b ) >= len (ciphertext ) {
195
+ dec , err := p .crypto .Decrypt (b [:0 ], ciphertext )
196
+ if err != nil {
197
+ return 0 , err
198
+ }
199
+ return len (dec ), nil
200
+ }
188
201
// Decrypt requires that if the dst and ciphertext alias, they
189
202
// must alias exactly. Code here used to use msg[:0], but msg
190
203
// starts MsgLenFieldSize+msgTypeFieldSize bytes earlier than
0 commit comments