diff --git a/attributes.go b/attributes.go index 0aa6f9de..0725588f 100644 --- a/attributes.go +++ b/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") diff --git a/pkg/ccfb/interceptor.go b/pkg/ccfb/interceptor.go index ec08cb7e..0438f61e 100644 --- a/pkg/ccfb/interceptor.go +++ b/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: diff --git a/pkg/ccfb/interceptor_test.go b/pkg/ccfb/interceptor_test.go index 4e513718..6632c3ef 100644 --- a/pkg/ccfb/interceptor_test.go +++ b/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") + } +} diff --git a/pkg/ccfb/recorder.go b/pkg/ccfb/recorder.go index 4caa4f34..95170aa7 100644 --- a/pkg/ccfb/recorder.go +++ b/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) diff --git a/pkg/ccfb/stream_log.go b/pkg/ccfb/stream_log.go index 30021a6d..72b67b11 100644 --- a/pkg/ccfb/stream_log.go +++ b/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, } diff --git a/pkg/ccfb/stream_log_test.go b/pkg/ccfb/stream_log_test.go index 7097b199..af42223d 100644 --- a/pkg/ccfb/stream_log_test.go +++ b/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) {