diff --git a/connsocket.go b/connsocket.go new file mode 100644 index 00000000..3b35256e --- /dev/null +++ b/connsocket.go @@ -0,0 +1,50 @@ +// SPDX-FileCopyrightText: 2025 The Pion community +// SPDX-License-Identifier: MIT + +package transport + +import ( + "net" + "time" +) + +type NetConnSocket interface { + net.Conn + + ReadWithAttributes(p []byte, attr *PacketAttributes) (n int, err error) +} + +type PacketConnSocket interface { + net.PacketConn + + ReadFromWithAttributes(p []byte, attr *PacketAttributes) (n int, addr net.Addr, err error) +} + +// NetConnToNetConnSocket wraps a net.Conn and implements the PacketStream interface by delegating +// calls to the underlying connection. ReadWithAttributes delegates to Read and +// ignores the provided PacketAttributes. +type NetConnToNetConnSocket struct { + conn net.Conn +} + +// NewNetConnToNetConnSocket returns a new Proxy that wraps the provided net.Conn. +func NewNetConnToNetConnSocket(conn net.Conn) *NetConnToNetConnSocket { + return &NetConnToNetConnSocket{conn: conn} +} + +// ReadWithAttributes reads from the underlying connection and ignores attributes. +func (p *NetConnToNetConnSocket) ReadWithAttributes(b []byte, _ *PacketAttributes) (int, error) { + return p.conn.Read(b) +} + +// Delegate net.Conn methods to the underlying connection. +func (p *NetConnToNetConnSocket) Read(b []byte) (int, error) { return p.conn.Read(b) } +func (p *NetConnToNetConnSocket) Write(b []byte) (int, error) { return p.conn.Write(b) } +func (p *NetConnToNetConnSocket) Close() error { return p.conn.Close() } +func (p *NetConnToNetConnSocket) LocalAddr() net.Addr { return p.conn.LocalAddr() } +func (p *NetConnToNetConnSocket) RemoteAddr() net.Addr { return p.conn.RemoteAddr() } +func (p *NetConnToNetConnSocket) SetDeadline(t time.Time) error { return p.conn.SetDeadline(t) } +func (p *NetConnToNetConnSocket) SetReadDeadline(t time.Time) error { return p.conn.SetReadDeadline(t) } +func (p *NetConnToNetConnSocket) SetWriteDeadline(t time.Time) error { + return p.conn.SetWriteDeadline(t) +} diff --git a/packetattributes.go b/packetattributes.go new file mode 100644 index 00000000..7f27d1bb --- /dev/null +++ b/packetattributes.go @@ -0,0 +1,42 @@ +// SPDX-FileCopyrightText: 2025 The Pion community +// SPDX-License-Identifier: MIT + +package transport + +const MaxAttributesLen = 1024 + +type PacketAttributes struct { + Buffer []byte + BytesCopied int +} + +func (p *PacketAttributes) Reset() { + p.BytesCopied = 0 +} + +func NewPacketAttributesWithLen(length int) *PacketAttributes { + buff := make([]byte, length) + + return &PacketAttributes{ + Buffer: buff, + BytesCopied: 0, + } +} + +func (p *PacketAttributes) Clone() *PacketAttributes { + b := make([]byte, p.BytesCopied) + copy(b, p.Buffer) + return &PacketAttributes{ + Buffer: b, + BytesCopied: p.BytesCopied, + } +} + +// Returns the read buffer. Just like when calling a read on a socket we have n, err := conn.Read(buf) +// and the read bytes are in buf[:n], we should use this method after calling the ReadWithAttributes method. +func (p *PacketAttributes) GetReadPacketAttributes() *PacketAttributes { + return &PacketAttributes{ + Buffer: p.Buffer[:p.BytesCopied], + BytesCopied: p.BytesCopied, + } +} diff --git a/packetio/buffer.go b/packetio/buffer.go index 3ac68d5f..9941b7ff 100644 --- a/packetio/buffer.go +++ b/packetio/buffer.go @@ -10,10 +10,12 @@ import ( "sync" "time" + "github.com/pion/transport/v3" "github.com/pion/transport/v3/deadline" ) var errPacketTooBig = errors.New("packet too big") +var errPacketAttributesBufferTooShort = errors.New("packet attributes buffer too short") // BufferPacketType allow the Buffer to know which packet protocol is writing. type BufferPacketType int @@ -61,15 +63,16 @@ func NewBuffer() *Buffer { } } -// available returns true if the buffer is large enough to fit a packet -// of the given size, taking overhead into account. +// available returns true if the buffer is large enough to fit a buffer +// of the given size, not taking overhead into account. func (b *Buffer) available(size int) bool { available := b.head - b.tail if available <= 0 { available += len(b.data) } // we interpret head=tail as empty, so always keep a byte free - if size+2+1 > available { + // the method + if size+1 > available { return false } @@ -119,12 +122,70 @@ func (b *Buffer) grow() error { return nil } +func (b *Buffer) consumeByte() byte { + byteRead := b.data[b.head] + b.head++ + if b.head >= len(b.data) { + b.head = 0 + } + + return byteRead +} + +func (b *Buffer) writeByte(x byte) { + b.data[b.tail] = x + b.tail++ + if b.tail >= len(b.data) { + b.tail = 0 + } +} + +func (b *Buffer) writeLengthHeaders(length int) { + n1 := uint8(length >> 8) //nolint:gosec + n2 := uint8(length) //nolint:gosec + b.writeByte(n1) + b.writeByte(n2) +} + +func (b *Buffer) getSegmentLength() int { + n1 := b.consumeByte() + n2 := b.consumeByte() + l := int((uint16(n1) << 8) | uint16(n2)) + + return l +} + +// The caller should make sure the input buffer (in []byte) +// can be completely accomodated in b.data . +func (b *Buffer) writeFromInputBuffer(in []byte) { + n := copy(b.data[b.tail:], in) + b.tail += n + if b.tail >= len(b.data) { + // we reached the end, wrap around + m := copy(b.data, in[n:]) + b.tail = m + } +} + // Write appends a copy of the packet data to the buffer. // Returns ErrFull if the packet doesn't fit. // // Note that the packet size is limited to 65536 bytes since v0.11.0 due to the internal data structure. -func (b *Buffer) Write(packet []byte) (int, error) { //nolint:cyclop - if len(packet) >= 0x10000 { +func (b *Buffer) Write(buff []byte) (n int, err error) { //nolint:cyclop + return b.WriteWithAttributes(buff, nil) +} + +// WriteWithAttributes works like Write, but it writes the +// additional packet attributes into the Buffer. +func (b *Buffer) WriteWithAttributes(packet []byte, + attr *transport.PacketAttributes) (int, error) { //nolint:cyclop + pLen := len(packet) + aLen := 0 + if attr != nil { + aLen = len(attr.Buffer) + } + + if pLen >= 0x10000 { return 0, errPacketTooBig } @@ -136,15 +197,17 @@ func (b *Buffer) Write(packet []byte) (int, error) { //nolint:cyclop return 0, io.ErrClosedPipe } + // two header bytes per segment indicating the length of the stored segment + lenWithHeaders := (2 + pLen) + (2 + aLen) if (b.limitCount > 0 && b.count >= b.limitCount) || - (b.limitSize > 0 && b.size()+2+len(packet) > b.limitSize) { + (b.limitSize > 0 && b.size()+lenWithHeaders > b.limitSize) { b.mutex.Unlock() return 0, ErrFull } // grow the buffer until the packet fits - for !b.available(len(packet)) { + for !b.available(lenWithHeaders) { err := b.grow() if err != nil { b.mutex.Unlock() @@ -154,41 +217,61 @@ func (b *Buffer) Write(packet []byte) (int, error) { //nolint:cyclop } // store the length of the packet - b.data[b.tail] = uint8(len(packet) >> 8) //nolint:gosec - b.tail++ - if b.tail >= len(b.data) { - b.tail = 0 - } - b.data[b.tail] = uint8(len(packet)) //nolint:gosec - b.tail++ - if b.tail >= len(b.data) { - b.tail = 0 - } + b.writeLengthHeaders(pLen) // store the packet - n := copy(b.data[b.tail:], packet) - b.tail += n - if b.tail >= len(b.data) { - // we reached the end, wrap around - m := copy(b.data, packet[n:]) - b.tail = m - } + b.writeFromInputBuffer(packet) + + // increment the number of packets in buffer b.count++ + // store the length of the attributes segment + b.writeLengthHeaders(aLen) + + if aLen > 0 { + // store the attributes buffer itself + b.writeFromInputBuffer(attr.Buffer) + } + select { case b.notify <- struct{}{}: default: } b.mutex.Unlock() - return len(packet), nil + return pLen, nil +} + +func (b *Buffer) writeToInputBuffer(in []byte, length int) { + if b.head+length < len(b.data) { + copy(in, b.data[b.head:b.head+length]) + } else { + k := copy(in, b.data[b.head:]) + copy(in[k:], b.data[:length-k]) + } +} + +func (b *Buffer) advanceHead(count int) { + b.head += count + if b.head >= len(b.data) { + b.head -= len(b.data) + } } // Read populates the given byte slice, returning the number of bytes read. // Blocks until data is available or the buffer is closed. // Returns io.ErrShortBuffer is the packet is too small to copy the Write. // Returns io.EOF if the buffer is closed. -func (b *Buffer) Read(packet []byte) (n int, err error) { //nolint:gocognit,cyclop +func (b *Buffer) Read(buff []byte) (n int, err error) { //nolint:gocognit,cyclop + n, err = b.ReadWithAttributes(buff, nil) + + return +} + +// ReadWithAttributes works like Read, but it also populates +// additional packet attributes into the attr field. +func (b *Buffer) ReadWithAttributes( + packet []byte, attr *transport.PacketAttributes) (n int, err error) { //nolint:gocognit,cyclop // Return immediately if the deadline is already exceeded. select { case <-b.readDeadline.Done(): @@ -201,17 +284,7 @@ func (b *Buffer) Read(packet []byte) (n int, err error) { //nolint:gocognit,cycl if b.head != b.tail { //nolint:nestif // decode the packet size - n1 := b.data[b.head] - b.head++ - if b.head >= len(b.data) { - b.head = 0 - } - n2 := b.data[b.head] - b.head++ - if b.head >= len(b.data) { - b.head = 0 - } - count := int((uint16(n1) << 8) | uint16(n2)) + count := b.getSegmentLength() // determine the number of bytes we'll actually copy copied := count @@ -220,18 +293,13 @@ func (b *Buffer) Read(packet []byte) (n int, err error) { //nolint:gocognit,cycl } // copy the data - if b.head+copied < len(b.data) { - copy(packet, b.data[b.head:b.head+copied]) - } else { - k := copy(packet, b.data[b.head:]) - copy(packet[k:], b.data[:copied-k]) - } + b.writeToInputBuffer(packet, copied) // advance head, discarding any data that wasn't copied - b.head += count - if b.head >= len(b.data) { - b.head -= len(b.data) - } + b.advanceHead(count) + + // read the attributes segment + attrErr := b.readPacketAttributes(attr) if b.head == b.tail { // the buffer is empty, reset to beginning @@ -244,9 +312,17 @@ func (b *Buffer) Read(packet []byte) (n int, err error) { //nolint:gocognit,cycl b.mutex.Unlock() if copied < count { + // The method still consumes the buffer even in the cases where + // packet buffer is short. but even in this case copies into the packet buffer + // and discards it. return copied, io.ErrShortBuffer } + // only if attr != nil do we care about attributes' buffer + if attr != nil && attrErr != nil { + return copied, attrErr + } + return copied, nil } @@ -265,6 +341,27 @@ func (b *Buffer) Read(packet []byte) (n int, err error) { //nolint:gocognit,cycl } } +func (b *Buffer) readPacketAttributes(attr *transport.PacketAttributes) error { + aLen := b.getSegmentLength() + if aLen == 0 { + return nil + } + if attr == nil { + b.advanceHead(aLen) + return nil + } + aBuffer := attr.Buffer + b.writeToInputBuffer(aBuffer, aLen) + b.advanceHead(aLen) + if len(aBuffer) >= aLen { + attr.BytesCopied = aLen + return nil + } + + attr.BytesCopied = len(aBuffer) + return errPacketAttributesBufferTooShort +} + // Close the buffer, unblocking any pending reads. // Data in the buffer can still be read, Read will return io.EOF only when empty. func (b *Buffer) Close() (err error) { diff --git a/packetio/buffer_test.go b/packetio/buffer_test.go index e35af368..4548585a 100644 --- a/packetio/buffer_test.go +++ b/packetio/buffer_test.go @@ -4,6 +4,7 @@ package packetio import ( + "crypto/rand" "errors" "fmt" "io" @@ -12,6 +13,7 @@ import ( "testing" "time" + "github.com/pion/transport/v3" "github.com/pion/transport/v3/test" "github.com/stretchr/testify/assert" ) @@ -257,7 +259,7 @@ func TestBufferLimitSize(t *testing.T) { assert := assert.New(t) buffer := NewBuffer() - buffer.SetLimitSize(11) + buffer.SetLimitSize(17) assert.Equal(0, buffer.Size()) @@ -265,23 +267,23 @@ func TestBufferLimitSize(t *testing.T) { n, err := buffer.Write([]byte{0, 1}) assert.NoError(err) assert.Equal(2, n) - assert.Equal(4, buffer.Size()) + assert.Equal(6, buffer.Size()) n, err = buffer.Write([]byte{2, 3}) assert.NoError(err) assert.Equal(2, n) - assert.Equal(8, buffer.Size()) + assert.Equal(12, buffer.Size()) // Over capacity _, err = buffer.Write([]byte{4, 5}) assert.Equal(ErrFull, err) - assert.Equal(8, buffer.Size()) + assert.Equal(12, buffer.Size()) // Cheeky write at exact size. n, err = buffer.Write([]byte{6}) assert.NoError(err) assert.Equal(1, n) - assert.Equal(11, buffer.Size()) + assert.Equal(17, buffer.Size()) // Read once packet := make([]byte, 4) @@ -289,31 +291,31 @@ func TestBufferLimitSize(t *testing.T) { assert.NoError(err) assert.Equal(2, n) assert.Equal([]byte{0, 1}, packet[:n]) - assert.Equal(7, buffer.Size()) + assert.Equal(11, buffer.Size()) // Write once n, err = buffer.Write([]byte{7, 8}) assert.NoError(err) assert.Equal(2, n) - assert.Equal(11, buffer.Size()) + assert.Equal(17, buffer.Size()) // Over capacity _, err = buffer.Write([]byte{9, 10}) assert.Equal(ErrFull, err) - assert.Equal(11, buffer.Size()) + assert.Equal(17, buffer.Size()) // Read everything n, err = buffer.Read(packet) assert.NoError(err) assert.Equal(2, n) assert.Equal([]byte{2, 3}, packet[:n]) - assert.Equal(7, buffer.Size()) + assert.Equal(11, buffer.Size()) n, err = buffer.Read(packet) assert.NoError(err) assert.Equal(1, n) assert.Equal([]byte{6}, packet[:n]) - assert.Equal(4, buffer.Size()) + assert.Equal(6, buffer.Size()) n, err = buffer.Read(packet) assert.NoError(err) @@ -717,3 +719,183 @@ func TestBufferConcurrentReadWrite(t *testing.T) { assert.NoError(buffer.Close()) } + +func newPacketAttributes(length int) *transport.PacketAttributes { + buff := make([]byte, length) + return &transport.PacketAttributes{ + Buffer: buff, + // EmptyBuffer: false, + // BytesCopied: length, + } +} + +func TestMixedReadWithAttributes(t *testing.T) { + assert := assert.New(t) + + buffer := NewBuffer() + + // prepare packets + p1 := make([]byte, 50) + p2 := make([]byte, 30) + p3 := make([]byte, 20) + _, err := rand.Read(p1) + assert.NoError(err) + _, err = rand.Read(p2) + assert.NoError(err) + _, err = rand.Read(p3) + assert.NoError(err) + + // prepare attributes for packets written with attributes + pa1 := newPacketAttributes(4) + pa3 := newPacketAttributes(2) + _, err = rand.Read(pa1.Buffer) + assert.NoError(err) + _, err = rand.Read(pa3.Buffer) + assert.NoError(err) + + // Write: with-attr, without-attr, with-attr + _, err = buffer.WriteWithAttributes(p1, pa1) + assert.NoError(err) + _, err = buffer.Write(p2) + assert.NoError(err) + _, err = buffer.WriteWithAttributes(p3, pa3) + assert.NoError(err) + + out := make([]byte, 1024) + + // 1) Read first packet using ReadWithAttributes to capture attributes + ra1 := newPacketAttributes(len(pa1.Buffer)) + var rattr1 *transport.PacketAttributes = ra1 + n, err := buffer.ReadWithAttributes(out, rattr1) + assert.NoError(err) + assert.Equal(len(p1), n) + assert.Equal(p1, out[:n]) + assert.Equal(pa1.Buffer, ra1.Buffer) + + // 2) Read second packet using ReadWithAttributes -> was written without attributes, + // attributes length should be zero and the provided buffer must remain unchanged. + ra2 := newPacketAttributes(3) + // pre-fill with sentinel + for i := range ra2.Buffer { + ra2.Buffer[i] = 0xFF + } + n, err = buffer.ReadWithAttributes(out, ra2) + assert.NoError(err) + assert.Equal(len(p2), n) + assert.Equal(p2, out[:n]) + // Ensure attributes buffer wasn't modified (aLen == 0 on write) + for _, v := range ra2.Buffer { + assert.Equal(byte(0xFF), v) + } + + // 3) Read third packet using ReadWithAttributes -> should return p3 and its attributes intact + ra3 := newPacketAttributes(len(pa3.Buffer)) + n, err = buffer.ReadWithAttributes(out, ra3) + assert.NoError(err) + assert.Equal(len(p3), n) + assert.Equal(p3, out[:n]) + assert.Equal(pa3.Buffer, ra3.Buffer) +} + +func TestShortPacketAttributesBuffer(t *testing.T) { + assert := assert.New(t) + + buffer := NewBuffer() + + // prepare packets + p1 := make([]byte, 32) + p2 := make([]byte, 16) + _, err := rand.Read(p1) + assert.NoError(err) + _, err = rand.Read(p2) + assert.NoError(err) + + // prepare attributes for packets + pa1 := newPacketAttributes(10) + pa2 := newPacketAttributes(2) + _, err = rand.Read(pa1.Buffer) + assert.NoError(err) + _, err = rand.Read(pa2.Buffer) + assert.NoError(err) + + // write two packets with attributes + _, err = buffer.WriteWithAttributes(p1, pa1) + assert.NoError(err) + _, err = buffer.WriteWithAttributes(p2, pa2) + assert.NoError(err) + + out := make([]byte, 1024) + + // Read first packet with too-small attributes buffer -> expect io.ErrShortBuffer + ra1 := newPacketAttributes(4) + // pre-fill with sentinel + for i := range ra1.Buffer { + ra1.Buffer[i] = 0x00 + } + n, err := buffer.ReadWithAttributes(out, ra1) + assert.Equal(len(p1), n) + assert.Equal(errPacketAttributesBufferTooShort, err) + // Ensure attributes buffer contains the leading bytes that fit + assert.Equal(pa1.Buffer[:len(ra1.Buffer)], ra1.Buffer) + // BytesCopied should reflect how many bytes were copied into the provided buffer + assert.Equal(len(ra1.Buffer), ra1.BytesCopied) + + // Read second packet with sufficiently large attributes buffer -> should succeed + ra2 := newPacketAttributes(4) + n, err = buffer.ReadWithAttributes(out, ra2) + assert.NoError(err) + assert.Equal(len(p2), n) + assert.Equal(pa2.Buffer, ra2.Buffer[:len(pa2.Buffer)]) + assert.Equal(len(pa2.Buffer), ra2.BytesCopied) +} + +func TestNilPacketAttributesWithReadAndWriteMixed(t *testing.T) { + assert := assert.New(t) + + buffer := NewBuffer() + + // prepare packets + p1 := make([]byte, 50) + p2 := make([]byte, 30) + p3 := make([]byte, 20) + _, err := rand.Read(p1) + assert.NoError(err) + _, err = rand.Read(p2) + assert.NoError(err) + _, err = rand.Read(p3) + assert.NoError(err) + + // prepare attributes for packets + pa1 := newPacketAttributes(4) + pa3 := newPacketAttributes(2) + _, err = rand.Read(pa1.Buffer) + assert.NoError(err) + _, err = rand.Read(pa3.Buffer) + assert.NoError(err) + + // Write: with-attr, without-attr, with-attr + _, err = buffer.WriteWithAttributes(p1, pa1) + assert.NoError(err) + _, err = buffer.Write(p2) + assert.NoError(err) + _, err = buffer.WriteWithAttributes(p3, pa3) + assert.NoError(err) + + out := make([]byte, 1024) + + // Simply Read (no attributes) and ensure payloads are returned intact. + n, err := buffer.Read(out) + assert.NoError(err) + assert.Equal(len(p1), n) + assert.Equal(p1, out[:n]) + + n, err = buffer.Read(out) + assert.NoError(err) + assert.Equal(len(p2), n) + assert.Equal(p2, out[:n]) + + n, err = buffer.Read(out) + assert.NoError(err) + assert.Equal(len(p3), n) + assert.Equal(p3, out[:n]) +}