From 1ac0ed51c7769e166ab7c70794ccf6ff2273002a Mon Sep 17 00:00:00 2001 From: Sasha Koshka Date: Fri, 4 Apr 2025 16:07:20 -0400 Subject: [PATCH] METADAPT-B supports setting a message length limit --- connection.go | 6 ++++++ metadaptb.go | 32 ++++++++++++++++++++++++++------ metadaptb_test.go | 17 +++++++++++++---- 3 files changed, 45 insertions(+), 10 deletions(-) diff --git a/connection.go b/connection.go index adefbdf..11f6fe9 100644 --- a/connection.go +++ b/connection.go @@ -3,6 +3,8 @@ package hopp import "net" // import "time" +const defaultSizeLimit = 1024 * 1024 // 1 megabyte + // Conn is a HOPP connection. type Conn interface { // Close closes the connection. Any blocked operations on the connection @@ -19,6 +21,10 @@ type Conn interface { // AcceptTrans accepts a transaction from the other party. This must // be called in a loop to avoid the connection locking up. AcceptTrans() (Trans, error) + + // SetSizeLimit sets a limit (in bytes) for how large messages can be. + // By default, this limit is 1 megabyte. + SetSizeLimit(limit int) } // Trans is a HOPP transaction. diff --git a/metadaptb.go b/metadaptb.go index 52e98fc..1c6eb9f 100644 --- a/metadaptb.go +++ b/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 } diff --git a/metadaptb_test.go b/metadaptb_test.go index dc0e8d4..add8043 100644 --- a/metadaptb_test.go +++ b/metadaptb_test.go @@ -9,7 +9,7 @@ import "testing" func TestEncodeMessageB(test *testing.T) { buffer := new(bytes.Buffer) payload := []byte { 0x00, 0x01, 0x02, 0x03, 0x04, 0x05 } - err := encodeMessageB(buffer, 0x6B12, payload) + err := encodeMessageB(buffer, defaultSizeLimit, 0x6B12, payload) correct := []byte { 0x6B, 0x12, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x06, @@ -23,12 +23,21 @@ func TestEncodeMessageB(test *testing.T) { } } +func TestEncodeMessageBErr(test *testing.T) { + buffer := new(bytes.Buffer) + payload := make([]byte, 0x10000) + err := encodeMessageB(buffer, 255, 0x6B12, payload) + if !errors.Is(err, ErrPayloadTooLarge) { + test.Fatalf("wrong error: %v", err) + } +} + func TestDecodeMessageB(test *testing.T) { method, payload, err := decodeMessageB(bytes.NewReader([]byte { 0x6B, 0x12, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x06, 0x00, 0x01, 0x02, 0x03, 0x04, 0x05, - })) + }), defaultSizeLimit) if err != nil { test.Fatal(err) } @@ -43,10 +52,10 @@ func TestDecodeMessageB(test *testing.T) { func TestDecodeMessageBErr(test *testing.T) { _, _, err := decodeMessageB(bytes.NewReader([]byte { - 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x6B, 0x12, + 0x6B, 0x12, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x06, 0x00, 0x01, 0x02, 0x03, 0x04, 0x05, - })) + }), defaultSizeLimit) if !errors.Is(err, io.ErrUnexpectedEOF) { test.Fatalf("wrong error: %v", err) }