Add test for PR94

Author: ioppermann

connection_test.go

@@ -6,6 +6,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/datarhei/gosrt/packet"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -146,6 +147,159 @@ func TestEncryption(t *testing.T) {
 	require.Equal(t, message, reader1)
 }
 
+// Test for https://github.com/datarhei/gosrt/pull/94
+func TestEncryptionRetransmit(t *testing.T) {
+	message := "Hello World!"
+	passphrase := "foobarfoobar"
+	channel := NewPubSub(PubSubConfig{})
+
+	config := DefaultConfig()
+	config.EnforcedEncryption = true
+
+	server := Server{
+		Addr:   "127.0.0.1:6003",
+		Config: &config,
+		HandleConnect: func(req ConnRequest) ConnType {
+			if req.IsEncrypted() {
+				if err := req.SetPassphrase(passphrase); err != nil {
+					return REJECT
+				}
+			}
+
+			streamid := req.StreamId()
+
+			if streamid == "publish" {
+				return PUBLISH
+			} else if streamid == "subscribe" {
+				return SUBSCRIBE
+			}
+
+			return REJECT
+		},
+		HandlePublish: func(conn Conn) {
+			channel.Publish(conn)
+
+			conn.Close()
+		},
+		HandleSubscribe: func(conn Conn) {
+			channel.Subscribe(conn)
+
+			conn.Close()
+		},
+	}
+
+	err := server.Listen()
+	require.NoError(t, err)
+
+	defer server.Shutdown()
+
+	go func() {
+		err := server.Serve()
+		if err == ErrServerClosed {
+			return
+		}
+		require.NoError(t, err)
+	}()
+
+	{
+		// Reject connection if wrong password is set
+		config := DefaultConfig()
+		config.StreamId = "subscribe"
+		config.Passphrase = "barfoobarfoo"
+
+		_, err := Dial("srt", "127.0.0.1:6003", config)
+		require.Error(t, err)
+	}
+
+	// Test transmitting an encrypted message
+
+	readerConnected := make(chan struct{})
+	readerDone := make(chan struct{})
+
+	dataReader1 := bytes.Buffer{}
+
+	go func() {
+		defer close(readerDone)
+
+		config := DefaultConfig()
+		config.StreamId = "subscribe"
+		config.Passphrase = "foobarfoobar"
+
+		conn, err := Dial("srt", "127.0.0.1:6003", config)
+		if !assert.NoError(t, err) {
+			panic(err.Error())
+		}
+
+		close(readerConnected)
+
+		buffer := make([]byte, 2048)
+
+		for {
+			n, err := conn.Read(buffer)
+			if n != 0 {
+				dataReader1.Write(buffer[:n])
+			}
+
+			if err != nil {
+				break
+			}
+		}
+
+		err = conn.Close()
+		require.NoError(t, err)
+	}()
+
+	<-readerConnected
+
+	writerDone := make(chan struct{})
+
+	go func() {
+		defer close(writerDone)
+
+		config := DefaultConfig()
+		config.StreamId = "publish"
+		config.Passphrase = "foobarfoobar"
+
+		conn, err := Dial("srt", "127.0.0.1:6003", config)
+		if !assert.NoError(t, err) {
+			panic(err.Error())
+		}
+
+		dialer, _ := conn.(*dialer)
+		originalOnSend := dialer.conn.onSend
+		dialer.conn.onSend = func(p packet.Packet) {
+			if !p.Header().IsControlPacket {
+				// Drop every 2nd original packet
+				if !p.Header().RetransmittedPacketFlag && p.Header().PacketSequenceNumber.Val()%2 == 1 {
+					return
+				}
+			}
+
+			originalOnSend(p)
+		}
+
+		for i := 0; i < 5; i++ {
+			n, err := conn.Write([]byte(message))
+			if !assert.NoError(t, err) {
+				panic(err.Error())
+			}
+			assert.Equal(t, 12, n)
+		}
+
+		time.Sleep(3 * time.Second)
+
+		err = conn.Close()
+		assert.NoError(t, err)
+	}()
+
+	<-writerDone
+	<-readerDone
+
+	reader1 := dataReader1.String()
+
+	require.Equal(t, message+message+message+message+message, reader1)
+}
+
 func TestEncryptionKeySwap(t *testing.T) {
 	message := "Hello World!"
 	passphrase := "foobarfoobar"