8000
Skip to content

rpc_util.go

Latest commit

 

History

History
1207 lines (1072 loc) · 37.1 KB

rpc_util.go

File metadata and controls

1207 lines (1072 loc) · 37.1 KB
Raw
type payloadInfo struct {
compressedLength int // The compressed length got from wire.
uncompressedBytes mem.BufferSlice
}
func (p *payloadInfo) free() {
if p != nil && p.uncompressedBytes != nil {
p.uncompressedBytes.Free()
}
}
// recvAndDecompress reads a message from the stream, decompressing it if necessary.
//
// Cancelling the returned cancel function releases the buffer back to the pool. So the caller should cancel as soon as
// the buffer is no longer needed.
// TODO: Refactor this function to reduce the number of arguments.
// See: https://google.github.io/styleguide/go/best-practices.html#function-argument-lists
func recvAndDecompress(p *parser, s recvCompressor, dc Decompressor, maxReceiveMessageSize int, payInfo *payloadInfo, compressor encoding.Compressor, isServer bool) (out mem.BufferSlice, err error) {
pf, compressed, err := p.recvMsg(maxReceiveMessageSize)
if err != nil {
return nil, err
}
compressedLength := compressed.Len()
if st := checkRecvPayload(pf, s.RecvCompress(), compressor != nil || dc != nil, isServer); st != nil {
compressed.Free()
return nil, st.Err()
}
if pf.isCompressed() {
defer compressed.Free()
// To match legacy behavior, if the decompressor is set by WithDecompressor or RPCDecompressor,
// use this decompressor as the default.
out, err = decompress(compressor, compressed, dc, maxReceiveMessageSize, p.bufferPool)
if err != nil {
return nil, err
}
} else {
out = compressed
}
if payInfo != nil {
payInfo.compressedLength = compressedLength
out.Ref()
payInfo.uncompressedBytes = out
}
return out, nil
}
// decompress processes the given data by decompressing it using either
// a custom decompressor or a standard compressor. If a custom decompressor
// is provided, it takes precedence. The function validates that
// the decompressed data does not exceed the specified maximum size and returns
// an error if this limit is exceeded. On success, it returns the decompressed
// data. Otherwise, it returns an error if decompression fails or the data
// exceeds the size limit.
func decompress(compressor encoding.Compressor, d mem.BufferSlice, dc Decompressor, maxReceiveMessageSize int, pool mem.BufferPool) (mem.BufferSlice, error) {
if dc != nil {
r := d.Reader()
// For the built-in gzip decompressor, bound the decompressed output
// at maxReceiveMessageSize+1 so that a small but highly compressed
// payload (a "zip bomb") cannot expand to gigabytes in memory before
// the post-decompression size check below has a chance to fire. The
// Decompressor interface does not accept an extra size parameter,
// so we type-assert to invoke a size-aware helper. Third-party
// Decompressor implementations keep the original Do behavior.
var uncompressed []byte
var err error
if gd, ok := dc.(*gzipDecompressor); ok {
uncompressed, err = gd.doWithMaxSize(r, int64(maxReceiveMessageSize))
} else {
uncompressed, err = dc.Do(r)
View remainder of file in raw view
0