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

Code before:
	if payInfo != nil {
		payInfo.compressedLength = compressedLength
		out.Ref()
		payInfo.uncompressedBytes = out
	}

	return out, nil
}

// Using compressor, decompress d, returning data and size.
// If the decompressed data exceeds maxReceiveMessageSize, it returns errMaxMessageSizeExceeded.
func decompress(compressor encoding.Compressor, d mem.BufferSlice, maxReceiveMessageSize int, pool mem.BufferPool) (mem.BufferSlice, error) {
	dcReader, err := compressor.Decompress(d.Reader())
	if err != nil {
		return nil, err
	}

	out, err := mem.ReadAll(io.LimitReader(dcReader, int64(maxReceiveMessageSize)), pool)
	if err != nil {
		out.Free()
		return nil, err
	}

	if doesReceiveMessageOverflow(out.Len(), maxReceiveMessageSize, dcReader) {
		out.Free()
		return nil, errMaxMessageSizeExceeded
	}
	return out, 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
	}
	return false
}

type recvCompressor interface {

Review comment (suggestion):
Can this just return the `status.Errorf(...)` instead of returning a magic error that gets converted into that?

(We'd want to make sure this method always returns an error that is appropriate to just return directly from `recvAndDecompress`.)

Code after:
		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
	}

	// If the codec wants its own reference to the data, it can get it. Otherwise, always
	// free the buffers.
	defer data.Free()

	if err := c.Unmarshal(data, m); err != nil {
		return status.Errorf(codes.Internal, "grpc: failed to unmarshal the received message: %v", err)
	}

	return nil
}

// Information about RPC
type rpcInfo struct {