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.
// Optionally, if data will be over maxReceiveMessageSize, just return the size.
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, err
	}
	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 {
		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
}

Review comment (bug):
Adding 1 to `out.Len()` can result in an overflow on 32 bit systems. The  following part of the function seems strange to me: 
> Optionally, if data will be over maxReceiveMessageSize, just return the size

I suggest making the following change to avoid this:
1. Declare a global error for indicating that the max receive size is exceeded:
```go
var	errMaxMessageSizeExceeded = errors.New("max message size exceeded")
```  
2. When the check here fails, `nil, 0, errMaxMessageSizeExceeded`. Update the godoc to mention the same.
3. In the caller, i.e. recvAndDecompress, instead of checking `if size > maxReceiveMessageSize`, check `if err == errMaxMessageSizeExceeded`. Also update the error message to not mention the actual size, because we didn't read the entire message anyways: `grpc: received message after decompression larger than max %d`.

This ensures we're using the returned error to indicate the failure instead of relying on special values of the bytes read count.

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

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