Fix WHIP reconnect logic

Old WHIP sessions would incorrectly cause new ones to get deleted. This
adds ID to WHIP sessions so we can properly identify them when cleaning
up.

Resolves #296

Author: Criveratech

go.mod

@@ -7,30 +7,34 @@ require (
 	github.com/joho/godotenv v1.5.1
 	github.com/pion/dtls/v3 v3.0.6
 	github.com/pion/ice/v3 v3.0.16
-	github.com/pion/interceptor v0.1.37
+	github.com/pion/interceptor v0.1.40
 	github.com/pion/rtcp v1.2.15
-	github.com/pion/rtp v1.8.17
+	github.com/pion/rtp v1.8.18
 	github.com/pion/sdp/v3 v3.0.13
-	github.com/pion/webrtc/v4 v4.1.1
+	github.com/pion/webrtc/v4 v4.1.2
+	github.com/stretchr/testify v1.10.0
 )
 
 require (
+	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/pion/datachannel v1.5.10 // indirect
 	github.com/pion/dtls/v2 v2.2.12 // indirect
 	github.com/pion/ice/v4 v4.0.10 // indirect
 	github.com/pion/logging v0.2.3 // indirect
 	github.com/pion/mdns/v2 v2.0.7 // indirect
 	github.com/pion/randutil v0.1.0 // indirect
 	github.com/pion/sctp v1.8.39 // indirect
-	github.com/pion/srtp/v3 v3.0.4 // indirect
+	github.com/pion/srtp/v3 v3.0.5 // indirect
 	github.com/pion/stun/v2 v2.0.0 // indirect
 	github.com/pion/stun/v3 v3.0.0 // indirect
 	github.com/pion/transport/v2 v2.2.10 // indirect
 	github.com/pion/transport/v3 v3.0.7 // indirect
 	github.com/pion/turn/v3 v3.0.3 // indirect
 	github.com/pion/turn/v4 v4.0.0 // indirect
+	github.com/pmezard/go-difflib v1.0.0 // indirect
 	github.com/wlynxg/anet v0.0.5 // indirect
 	golang.org/x/crypto v0.33.0 // indirect
 	golang.org/x/net v0.35.0 // indirect
 	golang.org/x/sys v0.30.0 // indirect
+	gopkg.in/yaml.v3 v3.0.1 // indirect
 )

go.sum

@@ -5,6 +5,8 @@ github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
 github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
 github.com/joho/godotenv v1.5.1 h1:7eLL/+HRGLY0ldzfGMeQkb7vMd0as4CfYvUVzLqw0N0=
 github.com/joho/godotenv v1.5.1/go.mod h1:f4LDr5Voq0i2e/R5DDNOoa2zzDfwtkZa6DnEwAbqwq4=
+github.com/kr/pretty v0.1.0 h1:L/CwN0zerZDmRFUapSPitk6f+Q3+0za1rQkzVuMiMFI=
+github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
 github.com/pion/datachannel v1.5.10 h1:ly0Q26K1i6ZkGf42W7D4hQYR90pZwzFOjTq5AuCKk4o=
 github.com/pion/datachannel v1.5.10/go.mod h1:p/jJfC9arb29W7WrxyKbepTU20CFgyx5oLo8Rs4Py/M=
 github.com/pion/dtls/v2 v2.2.7/go.mod h1:8WiMkebSHFD0T+dIU+UeBaoV7kDhOW5oDCzZ7WZ/F9s=
@@ -16,8 +18,8 @@ github.com/pion/ice/v3 v3.0.16 h1:YoPlNg3jU1UT/DDTa9v/g1vH6A2/pAzehevI1o66H8E=
 github.com/pion/ice/v3 v3.0.16/go.mod h1:SdmubtIsCcvdb1ZInrTUz7Iaqi90/rYd1pzbzlMxsZg=
 github.com/pion/ice/v4 v4.0.10 h1:P59w1iauC/wPk9PdY8Vjl4fOFL5B+USq1+xbDcN6gT4=
 github.com/pion/ice/v4 v4.0.10/go.mod h1:y3M18aPhIxLlcO/4dn9X8LzLLSma84cx6emMSu14FGw=
