Add support for RTP header extensions, and IVF videos now fully work!

examples/save-to-disk/ivf-writer.go

@@ -9,8 +9,6 @@ import (
 	"github.com/pions/webrtc/rtp"
 )
 
-var le = binary.LittleEndian
-
 type IVFWriter struct {
 	fd           *os.File
 	time         time.Time
@@ -49,16 +47,16 @@ func NewIVFWriter(fileName string) (*IVFWriter, error) {
 	}
 
 	header := make([]byte, 32)
-	copy(header[0:], []byte("DKIF")) // DKIF
-	le.PutUint16(header[4:], 0)      // Version
-	le.PutUint16(header[6:], 32)     // Header Size
-	copy(header[8:], []byte("VP80")) // FOURCC
-	le.PutUint16(header[12:], 640)   // Version
-	le.PutUint16(header[14:], 480)   // Header Size
-	le.PutUint32(header[16:], 30)    // Framerate numerator
-	le.PutUint32(header[20:], 1)     // Framerate Denominator
-	le.PutUint32(header[24:], 900)   // Frame count
-	le.PutUint32(header[28:], 0)     // Unused
+	copy(header[0:], []byte("DKIF"))                // DKIF
+	binary.LittleEndian.PutUint16(header[4:], 0)    // Version
+	binary.LittleEndian.PutUint16(header[6:], 32)   // Header Size
+	copy(header[8:], []byte("VP80"))                // FOURCC
+	binary.LittleEndian.PutUint16(header[12:], 640) // Version
+	binary.LittleEndian.PutUint16(header[14:], 480) // Header Size
+	binary.LittleEndian.PutUint32(header[16:], 30)  // Framerate numerator
+	binary.LittleEndian.PutUint32(header[20:], 1)   // Framerate Denominator
+	binary.LittleEndian.PutUint32(header[24:], 900) // Frame count
+	binary.LittleEndian.PutUint32(header[28:], 0)   // Unused
 
 	panicWrite(f, header)
 
@@ -103,7 +101,6 @@ func (i *IVFWriter) DecodeVP8RTPPacket(packet *rtp.Packet) (*VP8RTPPacket, error
 		payloadIndex++
 	}
 
-	fmt.Print(vp8Packet)
 	vp8Packet.Payload = p[payloadIndex:]
 
 	return vp8Packet, nil
@@ -123,9 +120,8 @@ func (i *IVFWriter) AddPacket(packet *rtp.Packet) {
 	}
 
 	frameHeader := make([]byte, 12)
-
-	le.PutUint32(frameHeader[0:], uint32(len(i.currentFrame))) // Frame length
-	le.PutUint64(frameHeader[4:], i.count)                     // PTS
+	binary.LittleEndian.PutUint32(frameHeader[0:], uint32(len(i.currentFrame))) // Frame length
+	binary.LittleEndian.PutUint64(frameHeader[4:], i.count)                     // PTS
 
 	i.count += 1
 

rtp/rtp_packet.go

@@ -14,23 +14,26 @@ type Packet struct {
 }
 
 const (
-	headerLength    = 4
-	versionShift    = 6
-	versionMask     = 0x3
-	paddingShift    = 5
-	paddingMask     = 0x1
-	extensionShift  = 4
-	extensionMask   = 0x1
-	ccMask          = 0xF
-	markerShift     = 7
-	markerMask      = 0x1
-	ptMask          = 0x7F
-	seqNumOffset    = 2
-	seqNumLength    = 2
-	timestampOffset = 4
-	timestampLength = 4
-	ssrcOffset      = 8
-	ssrcLength      = 4
-	csrcOffset      = 12
-	csrcLength      = 4
+	headerLength                 = 4
+	versionShift                 = 6
+	versionMask                  = 0x3
+	paddingShift                 = 5
+	paddingMask                  = 0x1
+	extensionShift               = 4
+	extensionMask                = 0x1
+	ccMask                       = 0xF
+	markerShift                  = 7
+	markerMask                   = 0x1
+	ptMask                       = 0x7F
+	seqNumOffset                 = 2
+	seqNumLength                 = 2
+	timestampOffset              = 4
+	timestampLength              = 4
+	ssrcOffset                   = 8
+	ssrcLength                   = 4
+	csrcOffset                   = 12
+	csrcLength                   = 4
+	extensionHeaderIdLength      = 2
+	extensionLengthFieldLength   = 2
+	extensionHeaderAssumedLength = 4
 )

rtp/unmarshal.go

@@ -2,6 +2,7 @@ package rtp
 
 import (
 	"encoding/binary"
+
 	"github.com/pkg/errors"
 )
 
@@ -37,9 +38,9 @@ func (p *Packet) Unmarshal(rawPacket []byte) error {
 	p.Timestamp = binary.BigEndian.Uint32(rawPacket[timestampOffset : timestampOffset+timestampLength])
 	p.SSRC = binary.BigEndian.Uint32(rawPacket[ssrcOffset : ssrcOffset+ssrcLength])
 
-	headerPlusCsrcs := headerLength + timestampLength + ssrcLength + (len(p.CSRC) * csrcLength)
-	if len(rawPacket) < headerPlusCsrcs {
-		return errors.Errorf("RTP header size insufficient; %d < %d", len(rawPacket), headerPlusCsrcs)
+	currOffset := headerLength + timestampLength + ssrcLength + (len(p.CSRC) * csrcLength)
+	if len(rawPacket) < currOffset {
+		return errors.Errorf("RTP header size insufficient; %d < %d", len(rawPacket), currOffset)
 	}
 
 	for i := range p.CSRC {
@@ -47,7 +48,15 @@ func (p *Packet) Unmarshal(rawPacket []byte) error {
 		p.CSRC[i] = binary.BigEndian.Uint32(rawPacket[offset:offset])
 	}
 
-	p.Payload = rawPacket[headerPlusCsrcs:]
+	if p.Extension {
+		currOffset += extensionHeaderIdLength
+		extensionLength := binary.BigEndian.Uint16(rawPacket[currOffset:])
+		currOffset += extensionLengthFieldLength
+
+		currOffset += (int(extensionLength) * extensionHeaderAssumedLength)
+	}
+
+	p.Payload = rawPacket[currOffset:]
 
 	return nil
 }