client,server: configurable wire message size limits.

Implement configurable limits for the maximum accepted message
size of the wire protocol. The default limit can be overridden
using the WithClientWireMessageLimit() option for clients and
using the WithServerWireMessageLimit() option for servers.

Signed-off-by: Krisztian Litkey <[email protected]>

Author: klihub

channel.go

@@ -31,6 +31,7 @@ import (
 const (
 	messageHeaderLength = 10
 	messageLengthMax    = 4 << 20
+	messageLengthMin    = 4 << 10
 )
 
 type messageType uint8
@@ -96,18 +97,23 @@ func writeMessageHeader(w io.Writer, p []byte, mh messageHeader) error {
 var buffers sync.Pool
 
 type channel struct {
-	conn  net.Conn
-	bw    *bufio.Writer
-	br    *bufio.Reader
-	hrbuf [messageHeaderLength]byte // avoid alloc when reading header
-	hwbuf [messageHeaderLength]byte
+	conn      net.Conn
+	bw        *bufio.Writer
+	br        *bufio.Reader
+	hrbuf     [messageHeaderLength]byte // avoid alloc when reading header
+	hwbuf     [messageHeaderLength]byte
+	maxMsgLen int
 }
 
-func newChannel(conn net.Conn) *channel {
+func newChannel(conn net.Conn, maxMsgLen int) *channel {
+	if maxMsgLen == 0 {
+		maxMsgLen = messageLengthMax
+	}
 	return &channel{
-		conn: conn,
-		bw:   bufio.NewWriter(conn),
-		br:   bufio.NewReader(conn),
+		conn:      conn,
+		bw:        bufio.NewWriter(conn),
+		br:        bufio.NewReader(conn),
+		maxMsgLen: maxMsgLen,
 	}
 }
 
@@ -123,12 +129,12 @@ func (ch *channel) recv() (messageHeader, []byte, error) {
 		return messageHeader{}, nil, err
 	}
 
-	if mh.Length > uint32(messageLengthMax) {
+	if mh.Length > uint32(ch.maxMsgLen) {
 		if _, err := ch.br.Discard(int(mh.Length)); err != nil {
 			return mh, nil, fmt.Errorf("failed to discard after receiving oversized message: %w", err)
 		}
 
-		return mh, nil, status.Errorf(codes.ResourceExhausted, "message length %v exceed maximum message size of %v", mh.Length, messageLengthMax)
+		return mh, nil, status.Errorf(codes.ResourceExhausted, "message length %v exceed maximum message size of %v", mh.Length, ch.maxMsgLen)
 	}
 
 	var p []byte
@@ -147,6 +153,7 @@ func (ch *channel) send(streamID uint32, t messageType, flags uint8, p []byte) e
 	//if len(p) > messageLengthMax {
 	//	return status.Errorf(codes.InvalidArgument, "refusing to send, message length %v exceed maximum message size of %v", len(p), messageLengthMax)
 	//}
+
 	if err := writeMessageHeader(ch.bw, ch.hwbuf[:], messageHeader{Length: uint32(len(p)), StreamID: streamID, Type: t, Flags: flags}); err != nil {
 		return err
 	}

client.go

@@ -35,9 +35,10 @@ import (
 
 // Client for a ttrpc server
 type Client struct {
-	codec   codec
-	conn    net.Conn
-	channel *channel
+	codec     codec
+	conn      net.Conn
+	channel   *channel
+	maxMsgLen int
 
 	streamLock   sync.RWMutex
 	streams      map[streamID]*stream
@@ -107,14 +108,25 @@ func chainUnaryInterceptors(interceptors []UnaryClientInterceptor, final Invoker
 	}
 }
 
+// WithClientWireMessageLimit sets the maximum allowed message length on the wire for the client.
+func WithClientWireMessageLimit(maxMsgLen int) ClientOpts {
+	switch {
+	case maxMsgLen == 0:
+		maxMsgLen = messageLengthMax
+	case maxMsgLen < messageLengthMin:
+		maxMsgLen = messageLengthMin
+	}
+	return func(c *Client) {
+		c.maxMsgLen = maxMsgLen
+	}
+}
+
 // NewClient creates a new ttrpc client using the given connection
 func NewClient(conn net.Conn, opts ...ClientOpts) *Client {
 	ctx, cancel := context.WithCancel(context.Background())
-	channel := newChannel(conn)
 	c := &Client{
 		codec:           codec{},
 		conn:            conn,
-		channel:         channel,
 		streams:         make(map[streamID]*stream),
 		nextStreamID:    1,
 		closed:          cancel,
@@ -127,6 +139,8 @@ func NewClient(conn net.Conn, opts ...ClientOpts) *Client {
 		o(c)
 	}
 
+	c.channel = newChannel(conn, c.maxMsgLen)
+
 	if c.interceptor == nil {
 		c.interceptor = defaultClientInterceptor
 	}

config.go

@@ -24,6 +24,7 @@ import (
 type serverConfig struct {
 	handshaker  Handshaker
 	interceptor UnaryServerInterceptor
+	maxMsgLen   int
 }
 
 // ServerOpt for configuring a ttrpc server
@@ -84,3 +85,17 @@ func chainUnaryServerInterceptors(info *UnaryServerInfo, method Method, intercep
 			chainUnaryServerInterceptors(info, method, interceptors[1:]))
 	}
 }
+
+// WithServerWireMessageLimit sets the maximum allowed message length on the wire for the server.
+func WithServerWireMessageLimit(maxMsgLen int) ServerOpt {
+	switch {
+	case maxMsgLen == 0:
+		maxMsgLen = messageLengthMax
+	case maxMsgLen < messageLengthMin:
+		maxMsgLen = messageLengthMin
+	}
+	return func(c *serverConfig) error {
+		c.maxMsgLen = maxMsgLen
+		return nil
+	}
+}

server.go

@@ -332,7 +332,7 @@ func (c *serverConn) run(sctx context.Context) {
 	)
 
 	var (
-		ch                     = newChannel(c.conn)
+		ch                     = newChannel(c.conn, c.server.config.maxMsgLen)
 		ctx, cancel            = context.WithCancel(sctx)
 		state        connState = connStateIdle
 		responses              = make(chan response)