-github.com/pion/interceptor v0.1.37 h1:aRA8Zpab/wE7/c0O3fh1PqY0AJI3fCSEM5lRWJVorwI=
-github.com/pion/interceptor v0.1.37/go.mod h1:JzxbJ4umVTlZAf+/utHzNesY8tmRkM2lVmkS82TTj8Y=
+github.com/pion/interceptor v0.1.40 h1:e0BjnPcGpr2CFQgKhrQisBU7V3GXK6wrfYrGYaU6Jq4=
+github.com/pion/interceptor v0.1.40/go.mod h1:Z6kqH7M/FYirg3frjGJ21VLSRJGBXB/KqaTIrdqnOic=
 github.com/pion/logging v0.2.2/go.mod h1:k0/tDVsRCX2Mb2ZEmTqNa7CWsQPc+YYCB7Q+5pahoms=
 github.com/pion/logging v0.2.3 h1:gHuf0zpoh1GW67Nr6Gj4cv5Z9ZscU7g/EaoC/Ke/igI=
 github.com/pion/logging v0.2.3/go.mod h1:z8YfknkquMe1csOrxK5kc+5/ZPAzMxbKLX5aXpbpC90=
@@ -27,14 +29,14 @@ github.com/pion/randutil v0.1.0 h1:CFG1UdESneORglEsnimhUjf33Rwjubwj6xfiOXBa3mA=
 github.com/pion/randutil v0.1.0/go.mod h1:XcJrSMMbbMRhASFVOlj/5hQial/Y8oH/HVo7TBZq+j8=
 github.com/pion/rtcp v1.2.15 h1:LZQi2JbdipLOj4eBjK4wlVoQWfrZbh3Q6eHtWtJBZBo=
 github.com/pion/rtcp v1.2.15/go.mod h1:jlGuAjHMEXwMUHK78RgX0UmEJFV4zUKOFHR7OP+D3D0=
-github.com/pion/rtp v1.8.17 h1:CFhaPN8Ikt9Sk7B3pic0kfwVia2dUMEtPSL34Gvihjw=
-github.com/pion/rtp v1.8.17/go.mod h1:bAu2UFKScgzyFqvUKmbvzSdPr+NGbZtv6UB2hesqXBk=
+github.com/pion/rtp v1.8.18 h1:yEAb4+4a8nkPCecWzQB6V/uEU18X1lQCGAQCjP+pyvU=
+github.com/pion/rtp v1.8.18/go.mod h1:bAu2UFKScgzyFqvUKmbvzSdPr+NGbZtv6UB2hesqXBk=
 github.com/pion/sctp v1.8.39 h1:PJma40vRHa3UTO3C4MyeJDQ+KIobVYRZQZ0Nt7SjQnE=
 github.com/pion/sctp v1.8.39/go.mod h1:cNiLdchXra8fHQwmIoqw0MbLLMs+f7uQ+dGMG2gWebE=
 github.com/pion/sdp/v3 v3.0.13 h1:uN3SS2b+QDZnWXgdr69SM8KB4EbcnPnPf2Laxhty/l4=
 github.com/pion/sdp/v3 v3.0.13/go.mod h1:88GMahN5xnScv1hIMTqLdu/cOcUkj6a9ytbncwMCq2E=
-github.com/pion/srtp/v3 v3.0.4 h1:2Z6vDVxzrX3UHEgrUyIGM4rRouoC7v+NiF1IHtp9B5M=
-github.com/pion/srtp/v3 v3.0.4/go.mod h1:1Jx3FwDoxpRaTh1oRV8A/6G1BnFL+QI82eK4ms8EEJQ=
+github.com/pion/srtp/v3 v3.0.5 h1:8XLB6Dt3QXkMkRFpoqC3314BemkpMQK2mZeJc4pUKqo=
+github.com/pion/srtp/v3 v3.0.5/go.mod h1:r1G7y5r1scZRLe2QJI/is+/O83W2d+JoEsuIexpw+uM=
 github.com/pion/stun/v2 v2.0.0 h1:A5+wXKLAypxQri59+tmQKVs7+l6mMM+3d+eER9ifRU0=
 github.com/pion/stun/v2 v2.0.0/go.mod h1:22qRSh08fSEttYUmJZGlriq9+03jtVmXNODgLccj8GQ=
 github.com/pion/stun/v3 v3.0.0 h1:4h1gwhWLWuZWOJIJR9s2ferRO+W3zA/b6ijOI6mKzUw=
