Add Peek to ReadStreamSRTP

This will be used while probing for PayloadType and Simulcast.

Relates to pion/webrtc#2777

Co-authored-by: cptpcrd <[email protected]>

Author: Sean-Der

stream_srtp.go

@@ -6,6 +6,7 @@ package srtp
 import (
 	"errors"
 	"io"
+	"slices"
 	"sync"
 	"time"
 
@@ -26,7 +27,8 @@ type ReadStreamSRTP struct {
 	ssrc     uint32
 	isInited bool
 
-	buffer io.ReadWriteCloser
+	buffer        io.ReadWriteCloser
+	peekedPackets [][]byte
 }
 
 // Used by getOrCreateReadStream.
@@ -74,8 +76,31 @@ func (r *ReadStreamSRTP) write(buf []byte) (n int, err error) {
 	return n, err
 }
 
+// Peek reads and decrypts full RTP packet from the nextConn.
+// It is then buffered so that a call to `Read` will return it.
+func (r *ReadStreamSRTP) Peek(buf []byte) (n int, err error) {
+	n, err = r.buffer.Read(buf)
+	if err == nil {
+		r.peekedPackets = append(r.peekedPackets, slices.Clone(buf[:n]))
+	}
+
+	return
+}
+
 // Read reads and decrypts full RTP packet from the nextConn.
 func (r *ReadStreamSRTP) Read(buf []byte) (int, error) {
+	if len(r.peekedPackets) != 0 {
+		if len(r.peekedPackets[0]) > len(buf) {
+			return 0, io.ErrShortBuffer
+		}
+
+		n := len(r.peekedPackets[0])
+		copy(buf, r.peekedPackets[0])
+		r.peekedPackets = r.peekedPackets[1:]
+
+		return n, nil
+	}
+
 	return r.buffer.Read(buf)
 }
 

stream_srtp_test.go

@@ -35,6 +35,105 @@ func (c *noopConn) SetDeadline(time.Time) error      { return nil }
 func (c *noopConn) SetReadDeadline(time.Time) error  { return nil }
 func (c *noopConn) SetWriteDeadline(time.Time) error { return nil }
 
+func TestPeek(t *testing.T) {
+	firstBuffer := []byte{0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA}
+	secondBuffer := []byte{0xBB, 0xBB, 0xBB}
+	thirdBuffer := []byte{0xCC, 0xCC, 0xCC}
+
+	buffer := packetio.NewBuffer()
+	stream := &ReadStreamSRTP{buffer: buffer}
+
+	t.Run("Short Peek", func(t *testing.T) {
+		_, err := buffer.Write(firstBuffer)
+		assert.NoError(t, err)
+
+		readBuff := make([]byte, 1)
+		_, err = stream.Peek(readBuff)
+		assert.Error(t, err, io.ErrShortBuffer)
+	})
+
+	t.Run("Short Read", func(t *testing.T) {
+		_, err := buffer.Write(firstBuffer)
+		assert.NoError(t, err)
+
+		readBuff := make([]byte, 6)
+		n, err := stream.Peek(readBuff)
+		assert.NoError(t, err)
+		assert.Equal(t, n, 6)
+		assert.Equal(t, readBuff, firstBuffer)
+
+		n, err = stream.Read([]byte{})
+		assert.Error(t, err, io.ErrShortBuffer)
+		assert.Equal(t, n, 0)
+		assert.Equal(t, readBuff, firstBuffer)
+
+		n, err = stream.Read(readBuff)
+		assert.NoError(t, err)
+		assert.Equal(t, n, 6)
+		assert.Equal(t, readBuff, firstBuffer)
+	})
+
+	t.Run("Single Peek", func(t *testing.T) {
+		_, err := buffer.Write(firstBuffer)
+		assert.NoError(t, err)
+
+		readBuff := make([]byte, 6)
+
+		n, err := stream.Peek(readBuff)
+		assert.NoError(t, err)
+		assert.Equal(t, n, 6)
+		assert.Equal(t, readBuff, firstBuffer)
+
+		n, err = stream.Read(readBuff)
+		assert.NoError(t, err)
+		assert.Equal(t, n, 6)
+		assert.Equal(t, readBuff, firstBuffer)
+	})
+
+	t.Run("Multi Peek", func(t *testing.T) {
+		_, err := buffer.Write(firstBuffer)
+		assert.NoError(t, err)
+
+		_, err = buffer.Write(secondBuffer)
+		assert.NoError(t, err)
+
+		_, err = buffer.Write(thirdBuffer)
+		assert.NoError(t, err)
+
+		readBuff := make([]byte, 6)
+
+		n, err := stream.Peek(readBuff)
+		assert.NoError(t, err)
+		assert.Equal(t, n, 6)
+		assert.Equal(t, readBuff[:n], firstBuffer)
+
+		n, err = stream.Peek(readBuff)
+		assert.NoError(t, err)
+		assert.Equal(t, n, 3)
+		assert.Equal(t, readBuff[:n], secondBuffer)
+
+		n, err = stream.Peek(readBuff)
+		assert.NoError(t, err)
+		assert.Equal(t, n, 3)
+		assert.Equal(t, readBuff[:n], thirdBuffer)
+
+		n, err = stream.Read(readBuff)
+		assert.NoError(t, err)
+		assert.Equal(t, n, 6)
+		assert.Equal(t, readBuff[:n], firstBuffer)
+
+		n, err = stream.Read(readBuff)
+		assert.NoError(t, err)
+		assert.Equal(t, n, 3)
+		assert.Equal(t, readBuff[:n], secondBuffer)
+
+		n, err = stream.Read(readBuff)
+		assert.NoError(t, err)
+		assert.Equal(t, n, 3)
+		assert.Equal(t, readBuff[:n], thirdBuffer)
+	})
+}
+
 func TestBufferFactory(t *testing.T) {
 	wg := sync.WaitGroup{}
 	wg.Add(2)