ECN support

attributes.go

@@ -10,11 +10,14 @@ import (
 	"github.com/pion/rtp"
 )
 
-type unmarshaledDataKeyType int
+type attributeKey int
 
 const (
-	rtpHeaderKey unmarshaledDataKeyType = iota
+	rtpHeaderKey attributeKey = iota
 	rtcpPacketsKey
+
+	// ECNKey identifies a packet's ECN marking, stored as an rtcp.ECN.
+	ECNKey
 )
 
 var errInvalidType = errors.New("found value of invalid type in attributes map")

pkg/ccfb/interceptor.go

@@ -83,7 +83,7 @@ type packet struct {
 	arrival        time.Time
 	ssrc           uint32
 	sequenceNumber uint16
-	ecn            uint8
+	ecn            rtcp.ECN
 }
 
 // BindRTCPWriter lets you modify any outgoing RTCP packets. It is called once per PeerConnection. The returned method
@@ -122,11 +122,12 @@ func (s *SenderInterceptor) BindRemoteStream(
 			return 0, nil, err
 		}
 
+		ecn, _ := attr.Get(interceptor.ECNKey).(rtcp.ECN)
 		p := packet{
 			arrival:        s.now(),
 			ssrc:           header.SSRC,
 			sequenceNumber: header.SequenceNumber,
-			ecn:            0, // ECN is not supported (yet).
+			ecn:            ecn,
 		}
 		select {
 		case <-s.close:

pkg/ccfb/interceptor_test.go

@@ -348,3 +348,67 @@ func TestConcurrentClose(t *testing.T) {
 	}
 	wg.Wait()
 }
+
+func TestInterceptorECN(t *testing.T) {
+	mTick := &test.MockTicker{C: make(chan time.Time)}
+	factory, err := NewSenderInterceptor(SenderTicker(func(time.Duration) ticker { return mTick }))
+	assert.NoError(t, err)
+	intcp, err := factory.NewInterceptor("")
+	assert.NoError(t, err)
+	defer func() { assert.NoError(t, intcp.Close()) }()
+
+	reports := make(chan []rtcp.Packet, 1)
+	intcp.BindRTCPWriter(interceptor.RTCPWriterFunc(
+		func(packets []rtcp.Packet, _ interceptor.Attributes) (int, error) {
+			reports <- packets
+
+			return 0, nil
+		},
+	))
+	header := rtp.Header{Version: 2, SSRC: 123456}
+	reader := intcp.BindRemoteStream(&interceptor.StreamInfo{SSRC: header.SSRC}, interceptor.RTPReaderFunc(
+		func(buf []byte, attrs interceptor.Attributes) (int, interceptor.Attributes, error) {
+			n, marshalErr := header.MarshalTo(buf)
+			header.SequenceNumber++
+
+			return n, attrs, marshalErr
+		},
+	))
+	cases := []struct {
+		attrs interceptor.Attributes
+		ecn   rtcp.ECN
+	}{
+		{interceptor.Attributes{interceptor.ECNKey: rtcp.ECNNonECT}, rtcp.ECNNonECT},
+		{interceptor.Attributes{interceptor.ECNKey: rtcp.ECNECT1}, rtcp.ECNECT1},
+		{interceptor.Attributes{interceptor.ECNKey: rtcp.ECNECT0}, rtcp.ECNECT0},
+		{interceptor.Attributes{interceptor.ECNKey: rtcp.ECNCE}, rtcp.ECNCE},
+		{nil, rtcp.ECNNonECT},
+		{interceptor.Attributes{interceptor.ECNKey: byte(3)}, rtcp.ECNNonECT},
+		{interceptor.Attributes{"ECN": rtcp.ECNCE, int(interceptor.ECNKey): rtcp.ECNCE}, rtcp.ECNNonECT},
+	}
+	for _, testCase := range cases {
+		_, _, err = reader.Read(make([]byte, 1500), testCase.attrs)
+		assert.NoError(t, err)
+	}
+	mTick.Tick(time.Now())
+
+	select {
+	case packets := <-reports:
+		if !assert.Len(t, packets, 1) {
+			return
+		}
+		report, ok := packets[0].(*rtcp.CCFeedbackReport)
+		if !assert.True(t, ok) || !assert.Len(t, report.ReportBlocks, 1) {
+			return
+		}
+		metrics := report.ReportBlocks[0].MetricBlocks
+		if !assert.Len(t, metrics, len(cases)) {
+			return
+		}
+		for i, testCase := range cases {
+			assert.Equal(t, testCase.ecn, metrics[i].ECN)
+		}
+	case <-time.After(time.Second):
+		assert.Fail(t, "ECN feedback report not received")
+	}
+}

pkg/ccfb/recorder.go

@@ -12,7 +12,7 @@ import (
 
 type packetReport struct {
 	arrivalTime time.Time
-	ecn         uint8
+	ecn         rtcp.ECN
 }
 
 // Recorder records incoming RTP packets and their arrival times. Recorder can
@@ -30,7 +30,7 @@ func NewRecorder() *Recorder {
 }
 
 // AddPacket writes a packet to the underlying stream.
-func (r *Recorder) AddPacket(ts time.Time, ssrc uint32, seq uint16, ecn uint8) {
+func (r *Recorder) AddPacket(ts time.Time, ssrc uint32, seq uint16, ecn rtcp.ECN) {
 	stream, ok := r.streams[ssrc]
 	if !ok {
 		stream = newStreamLog(ssrc)

pkg/ccfb/stream_log.go

@@ -32,7 +32,7 @@ func newStreamLog(ssrc uint32) *streamLog {
 	}
 }
 
-func (l *streamLog) add(ts time.Time, sequenceNumber uint16, ecn uint8) {
+func (l *streamLog) add(ts time.Time, sequenceNumber uint16, ecn rtcp.ECN) {
 	unwrappedSequenceNumber := l.sequence.Unwrap(sequenceNumber)
 	if !l.init {
 		l.init = true
@@ -81,7 +81,7 @@ func (l *streamLog) metricsAfter(reference time.Time, maxReportBlocks int64) rtc
 	gapDetected := false
 	for i := offset; i <= l.lastSequenceNumberReceived; i++ { //nolint:varnamelen // i int64
 		received := false
-		ecn := uint8(0)
+		ecn := rtcp.ECNNonECT
 		ato := uint16(0)
 		if report, ok := l.log[i]; ok {
 			received = true
@@ -90,7 +90,7 @@ func (l *streamLog) metricsAfter(reference time.Time, maxReportBlocks int64) rtc
 		}
 		metricBlocks[i-offset] = rtcp.CCFeedbackMetricBlock{
 			Received:          received,
-			ECN:               rtcp.ECN(ecn),
+			ECN:               ecn,
 			ArrivalTimeOffset: ato,
 		}
 

pkg/ccfb/stream_log_test.go

@@ -14,7 +14,7 @@ import (
 type input struct {
 	ts  time.Time
 	nr  uint16
-	ecn uint8
+	ecn rtcp.ECN
 }
 
 func TestStreamLogAdd(t *testing.T) {