@@ -50,8 +52,8 @@ github.com/pion/turn/v3 v3.0.3 h1:1e3GVk8gHZLPBA5LqadWYV60lmaKUaHCkm9DX9CkGcE=
 github.com/pion/turn/v3 v3.0.3/go.mod h1:vw0Dz420q7VYAF3J4wJKzReLHIo2LGp4ev8nXQexYsc=
 github.com/pion/turn/v4 v4.0.0 h1:qxplo3Rxa9Yg1xXDxxH8xaqcyGUtbHYw4QSCvmFWvhM=
 github.com/pion/turn/v4 v4.0.0/go.mod h1:MuPDkm15nYSklKpN8vWJ9W2M0PlyQZqYt1McGuxG7mA=
-github.com/pion/webrtc/v4 v4.1.1 h1:PMFPtLg1kpD2pVtun+LGUzA3k54JdFl87WO0Z1+HKug=
-github.com/pion/webrtc/v4 v4.1.1/go.mod h1:cgEGkcpxGkT6Di2ClBYO5lP9mFXbCfEOrkYUpjjCQO4=
+github.com/pion/webrtc/v4 v4.1.2 h1:mpuUo/EJ1zMNKGE79fAdYNFZBX790KE7kQQpLMjjR54=
+github.com/pion/webrtc/v4 v4.1.2/go.mod h1:xsCXiNAmMEjIdFxAYU0MbB3RwRieJsegSB2JZsGN+8U=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
@@ -62,6 +64,7 @@ github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO
 github.com/stretchr/testify v1.8.3/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo=
 github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo=
 github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA=
+github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
 github.com/wlynxg/anet v0.0.3/go.mod h1:eay5PRQr7fIVAMbTbchTnO9gG65Hg/uYGdc7mguHxoA=
 github.com/wlynxg/anet v0.0.5 h1:J3VJGi1gvo0JwZ/P1/Yc/8p63SoW98B5dHkYDmpgvvU=
 github.com/wlynxg/anet v0.0.5/go.mod h1:eay5PRQr7fIVAMbTbchTnO9gG65Hg/uYGdc7mguHxoA=
@@ -120,6 +123,7 @@ golang.org/x/tools v0.1.12/go.mod h1:hNGJHUnrk76NpqgfD5Aqm5Crs+Hm0VOH/i9J2+nxYbc
 golang.org/x/tools v0.6.0/go.mod h1:Xwgl3UAJ/d3gWutnCtw505GrjyAbvKui8lOU390QaIU=
 golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15 h1:YR8cESwS4TdDjEe65xsg0ogRM/Nc3DYOhEAlW+xobZo=
 gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
 gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
 gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

internal/webrtc/webrtc.go

