Refactored samplebuilder logic

Many corner cases would cause samplebuilder to fail and return invalid
results. This refactoring is more reliable in all cases.

Fixed bug in H264 writer by reusing the packet object in  H264 writer.

Author: robin-raymond

README.md

@@ -275,6 +275,7 @@ Check out the **[contributing wiki](https://github.com/pion/webrtc/wiki/Contribu
 * [Jin Gong](https://github.com/cgojin)
 * [yusuke](https://github.com/yusukem99)
 * [Patryk Rogalski](https://github.com/digitalix)
+* [Robin Raymond](https://github.com/robin-raymond)
 
 ### License
 MIT License - see [LICENSE](LICENSE) for full text

examples/custom-logger/main.go

@@ -14,8 +14,7 @@ import (
 // customLogger satisfies the interface logging.LeveledLogger
 // a logger is created per subsystem in Pion, so you can have custom
 // behavior per subsystem (ICE, DTLS, SCTP...)
-type customLogger struct {
-}
+type customLogger struct{}
 
 // Print all messages except trace
 func (c customLogger) Trace(msg string)                          {}
@@ -41,8 +40,7 @@ func (c customLogger) Errorf(format string, args ...interface{}) {
 // customLoggerFactory satisfies the interface logging.LoggerFactory
 // This allows us to create different loggers per subsystem. So we can
 // add custom behavior
-type customLoggerFactory struct {
-}
+type customLoggerFactory struct{}
 
 func (c customLoggerFactory) NewLogger(subsystem string) logging.LeveledLogger {
 	fmt.Printf("Creating logger for %s \n", subsystem)

go.mod

@@ -12,7 +12,7 @@ require (
 	github.com/pion/logging v0.2.2
 	github.com/pion/randutil v0.1.0
 	github.com/pion/rtcp v1.2.6
-	github.com/pion/rtp v1.6.2
+	github.com/pion/rtp v1.6.5
 	github.com/pion/sctp v1.7.12
 	github.com/pion/sdp/v3 v3.0.4
 	github.com/pion/srtp/v2 v2.0.2

go.sum

@@ -53,8 +53,9 @@ 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.6 h1:1zvwBbyd0TeEuuWftrd/4d++m+/kZSeiguxU61LFWpo=
 github.com/pion/rtcp v1.2.6/go.mod h1:52rMNPWFsjr39z9B9MhnkqhPLoeHTv1aN63o/42bWE0=
-github.com/pion/rtp v1.6.2 h1:iGBerLX6JiDjB9NXuaPzHyxHFG9JsIEdgwTC0lp5n/U=
 github.com/pion/rtp v1.6.2/go.mod h1:bDb5n+BFZxXx0Ea7E5qe+klMuqiBrP+w8XSjiWtCUko=
+github.com/pion/rtp v1.6.5 h1:o2cZf8OascA5HF/b0PAbTxRKvOWxTQxWYt7SlToxFGI=
+github.com/pion/rtp v1.6.5/go.mod h1:bDb5n+BFZxXx0Ea7E5qe+klMuqiBrP+w8XSjiWtCUko=
 github.com/pion/sctp v1.7.10/go.mod h1:EhpTUQu1/lcK3xI+eriS6/96fWetHGCvBi9MSsnaBN0=
 github.com/pion/sctp v1.7.12 h1:GsatLufywVruXbZZT1CKg+Jr8ZTkwiPnmUC/oO9+uuY=
 github.com/pion/sctp v1.7.12/go.mod h1:xFe9cLMZ5Vj6eOzpyiKjT9SwGM4KpK/8Jbw5//jc+0s=

pkg/media/h264writer/h264writer.go

@@ -18,8 +18,9 @@ type (
 	// Therefore, only 1-23, 24 (STAP-A), 28 (FU-A) NAL types are allowed.
 	// https://tools.ietf.org/html/rfc6184#section-5.2
 	H264Writer struct {
-		writer      io.Writer
-		hasKeyFrame bool
+		writer       io.Writer
+		hasKeyFrame  bool
+		cachedPacket *codecs.H264Packet
 	}
 )
 
@@ -53,7 +54,11 @@ func (h *H264Writer) WriteRTP(packet *rtp.Packet) error {
 		}
 	}
 
-	data, err := (&codecs.H264Packet{}).Unmarshal(packet.Payload)
+	if h.cachedPacket == nil {
+		h.cachedPacket = &codecs.H264Packet{}
+	}
+
+	data, err := h.cachedPacket.Unmarshal(packet.Payload)
 	if err != nil {
 		return err
 	}
@@ -65,6 +70,7 @@ func (h *H264Writer) WriteRTP(packet *rtp.Packet) error {
 
 // Close closes the underlying writer
 func (h *H264Writer) Close() error {
+	h.cachedPacket = nil
 	if h.writer != nil {
 		if closer, ok := h.writer.(io.Closer); ok {
 			return closer.Close()

pkg/media/h264writer/h264writer_test.go

@@ -59,51 +59,61 @@ func TestWriteRTP(t *testing.T) {
 		hasKeyFrame bool
 		wantBytes   []byte
 		wantErr     error
+		reuseWriter bool
 	}{
 		{
 			"When given an empty payload; it should return nil",
 			[]byte{},
 			false,
 			[]byte{},
 			nil,
+			false,
 		},
 		{
 			"When no keyframe is defined; it should discard the packet",
 			[]byte{0x25, 0x90, 0x90},
 			false,
 			[]byte{},
 			nil,
+			false,
 		},
 		{
 			"When a valid Single NAL Unit packet is given; it should unpack it without error",
 			[]byte{0x27, 0x90, 0x90},
 			true,
 			[]byte{0x00, 0x00, 0x00, 0x01, 0x27, 0x90, 0x90},
 			nil,
+			false,
 		},
 		{
 			"When a valid STAP-A packet is given; it should unpack it without error",
 			[]byte{0x38, 0x00, 0x03, 0x27, 0x90, 0x90, 0x00, 0x05, 0x28, 0x90, 0x90, 0x90, 0x90},
 			true,
 			[]byte{0x00, 0x00, 0x00, 0x01, 0x27, 0x90, 0x90, 0x00, 0x00, 0x00, 0x01, 0x28, 0x90, 0x90, 0x90, 0x90},
 			nil,
+			false,
 		},
 		{
 			"When a valid FU-A start packet is given; it should unpack it without error",
 			[]byte{0x3C, 0x85, 0x90, 0x90, 0x90},
 			true,
-			[]byte{0x00, 0x00, 0x00, 0x01, 0x25, 0x90, 0x90, 0x90},
+			[]byte{},
 			nil,
+			true,
 		},
 		{
 			"When a valid FU-A end packet is given; it should unpack it without error",
 			[]byte{0x3C, 0x45, 0x90, 0x90, 0x90},
 			true,
-			[]byte{0x90, 0x90, 0x90},
+			[]byte{0x00, 0x00, 0x00, 0x01, 0x25, 0x90, 0x90, 0x90, 0x90, 0x90, 0x90},
 			nil,
+			false,
 		},
 	}
 
+	var reuseWriter *bytes.Buffer
+	var reuseH264Writer *H264Writer
+
 	for _, tt := range tests {
 		tt := tt
 		t.Run(tt.name, func(t *testing.T) {
@@ -112,6 +122,12 @@ func TestWriteRTP(t *testing.T) {
 				hasKeyFrame: tt.hasKeyFrame,
 				writer:      writer,
 			}
+			if reuseWriter != nil {
+				writer = reuseWriter
+			}
+			if reuseH264Writer != nil {
+				h264Writer = reuseH264Writer
+			}
 			packet := &rtp.Packet{
 				Payload: tt.payload,
 			}
@@ -120,7 +136,14 @@ func TestWriteRTP(t *testing.T) {
 
 			assert.Equal(t, tt.wantErr, err)
 			assert.True(t, bytes.Equal(tt.wantBytes, writer.Bytes()))
-			assert.Nil(t, h264Writer.Close())
+			if !tt.reuseWriter {
+				assert.Nil(t, h264Writer.Close())
+				reuseWriter = nil
+				reuseH264Writer = nil
+			} else {
+				reuseWriter = writer
+				reuseH264Writer = h264Writer
+			}
 		})
 	}
 }

pkg/media/media.go

@@ -9,9 +9,11 @@ import (
 
 // A Sample contains encoded media and timing information
 type Sample struct {
-	Data      []byte
-	Timestamp time.Time
-	Duration  time.Duration
+	Data               []byte
+	Timestamp          time.Time
+	Duration           time.Duration
+	PacketTimestamp    uint32
+	PrevDroppedPackets uint16
 }
 
 // Writer defines an interface to handle

pkg/media/samplebuilder/sampleSequenceLocation.go

@@ -0,0 +1,82 @@
+// Package samplebuilder provides functionality to reconstruct media frames from RTP packets.
+package samplebuilder
+
+import "math"
+
+type sampleSequenceLocation struct {
+	// head is the first packet in a sequence
+	head uint16
+	// tail is always set to one after the final sequence number,
+	// so if head == tail then the sequence is empty
+	tail uint16
+}
+
+func (l sampleSequenceLocation) empty() bool {
+	return l.head == l.tail
+}
+
+func (l sampleSequenceLocation) hasData() bool {
+	return l.head != l.tail
+}
+
+func (l sampleSequenceLocation) count() uint16 {
+	return seqnumDistance(l.head, l.tail)
+}
+
+const (
+	slCompareVoid = iota
+	slCompareBefore
+	slCompareInside
+	slCompareAfter
+)
+
+func minUint32(x, y uint32) uint32 {
+	if x < y {
+		return x
+	}
+	return y
+}
+
+// Distance between two seqnums
+func seqnumDistance32(x, y uint32) uint32 {
+	diff := int32(x - y)
+	if diff < 0 {
+		return uint32(-diff)
+	}
+
+	return uint32(diff)
+}
+
+func (l sampleSequenceLocation) compare(pos uint16) int {
+	if l.empty() {
+		return slCompareVoid
+	}
+
+	head32 := uint32(l.head)
+	count32 := uint32(l.count())
+	tail32 := head32 + count32
+
+	// pos32 is possibly two values, the normal value or a wrap
+	// around the start value, figure out which it is...
+
+	pos32Normal := uint32(pos)
+	pos32Wrap := uint32(pos) + math.MaxUint16 + 1
+
+	distNormal := minUint32(seqnumDistance32(head32, pos32Normal), seqnumDistance32(tail32, pos32Normal))
+	distWrap := minUint32(seqnumDistance32(head32, pos32Wrap), seqnumDistance32(tail32, pos32Wrap))
+
+	pos32 := pos32Normal
+	if distWrap < distNormal {
+		pos32 = pos32Wrap
+	}
+
+	if pos32 < head32 {
+		return slCompareBefore
+	}
+
+	if pos32 >= tail32 {
+		return slCompareAfter
+	}
+
+	return slCompareInside
+}