Skip to content

Commit

Permalink
Fix pooling when unmarshaling length-prefixed messages
Browse files Browse the repository at this point in the history
When we want to use a pooled buffer as a byte slice, we have to first
call `Grow` (to ensure adequate capacity) and then reslice to set
`len(raw)` to `cap(raw)`.
  • Loading branch information
akshayjshah committed Apr 9, 2022
1 parent ac70f68 commit 09a7397
Showing 1 changed file with 10 additions and 8 deletions.
18 changes: 10 additions & 8 deletions protocol_grpc_lpm.go
Original file line number Diff line number Diff line change
Expand Up @@ -177,16 +177,18 @@ func (u *unmarshaler) Unmarshal(message any) (retErr *Error) {
if size < 0 {
return errorf(CodeInvalidArgument, "message size %d overflowed uint32", size)
}
raw := u.bufferPool.Get()
defer u.bufferPool.Put(raw)
rawBuffer := u.bufferPool.Get()
defer u.bufferPool.Put(rawBuffer)
rawBuffer.Grow(size)
raw := rawBuffer.Bytes()[0:size]
if size > 0 {
// At layer 7, we don't know exactly what's happening down in L4. Large
// length-prefixed messages may arrive in chunks, so we may need to read
// the request body past EOF. We also need to take care that we don't retry
// forever if the LPM is malformed.
remaining := size
for remaining > 0 {
bytesRead, err := u.reader.Read(raw.Bytes()[size-remaining : size])
bytesRead, err := u.reader.Read(raw[size-remaining : size])
if err != nil && !errors.Is(err, io.EOF) {
return errorf(CodeUnknown, "error reading length-prefixed message data: %w", err)
}
Expand Down Expand Up @@ -225,30 +227,30 @@ func (u *unmarshaler) Unmarshal(message any) (retErr *Error) {
if err := u.compressionPool.PutDecompressor(decompressor); err != nil {
return errorf(CodeUnknown, "recycle decompressor: %w", err)
}
raw = decompressed
raw = decompressed.Bytes()
}

if isWebTrailer {
// Per the gRPC-Web specification, trailers should be encoded as an HTTP/1
// headers block _without_ the terminating newline. To make the headers
// parseable by net/textproto, we need to add the newline.
raw.WriteRune('\n')
bufferedReader := bufio.NewReader(raw)
raw = append(raw, '\n')
bufferedReader := bufio.NewReader(bytes.NewReader(raw))
mimeReader := textproto.NewReader(bufferedReader)
mimeHeader, err := mimeReader.ReadMIMEHeader()
if err != nil {
return errorf(
CodeInvalidArgument,
"gRPC-Web protocol error: received invalid trailers %q: %w",
raw.String(),
string(raw),
err,
)
}
u.webTrailer = http.Header(mimeHeader)
return errGotWebTrailers
}

if err := u.codec.Unmarshal(raw.Bytes(), message); err != nil {
if err := u.codec.Unmarshal(raw, message); err != nil {
return errorf(CodeInvalidArgument, "can't unmarshal into %T: %w", message, err)
}

Expand Down

0 comments on commit 09a7397

Please sign in to comment.