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

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

	out, err := mem.ReadAll(io.LimitReader(dcReader, int64(maxReceiveMessageSize)), pool)
	if err != nil {
		out.Free()
		return nil, 0, ErrMaxMessageSizeExceeded
	}
	if err = checkReceiveMessageOverflow(int64(out.Len()), int64(maxReceiveMessageSize), dcReader); err != nil {
		return nil, out.Len() + 1, err
	}
	return out, out.Len(), nil
}

// checkReceiveMessageOverflow checks if the number of bytes read from the stream exceeds
// the maximum receive message size allowed by the client. If the `readBytes` equals
// `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 an error indicating that the message size has exceeded the permissible limit.
func checkReceiveMessageOverflow(readBytes, maxReceiveMessageSize int64, dcReader io.Reader) error {
	if readBytes < maxReceiveMessageSize {
		return nil
	}

	b := make([]byte, 1)
	if n, err := dcReader.Read(b); n > 0 || err != io.EOF {
		return fmt.Errorf("overflow: received message size is larger than the allowed maxReceiveMessageSize (%d bytes)", maxReceiveMessageSize)
	}
	return nil
}

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

Review comment (bug):
I recommend changing the name to `doesReceiveMessageOverflow` and changing the return type to a boolean. Returning an error here makes the callers handle any error.

Code after:
		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 {
	failfast      bool
	preloaderInfo *compressorInfo
}