@@ -9,6 +9,7 @@ import (
 	"net"
 	"net/http"
 	"os"
+	"slices"
 	"strconv"
 	"strings"
 	"sync"
@@ -36,6 +37,7 @@ type (
 		// Does this stream have a publisher?
 		// If stream was created by a WHEP request hasWHIPClient == false
 		hasWHIPClient atomic.Bool
+		sessionId     string
 
 		firstSeenEpoch uint64
 
@@ -54,6 +56,7 @@ type (
 	}
 
 	videoTrack struct {
+		sessionId        string
 		rid              string
 		packetsReceived  atomic.Uint64
 		lastKeyFrameSeen atomic.Value
@@ -89,7 +92,7 @@ func getVideoTrackCodec(in string) videoTrackCodec {
 	return 0
 }
 
-func getStream(streamKey string, forWHIP bool) (*stream, error) {
+func getStream(streamKey string, whipSessionId string) (*stream, error) {
 	foundStream, ok := streamMap[streamKey]
 	if !ok {
 		audioTrack, err := webrtc.NewTrackLocalStaticRTP(webrtc.RTPCodecCapability{MimeType: webrtc.MimeTypeOpus}, "audio", "pion")
@@ -110,14 +113,15 @@ func getStream(streamKey string, forWHIP bool) (*stream, error) {
 		streamMap[streamKey] = foundStream
 	}
 
-	if forWHIP {
+	if whipSessionId != "" {
 		foundStream.hasWHIPClient.Store(true)
+		foundStream.sessionId = whipSessionId
 	}
 
 	return foundStream, nil
 }
 
-func peerConnectionDisconnected(streamKey string, whepSessionId string) {
+func peerConnectionDisconnected(forWHIP bool, streamKey string, sessionId string) {
 	streamMapLock.Lock()
 	defer streamMapLock.Unlock()
 
@@ -129,11 +133,20 @@ func peerConnectionDisconnected(streamKey string, whepSessionId string) {
 	stream.whepSessionsLock.Lock()
 	defer stream.whepSessionsLock.Unlock()
 
-	if whepSessionId != "" {
-		delete(stream.whepSessions, whepSessionId)
+	if !forWHIP {
+		delete(stream.whepSessions, sessionId)
 	} else {
+		stream.videoTracks = slices.DeleteFunc(stream.videoTracks, func(v *videoTrack) bool {
+			return v.sessionId == sessionId
+		})
+
+		// A PeerConnection for a old WHIP session has gone to disconnected
+		// closed. Cleanup the state associated with that session, but
+		// don't modify the current session
+		if stream.sessionId != sessionId {
+			return
+		}
 		stream.hasWHIPClient.Store(false)
-		stream.videoTracks = nil
 	}
 
 	// Only delete stream if all WHEP Sessions are gone and have no WHIP Client
@@ -145,7 +158,7 @@ func peerConnectionDisconnected(streamKey string, whepSessionId string) {
 	delete(streamMap, streamKey)
 }
 
-func addTrack(stream *stream, rid string) (*videoTrack, error) {
+func addTrack(stream *stream, rid, sessionId string) (*videoTrack, error) {
 	streamMapLock.Lock()
 	defer streamMapLock.Unlock()
 
@@ -155,7 +168,7 @@ func addTrack(stream *stream, rid string) (*videoTrack, error) {
 		}
 	}
 
-	t := &videoTrack{rid: rid}
+	t := &videoTrack{rid: rid, sessionId: sessionId}
 	t.lastKeyFrameSeen.Store(time.Time{})
 	stream.videoTracks = append(stream.videoTracks, t)
 	return t, nil

internal/webrtc/whep.go

@@ -78,7 +78,7 @@ func WHEP(offer, streamKey string) (string, string, error) {
 
 	streamMapLock.Lock()
 	defer streamMapLock.Unlock()
-	stream, err := getStream(streamKey, false)
+	stream, err := getStream(streamKey, "")
 	if err != nil {
 		return "", "", err
 	}
@@ -98,7 +98,7 @@ func WHEP(offer, streamKey string) (string, string, error) {
 				log.Println(err)
 			}
 
-			peerConnectionDisconnected(streamKey, whepSessionId)
+			peerConnectionDisconnected(false, streamKey, whepSessionId)
 		}
 	})
 

internal/webrtc/whip.go

@@ -8,6 +8,7 @@ import (
 	"strings"
 	"time"
 
+	"github.com/google/uuid"
 	"github.com/pion/rtcp"
 	"github.com/pion/rtp"
 	"github.com/pion/rtp/codecs"
@@ -34,13 +35,13 @@ func audioWriter(remoteTrack *webrtc.TrackRemote, stream *stream) {
 	}
 }
 
-func videoWriter(remoteTrack *webrtc.TrackRemote, stream *stream, peerConnection *webrtc.PeerConnection, s *stream) {
+func videoWriter(remoteTrack *webrtc.TrackRemote, stream *stream, peerConnection *webrtc.PeerConnection, s *stream, sessionId string) {
 	id := remoteTrack.RID()
 	if id == "" {
 		id = videoTrackLabelDefault
 	}
 
-	videoTrack, err := addTrack(s, id)
+	videoTrack, err := addTrack(s, id, sessionId)
 	if err != nil {
 		log.Println(err)
 		return
@@ -142,14 +143,16 @@ func videoWriter(remoteTrack *webrtc.TrackRemote, stream *stream, peerConnection
 func WHIP(offer, streamKey string) (string, error) {
 	maybePrintOfferAnswer(offer, true)
 
+	whipSessionId := uuid.New().String()
+
 	peerConnection, err := newPeerConnection(apiWhip)
 	if err != nil {
 		return "", err
 	}
 
 	streamMapLock.Lock()
 	defer streamMapLock.Unlock()
-	stream, err := getStream(streamKey, true)
+	stream, err := getStream(streamKey, whipSessionId)
 	if err != nil {
 		return "", err
 	}
@@ -158,7 +161,7 @@ func WHIP(offer, streamKey string) (string, error) {
 		if strings.HasPrefix(remoteTrack.Codec().MimeType, "audio") {
 			audioWriter(remoteTrack, stream)
 		} else {
-			videoWriter(remoteTrack, stream, peerConnection, stream)
+			videoWriter(remoteTrack, stream, peerConnection, stream, whipSessionId)
 
 		}
 	})
@@ -168,7 +171,7 @@ func WHIP(offer, streamKey string) (string, error) {
 			if err := peerConnection.Close(); err != nil {
 				log.Println(err)
 			}
-			peerConnectionDisconnected(streamKey, "")
+			peerConnectionDisconnected(true, streamKey, whipSessionId)
 		}
 	})
 

internal/webrtc/whip_test.go

@@ -0,0 +1,102 @@
+package webrtc
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/pion/webrtc/v4"
+	"github.com/stretchr/testify/require"
+)
+
+const testStreamKey = "test"
+
+func doesWHIPSessionExist() (ok bool) {
+	streamMapLock.Lock()
+	defer streamMapLock.Unlock()
+
+	_, ok = streamMap[testStreamKey]
+	return
+}
+
+// Asserts that a old PeerConnection doesn't destroy the new one
+// when it disconnects
+func TestReconnect(t *testing.T) {
+	Configure()
+	localTrack, err := webrtc.NewTrackLocalStaticSample(
+		webrtc.RTPCodecCapability{MimeType: webrtc.MimeTypeH264}, "video", "pion",
+	)
+	require.NoError(t, err)
+
+	// Create the first WHIP Session
+	firstPublisherConnected, firstPublisherConnectedDone := context.WithCancel(context.TODO())
+
+	firstPublisher, err := webrtc.NewPeerConnection(webrtc.Configuration{})
+	require.NoError(t, err)
+
+	firstPublisher.OnConnectionStateChange(func(c webrtc.PeerConnectionState) {
+		if c == webrtc.PeerConnectionStateConnected {
+			firstPublisherConnectedDone()
+
+		}
+	})
+
+	_, err = firstPublisher.AddTrack(localTrack)
+	require.NoError(t, err)
+
+	offer, err := firstPublisher.CreateOffer(nil)
+	require.NoError(t, err)
+	require.NoError(t, firstPublisher.SetLocalDescription(offer))
+
+	answer, err := WHIP(offer.SDP, testStreamKey)
+	require.NoError(t, err)
+
+	firstPublisher.SetRemoteDescription(webrtc.SessionDescription{
+		Type: webrtc.SDPTypeAnswer,
+		SDP:  answer,
+	})
+
+	require.True(t, doesWHIPSessionExist())
+	<-firstPublisherConnected.Done()
+
+	// Create the second WHIP Session
+	secondPublisherConnected, secondPublisherConnectedDone := context.WithCancel(context.TODO())
+
+	secondPublisher, err := webrtc.NewPeerConnection(webrtc.Configuration{})
+	require.NoError(t, err)
+
+	secondPublisher.OnConnectionStateChange(func(c webrtc.PeerConnectionState) {
+		if c == webrtc.PeerConnectionStateConnected {
+			secondPublisherConnectedDone()
+
+		}
+	})
+
+	_, err = secondPublisher.AddTrack(localTrack)
+	require.NoError(t, err)
+
+	offer, err = secondPublisher.CreateOffer(nil)
+	require.NoError(t, err)
+	require.NoError(t, secondPublisher.SetLocalDescription(offer))
+
+	answer, err = WHIP(offer.SDP, testStreamKey)
+	require.NoError(t, err)
+
+	secondPublisher.SetRemoteDescription(webrtc.SessionDescription{
+		Type: webrtc.SDPTypeAnswer,
+		SDP:  answer,
+	})
+
+	require.True(t, doesWHIPSessionExist())
+	<-secondPublisherConnected.Done()
+
+	// Close the first WHIP Session, the session must still exist
+	require.NoError(t, firstPublisher.Close())
+	time.Sleep(time.Second)
+	require.True(t, doesWHIPSessionExist())
+
+	// Close the second WHIP Session, the session must be gone
+	require.NoError(t, secondPublisher.Close())
+	time.Sleep(time.Second)
+	require.False(t, doesWHIPSessionExist())
+}