Repository: grpc/grpc-go
File: rpc_util.go (Go)

Code before:
// 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 {
		uncompressed, err := dc.Do(d.Reader())
		if err != nil {
			return nil, status.Errorf(codes.Internal, "grpc: failed to decompress the received message: %v", err)
		}
		if len(uncompressed) > maxReceiveMessageSize {
			return nil, status.Errorf(codes.ResourceExhausted, "grpc: message after decompression larger than max (%d vs. %d)", len(uncompressed), maxReceiveMessageSize)
		}
		return mem.BufferSlice{mem.SliceBuffer(uncompressed)}, nil
	}
	if compressor != nil {
		dcReader, err := compressor.Decompress(d.Reader())
		if err != nil {
			return nil, status.Errorf(codes.Internal, "grpc: failed to decompress the message: %v", err)
		}

		out, err := mem.ReadAll(io.LimitReader(dcReader, int64(maxReceiveMessageSize)), pool)
		if err != nil {
			out.Free()
			return nil, status.Errorf(codes.Internal, "grpc: failed to read decompressed data: %v", err)
		}

		if doesReceiveMessageOverflow(out.Len(), maxReceiveMessageSize, dcReader) {
			out.Free()
			return nil, status.Errorf(codes.ResourceExhausted, "grpc: received message after decompression larger than max %d", maxReceiveMessageSize)
		}
		return out, nil
	}
	return d, nil
}

// doesReceiveMessageOverflow checks if the number of bytes read from the stream
// exceed the maximum receive message size allowed by the receiver. If the `readBytes`
// is greater than or equal to `maxReceiveMessageSize`, the function attempts to read
// one more byte from the `dcReader` to detect if there's an overflow.
//
// If additional data is read, or an error other than `io.EOF` is encountered, the function
// returns `true` to indicate that the message size has exceeded the permissible limit.
// Otherwise, it returns `false` indicating no overflow.
func doesReceiveMessageOverflow(readBytes, maxReceiveMessageSize int, dcReader io.Reader) bool {
	if readBytes < maxReceiveMessageSize {
		return false
	}

	b := make([]byte, 1)
	if n, err := dcReader.Read(b); n > 0 || err != io.EOF {
		return true
	}

Review comment (bug):
Nit:

How about something like this instead:

```go
if out.Len() == maxReceiveMessageSize && !atEOF(dcReader) {
```

```go
// atEOF reads data from r and returns true iff zero bytes could be read and r.Read returns EOF.
func atEOF(r io.Reader) bool {
	n, err := dcReader.Read(make([]byte, 1))
	return n == 0 && err == io.EOF
}
```

(Reasoning: "doesReceiveMessageOverflow" is very special-purpose, combines multiple operations into a single thing, and hides/splits relevant logic (reading & checking the length), whereas "atEOF" is potentially reusable, and does exactly one thing.)

Code after:
// 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 {
		uncompressed, err := dc.Do(d.Reader())
		if err != nil {
			return nil, status.Errorf(codes.Internal, "grpc: failed to decompress the received message: %v", err)
		}
		if len(uncompressed) > maxReceiveMessageSize {
			return nil, status.Errorf(codes.ResourceExhausted, "grpc: message after decompression larger than max (%d vs. %d)", len(uncompressed), maxReceiveMessageSize)
		}
		return mem.BufferSlice{mem.SliceBuffer(uncompressed)}, nil
	}
	if compressor != nil {
		dcReader, err := compressor.Decompress(d.Reader())
		if err != nil {
			return nil, status.Errorf(codes.Internal, "grpc: failed to decompress the message: %v", err)
		}

		out, err := mem.ReadAll(io.LimitReader(dcReader, int64(maxReceiveMessageSize)), pool)
		if err != nil {
			out.Free()
			return nil, status.Errorf(codes.Internal, "grpc: failed to read decompressed data: %v", err)
		}

		if out.Len() == maxReceiveMessageSize && !atEOF(dcReader) {
			out.Free()
			return nil, status.Errorf(codes.ResourceExhausted, "grpc: received message after decompression larger than max %d", maxReceiveMessageSize)
		}
		return out, nil
	}
	return nil, status.Errorf(codes.Internal, "grpc: no decompressor available for compressed payload")
}

// atEOF reads data from r and returns true if zero bytes could be read and r.Read returns EOF.
func atEOF(dcReader io.Reader) bool {
	n, err := dcReader.Read(make([]byte, 1))
	return n == 0 && err == io.EOF
}

type recvCompressor interface {
	RecvCompress() string
}

// For the two compressor parameters, both should not be set, but if they are,
// dc takes precedence over compressor.
// TODO(dfawley): wrap the old compressor/decompressor using the new API?
func recv(p *parser, c baseCodec, s recvCompressor, dc Decompressor, m any, maxReceiveMessageSize int, payInfo *payloadInfo, compressor encoding.Compressor, isServer bool) error {
	data, err := recvAndDecompress(p, s, dc, maxReceiveMessageSize, payInfo, compressor, isServer)
	if err != nil {
		return err