METADAPT-B supports setting a message length limit

metadaptb.go

@@ -8,6 +8,7 @@ import "git.tebibyte.media/sashakoshka/hopp/tape"
 // B implements METADAPT-B over a multiplexed stream-oriented transport such as
 // QUIC.
 type b struct {
+	sizeLimit  int
 	underlying MultiConn
 }
 
@@ -15,6 +16,7 @@ type b struct {
 // oriented transport such as TCP or UNIX domain stream sockets.
 func AdaptB(underlying MultiConn) Conn {
 	return &b {
+		sizeLimit:  defaultSizeLimit,
 		underlying: underlying,
 	}
 }
@@ -34,16 +36,28 @@ func (this *b) RemoteAddr() net.Addr {
 func (this *b) OpenTrans() (Trans, error) {
 	stream, err := this.underlying.OpenStream()
 	if err != nil { return nil, err }
-	return transB { underlying: stream }, nil
+	return this.newTrans(stream), nil
 }
 
 func (this *b) AcceptTrans() (Trans, error) {
 	stream, err := this.underlying.AcceptStream(context.Background())
 	if err != nil { return nil, err }
-	return transB { underlying: stream }, nil
+	return this.newTrans(stream), nil
+}
+
+func (this *b) SetSizeLimit(limit int) {
+	this.sizeLimit = limit
+}
+
+func (this *b) newTrans(underlying Stream) transB {
+	return transB {
+		sizeLimit: this.sizeLimit,
+		underlying: underlying,
+	}
 }
 
 type transB struct {
+	sizeLimit int
 	underlying Stream
 }
 
@@ -56,11 +70,11 @@ func (trans transB) ID() int64 {
 }
 
 func (trans transB) Send(method uint16, data []byte) error {
-	return encodeMessageB(trans.underlying, method, data)
+	return encodeMessageB(trans.underlying, trans.sizeLimit, method, data)
 }
 
 func (trans transB) Receive() (uint16, []byte, error) {
-	return decodeMessageB(trans.underlying)
+	return decodeMessageB(trans.underlying, trans.sizeLimit)
 }
 
 // MultiConn represens a multiplexed stream-oriented transport for use in
@@ -84,7 +98,10 @@ type Stream interface {
 	ID() int64
 }
 
-func encodeMessageB(writer io.Writer, method uint16, data []byte) error {
+func encodeMessageB(writer io.Writer, sizeLimit int, method uint16, data []byte) error {
+	if len(data) > sizeLimit {
+		return ErrPayloadTooLarge
+	}
 	buffer := make([]byte, 10 + len(data))
 	tape.EncodeI16(buffer[:2], method)
 	tape.EncodeI64(buffer[2:10], uint64(len(data)))
@@ -93,7 +110,7 @@ func encodeMessageB(writer io.Writer, method uint16, data []byte) error {
 	return err
 }
 
-func decodeMessageB(reader io.Reader) (uint16, []byte, error) {
+func decodeMessageB(reader io.Reader, sizeLimit int) (uint16, []byte, error) {
 	headerBuffer := [10]byte { }
 	_, err := io.ReadFull(reader, headerBuffer[:])
 	if err != nil { return 0, nil, err }
@@ -101,6 +118,9 @@ func decodeMessageB(reader io.Reader) (uint16, []byte, error) {
 	if err != nil { return 0, nil, err }
 	length, err := tape.DecodeI64[uint64](headerBuffer[2:10])
 	if err != nil { return 0, nil, err }
+	if length > uint64(sizeLimit) {
+		return 0, nil, ErrPayloadTooLarge
+	}
 	payloadBuffer := make([]byte, int(length))
 	_, err = io.ReadFull(reader, payloadBuffer)
 	if err != nil { return 0, nil, err }