Fix writer deadlocks in nack

Author: philipch07

pkg/nack/generator_interceptor.go

@@ -156,66 +156,71 @@ func (n *GeneratorInterceptor) loop(rtcpWriter interceptor.RTCPWriter) {
 	for {
 		select {
 		case <-ticker.C:
-			func() {
-				n.receiveLogsMu.Lock()
-				defer n.receiveLogsMu.Unlock()
+			// save NACKs to send without holding the mutex during Write
+			var toSend []rtcp.Packet
 
-				for ssrc, receiveLog := range n.receiveLogs {
-					missing := receiveLog.missingSeqNumbers(n.skipLastN, missingPacketSeqNums)
+			n.receiveLogsMu.Lock()
+			for ssrc, receiveLog := range n.receiveLogs {
+				missing := receiveLog.missingSeqNumbers(n.skipLastN, missingPacketSeqNums)
 
-					if len(missing) == 0 || n.nackCountLogs[ssrc] == nil {
-						n.nackCountLogs[ssrc] = map[uint16]uint16{}
-					}
-					if len(missing) == 0 {
-						continue
-					}
-
-					nack := &rtcp.TransportLayerNack{} // nolint:ineffassign,wastedassign
-
-					c := 0 // nolint:varnamelen,
-					if n.maxNacksPerPacket > 0 {
-						for _, missingSeq := range missing {
-							if n.nackCountLogs[ssrc][missingSeq] < n.maxNacksPerPacket {
-								filteredMissingPacket[c] = missingSeq
-								c++
-							}
-							n.nackCountLogs[ssrc][missingSeq]++
-						}
+				if len(missing) == 0 || n.nackCountLogs[ssrc] == nil {
+					n.nackCountLogs[ssrc] = map[uint16]uint16{}
+				}
+				if len(missing) == 0 {
+					continue
+				}
 
-						if c == 0 {
-							continue
-						}
+				var nack *rtcp.TransportLayerNack
 
-						nack = &rtcp.TransportLayerNack{
-							SenderSSRC: senderSSRC,
-							MediaSSRC:  ssrc,
-							Nacks:      rtcp.NackPairsFromSequenceNumbers(filteredMissingPacket[:c]),
-						}
-					} else {
-						nack = &rtcp.TransportLayerNack{
-							SenderSSRC: senderSSRC,
-							MediaSSRC:  ssrc,
-							Nacks:      rtcp.NackPairsFromSequenceNumbers(missing),
+				count := 0
+				if n.maxNacksPerPacket > 0 {
+					for _, missingSeq := range missing {
+						if n.nackCountLogs[ssrc][missingSeq] < n.maxNacksPerPacket {
+							filteredMissingPacket[count] = missingSeq
+							count++
 						}
+						n.nackCountLogs[ssrc][missingSeq]++
 					}
 
-					for nackSeq := range n.nackCountLogs[ssrc] {
-						isMissing := slices.Contains(missing, nackSeq)
-						if !isMissing {
-							delete(n.nackCountLogs[ssrc], nackSeq)
-						}
+					if count == 0 {
+						continue
 					}
 
-					// clean up the count log for the ssrc if it's empty
-					if len(n.nackCountLogs[ssrc]) == 0 {
-						delete(n.nackCountLogs, ssrc)
+					nack = &rtcp.TransportLayerNack{
+						SenderSSRC: senderSSRC,
+						MediaSSRC:  ssrc,
+						Nacks:      rtcp.NackPairsFromSequenceNumbers(filteredMissingPacket[:count]),
+					}
+				} else {
+					nack = &rtcp.TransportLayerNack{
+						SenderSSRC: senderSSRC,
+						MediaSSRC:  ssrc,
+						Nacks:      rtcp.NackPairsFromSequenceNumbers(missing),
 					}
+				}
 
-					if _, err := rtcpWriter.Write([]rtcp.Packet{nack}, interceptor.Attributes{}); err != nil {
-						n.log.Warnf("failed sending nack: %+v", err)
+				for nackSeq := range n.nackCountLogs[ssrc] {
+					if !slices.Contains(missing, nackSeq) {
+						delete(n.nackCountLogs[ssrc], nackSeq)
 					}
 				}
-			}()
+
+				// clean up the count log for the ssrc if it's empty
+				if len(n.nackCountLogs[ssrc]) == 0 {
+					delete(n.nackCountLogs, ssrc)
+				}
+
+				toSend = append(toSend, nack)
+			}
+			n.receiveLogsMu.Unlock()
+
+			// send RTCP without holding receiveLogsMu
+			for _, pkt := range toSend {
+				if _, err := rtcpWriter.Write([]rtcp.Packet{pkt}, interceptor.Attributes{}); err != nil {
+					n.log.Warnf("failed sending nack: %+v", err)
+				}
+			}
+
 		case <-n.close:
 			return
 		}

pkg/nack/generator_interceptor_test.go

@@ -179,3 +179,85 @@ func TestGeneratorInterceptor_UnbindRemovesCorrespondingSSRC(t *testing.T) {
 	_, ok = gen.nackCountLogs[ssrc]
 	assert.False(t, ok, "ssrc should not be present in nackCountLogs")
 }
+
+// reentrantRTCPWriter tries to re-acquire GeneratorInterceptor.receiveLogsMu
+// inside Write. If loop() calls Write while holding that mutex, this will
+// cause a deadlock.
+type reentrantRTCPWriter struct {
+	n      *GeneratorInterceptor
+	called chan struct{}
+}
+
+func (w *reentrantRTCPWriter) Write(pkts []rtcp.Packet, attrs interceptor.Attributes) (int, error) {
+	// signal to the test that Write was entered.
+	select {
+	case <-w.called:
+		// already closed
+	default:
+		close(w.called)
+	}
+
+	// re-enter the interceptor's lock
+	w.n.receiveLogsMu.Lock()
+	defer w.n.receiveLogsMu.Unlock()
+
+	return len(pkts), nil
+}
+
+// this fails if loop() calls rtcpWriter.Write while holding receiveLogsMu
+// but will pass if Write is called after releasing receiveLogsMu.
+func TestGeneratorInterceptor_NoDeadlockWithReentrantRTCPWriter(t *testing.T) {
+	const interval = time.Millisecond * 5
+
+	f, err := NewGeneratorInterceptor(
+		GeneratorSize(64),
+		GeneratorInterval(interval),
+	)
+	assert.NoError(t, err)
+
+	i, err := f.NewInterceptor("")
+	assert.NoError(t, err)
+	gen, ok := i.(*GeneratorInterceptor)
+	assert.True(t, ok, "expected *GeneratorInterceptor, got %T", i)
+
+	writer := &reentrantRTCPWriter{
+		n:      gen,
+		called: make(chan struct{}),
+	}
+	_ = gen.BindRTCPWriter(writer)
+
+	// set receiveLog with a gap so that missingSeqNumbers()
+	// returns something and causes a NACK -> Write() call.
+	rl, err := newReceiveLog(gen.size)
+	assert.NoError(t, err)
+
+	gen.receiveLogsMu.Lock()
+	gen.receiveLogs[1] = rl
+	gen.receiveLogsMu.Unlock()
+
+	// 100 and 102 received -> 101 is missing.
+	rl.add(100)
+	rl.add(102)
+
+	// wait until the writer was actually called at least once.
+	select {
+	case <-writer.called:
+		// good: generator loop attempted to send a NACK
+	case <-time.After(time.Second):
+		assert.Fail(t, "generator did not call RTCP writer")
+	}
+
+	// verify that Close() does not deadlock.
+	done := make(chan struct{})
+	go func() {
+		_ = gen.Close()
+		close(done)
+	}()
+
+	select {
+	case <-done:
+		// no deadlock with reentrant writer
+	case <-time.After(time.Second):
+		assert.Fail(t, "GeneratorInterceptor.Close deadlocked with reentrant RTCP writer")
+	}
+}

pkg/nack/responder_interceptor.go

@@ -126,10 +126,10 @@ func (n *ResponderInterceptor) BindLocalStream(
 			if err != nil {
 				return 0, err
 			}
-			stream.rtpBufferMutex.Lock()
-			defer stream.rtpBufferMutex.Unlock()
 
-			rtpBuffer.Add(pkt)
+			stream.rtpBufferMutex.Lock()
+			stream.rtpBuffer.Add(pkt)
+			stream.rtpBufferMutex.Unlock()
 
 			return writer.Write(header, payload, attributes)
 		},
@@ -153,10 +153,13 @@ func (n *ResponderInterceptor) resendPackets(nack *rtcp.TransportLayerNack) {
 
 	for i := range nack.Nacks {
 		nack.Nacks[i].Range(func(seq uint16) bool {
+			// save the packet under the buffer lock
 			stream.rtpBufferMutex.Lock()
-			defer stream.rtpBufferMutex.Unlock()
+			p := stream.rtpBuffer.Get(seq)
+			stream.rtpBufferMutex.Unlock()
 
-			if p := stream.rtpBuffer.Get(seq); p != nil {
+			if p != nil {
+				// send without holding rtpBufferMutex
 				if _, err := stream.rtpWriter.Write(p.Header(), p.Payload(), interceptor.Attributes{}); err != nil {
 					n.log.Warnf("failed resending nacked packet: %+v", err)
 				}

pkg/nack/responder_interceptor_test.go

@@ -398,3 +398,87 @@ func TestResponderInterceptor_BypassUnknownSSRCs(t *testing.T) {
 		}
 	}
 }
+
+// reentrantRTPWriter tries to re-acquire localStream.rtpBufferMutex inside Write.
+// If BindLocalStream's wrapper calls writer.Write while holding that mutex, this
+// will deadlock.
+type reentrantRTPWriter struct {
+	stream *localStream
+	called chan struct{}
+}
+
+func (w *reentrantRTPWriter) Write(header *rtp.Header, payload []byte, attrs interceptor.Attributes) (int, error) {
+	// signal to the test that Write was entered.
+	if w.called != nil {
+		select {
+		case <-w.called:
+			// already closed
+		default:
+			close(w.called)
+		}
+	}
+
+	// re-enter the same mutex.
+	w.stream.rtpBufferMutex.Lock()
+	defer w.stream.rtpBufferMutex.Unlock()
+
+	return len(payload), nil
+}
+
+// this fails if writer.Write is called while holding rtpBufferMutex and
+// will pass when the mutex is released before calling writer.Write.
+func TestResponderInterceptor_NoDeadlockWithReentrantRTPWriter(t *testing.T) {
+	f, err := NewResponderInterceptor(
+		ResponderSize(8),
+		ResponderLog(logging.NewDefaultLoggerFactory().NewLogger("test")),
+	)
+	require.NoError(t, err)
+
+	i, err := f.NewInterceptor("")
+	require.NoError(t, err)
+
+	resp, ok := i.(*ResponderInterceptor)
+	require.True(t, ok, "expected *ResponderInterceptor, got %T", i)
+
+	info := &interceptor.StreamInfo{
+		SSRC:         1,
+		RTCPFeedback: []interceptor.RTCPFeedback{{Type: "nack"}},
+	}
+
+	writer := &reentrantRTPWriter{
+		called: make(chan struct{}),
+	}
+
+	// BindLocalStream wraps the writer and stores a localStream in resp.streams.
+	wrapped := resp.BindLocalStream(info, writer)
+
+	// fill the writer with the actual localStream instance.
+	resp.streamsMu.Lock()
+	writer.stream = resp.streams[info.SSRC]
+	resp.streamsMu.Unlock()
+	require.NotNil(t, writer.stream, "localStream should not be nil")
+
+	header := &rtp.Header{SSRC: 1, SequenceNumber: 1}
+	payload := []byte{0x01}
+
+	done := make(chan struct{})
+	go func() {
+		_, _ = wrapped.Write(header, payload, interceptor.Attributes{})
+		close(done)
+	}()
+
+	// make sure the reentrant writer was actually invoked.
+	select {
+	case <-writer.called:
+		// good: reentrant path hit
+	case <-time.After(time.Second):
+		assert.Fail(t, "wrapped writer was never called")
+	}
+
+	select {
+	case <-done:
+		// no deadlock with reentrant writer
+	case <-time.After(time.Second):
+		assert.Fail(t, "ResponderInterceptor.Write deadlocked with reentrant RTP writer")
+	}
+}