protocol/rtp: extended functionality of rtp pkg

rtp.Bytes can now deal with CSRC fields, extension fields and padding. Tests have been updated to be more extensive in checks.

protocol/rtp/encoder.go

@@ -34,8 +34,6 @@ import (
 )
 
 const (
-	yes            = 1
-	no             = 0
 	defaultPktType = 33
 	timestampFreq  = 90000 // Hz
 	mtsSize        = 188
@@ -101,15 +99,15 @@ func min(a, b int) int {
 func (e *Encoder) Encode(payload []byte) error {
 	pkt := Pkt{
 		V:       rtpVer,           // version
-		X:       no,               // header extension
+		X:       false,            // header extension
-		CC:      no,               // CSRC count
+		CC:      0,                // CSRC count
-		M:       no,               // NOTE: need to check if this works (decoders should ignore this)
+		M:       false,            // NOTE: need to check if this works (decoders should ignore this)
 		PT:      defaultPktType,   // 33 for mpegts
 		SN:      e.nxtSeqNo(),     // sequence number
 		TS:      e.nxtTimestamp(), // timestamp
 		SSRC:    e.ssrc,           // source identifier
 		Payload: payload,
-		Padding: no,
+		Padding: nil,
 	}
 	_, err := e.dst.Write(pkt.Bytes(e.pktSpace[:defPktSize]))
 	if err != nil {

protocol/rtp/rtp.go

@@ -32,75 +32,109 @@ for rtp-h264 and rtp standards.
 */
 package rtp
 
+import (
+	"encoding/binary"
+)
+
 const (
 	rtpVer          = 2
-	headSize       = 3 * 4                     // Header size of an rtp packet.
+	defaultHeadSize = 3 * 4                            // Header size of an rtp packet.
 	defPayloadSize  = sendSize                         // Default payload size for the rtp packet.
-	defPktSize     = headSize + defPayloadSize // Default packet size is header size + payload size.
+	defPktSize      = defaultHeadSize + defPayloadSize // Default packet size is header size + payload size.
 )
 
 // Pkt provides fields consistent with RFC3550 definition of an rtp packet
 // The padding indicator does not need to be set manually, only the padding length
 type Pkt struct {
-	V       byte   // Version (currently 2)
+	V         uint8           // Version (currently 2).
-	p       byte   // Padding indicator (0 => padding, 1 => padding)
+	p         bool            // Padding indicator (0 => padding, 1 => padding).
-	X       byte   // Extension header indicator
+	X         bool            // Extension header indicator.
-	CC      byte   // CSRC count
+	CC        uint8           // CSRC count.
-	M       byte   // Marker bit
+	M         bool            // Marker bit.
-	PT      byte   // Packet type
+	PT        uint8           // Packet type.
-	SN      uint16 // Synch number
+	SN        uint16          // Synch number.
-	TS      uint32 // Timestamp
+	TS        uint32          // Timestamp.
-	SSRC    uint32 // Synchronisation source identifier
+	SSRC      uint32          // Synchronisation source identifier.
-	Payload []byte // H264 Payload data
+	CSRC      [][4]byte       // Contributing source identifier.
-	Padding byte   // No of bytes of padding
+	Extension ExtensionHeader // Header extension.
+	Payload   []byte          // H264 Payload data.
+	Padding   []byte          // No of bytes of padding.
+}
+
+// ExtensionHeader header provides fields for an RTP packet extension header.
+type ExtensionHeader struct {
+	ID     uint16
+	Len    uint16
+	Header [][4]byte
 }
 
 // Bytes provides a byte slice of the packet
 func (p *Pkt) Bytes(buf []byte) []byte {
-	if buf == nil || len(buf) != defPktSize {
+	// Calculate the required length for the RTP packet.
-		buf = make([]byte, headSize, defPktSize)
+	headerExtensionLen := 0
+	if p.X {
+		if p.Extension.Len != uint16(len(p.Extension.Header)) {
+			panic("bad rtp packet extension length field")
 		}
-	buf = buf[:headSize]
+		headerExtensionLen = int(4 + 4*p.Extension.Len)
-	if p.V == 0 {
-		p.V = rtpVer
 	}
+	requiredPktLen := defaultHeadSize + uint8(4*p.CC) + uint8(headerExtensionLen) + uint8(len(p.Payload)) + uint8(len(p.Padding))
 
-	if p.Padding > 0 {
+	// Create new space if no buffer is given, or it doesn't have sufficient capacity.
-		p.p = 1
+	if buf == nil || requiredPktLen > uint8(cap(buf)) {
+		buf = make([]byte, requiredPktLen, defPktSize)
 	}
+	buf = buf[:requiredPktLen]
 
+	// Start encoding fields into the buffer.
+	buf[0] = p.V<<6 | asByte(p.p)<<5 | asByte(p.X)<<4 | p.CC
+	buf[1] = asByte(p.M)<<7 | p.PT
+	binary.BigEndian.PutUint16(buf[2:4], p.SN)
+	binary.BigEndian.PutUint32(buf[4:8], p.TS)
+	binary.BigEndian.PutUint32(buf[8:12], p.SSRC)
+
+	// If there is a CSRC count, add the CSRC to the buffer.
 	if p.CC != 0 {
-		panic("CC has been set to something other than 0 - this is not supported yet.")
+		if p.CC != uint8(len(p.CSRC)) {
+			panic("CSRC count in RTP packet is incorrect")
+		}
+		for i := 0; i < int(p.CC); i++ {
+			copy(buf[12+i*4:], p.CSRC[i][:])
+		}
 	}
 
-	if p.X != 0 {
+	// This is our current index for writing to the buffer.
-		panic("rtp: X (extension header indicator) not 0, but extensiion headers not currently supported.")
+	idx := int(12 + 4*p.CC)
+
+	// If there is an extension field, add this to the buffer.
+	if p.X {
+		binary.BigEndian.PutUint16(buf[idx:idx+2], p.Extension.ID)
+		idx += 2
+		binary.BigEndian.PutUint16(buf[idx:idx+2], p.Extension.Len)
+		idx += 2
+		for i := 0; i < int(p.Extension.Len); i++ {
+			copy(buf[idx+i*4:], p.Extension.Header[i][:])
+		}
+		idx += int(p.Extension.Len * 4)
 	}
 
-	if p.CC != 0 {
+	// If there is payload, add to the buffer.
-		panic("rtp: CC (CSRC count) not 0, but CSRC headers not yet supported.")
+	if p.Payload != nil {
+		copy(buf[idx:], p.Payload)
+		idx += len(p.Payload)
 	}
 
-	buf[0] = p.V<<6 | p.p<<5 | p.CC
+	// Finally, if there is padding, add to the buffer.
-	buf[1] = p.M<<7 | p.PT
+	if p.Padding != nil {
-	buf[2] = byte(p.SN >> 8)
+		copy(buf[idx:], p.Padding)
-	buf[3] = byte(p.SN)
-	buf[4] = byte(p.TS >> 24)
-	buf[5] = byte(p.TS >> 16)
-	buf[6] = byte(p.TS >> 8)
-	buf[7] = byte(p.TS)
-	buf[8] = byte(p.SSRC >> 24)
-	buf[9] = byte(p.SSRC >> 16)
-	buf[10] = byte(p.SSRC >> 8)
-	buf[11] = byte(p.SSRC)
-
-	buf = append(buf, p.Payload...)
-	// see https://tools.ietf.org/html/rfc3550 section 5.1 (padding). At end of
-	// rtp packet, padding may exist, with the last octet being the length of the
-	// padding including itself.
-	if p.Padding != 0 {
-		buf = buf[:cap(buf)]
-		buf[len(buf)-1] = byte(p.Padding)
 	}
+
 	return buf
 }
+
+func asByte(b bool) byte {
+	if b {
+		return 0x01
+	}
+	return 0x00
+}

protocol/rtp/rtp_test.go

@@ -38,20 +38,23 @@ var rtpTests = []struct {
 	pkt  Pkt
 	want []byte
 }{
+	// No padding, no CSRC and no extension.
 	{
 		num: 1,
 		pkt: Pkt{
 			V:    2,
-			p:       0,
+			p:    false,
-			X:       0,
+			X:    false,
 			CC:   0,
-			M:       0,
+			M:    false,
 			PT:   6,
 			SN:   167,
 			TS:   160,
 			SSRC: 10,
-			Payload: []byte{0x00, 0x01, 0x07, 0xf0, 0x56, 0x37, 0x0a, 0x0f},
+			Payload: []byte{
-			Padding: 0,
+				0x00, 0x01, 0x07, 0xf0,
+				0x56, 0x37, 0x0a, 0x0f,
+			},
 		},
 		want: []byte{
 			0x80, 0x06, 0x00, 0xa7,
@@ -61,13 +64,138 @@ var rtpTests = []struct {
 			0x56, 0x37, 0x0a, 0x0f,
 		},
 	},
+	// With padding.
+	{
+		num: 2,
+		pkt: Pkt{
+			V:    2,
+			p:    true,
+			X:    false,
+			CC:   0,
+			M:    false,
+			PT:   6,
+			SN:   167,
+			TS:   160,
+			SSRC: 10,
+			Payload: []byte{
+				0x00, 0x01, 0x07, 0xf0,
+				0x56, 0x37, 0x0a, 0x0f,
+			},
+			Padding: []byte{
+				0x00, 0x00, 0x00, 0x00,
+				0x00, 0x00, 0x00, 0x00,
+				0x00, 0x00, 0x0b,
+			},
+		},
+		want: []byte{
+			0xa0, 0x06, 0x00, 0xa7,
+			0x00, 0x00, 0x00, 0xa0,
+			0x00, 0x00, 0x00, 0x0a,
+			0x00, 0x01, 0x07, 0xf0,
+			0x56, 0x37, 0x0a, 0x0f,
+			0x00, 0x00, 0x00, 0x00,
+			0x00, 0x00, 0x00, 0x00,
+			0x00, 0x00, 0x0b,
+		},
+	},
+	// With padding and CSRC.
+	{
+		num: 3,
+		pkt: Pkt{
+			V:    2,
+			p:    true,
+			X:    false,
+			CC:   2,
+			M:    false,
+			PT:   6,
+			SN:   167,
+			TS:   160,
+			SSRC: 10,
+			CSRC: [][4]byte{
+				{0x01, 0x02, 0x03, 0x04},
+				{0x05, 0x06, 0x07, 0x08},
+			},
+			Payload: []byte{
+				0x00, 0x01, 0x07, 0xf0,
+				0x56, 0x37, 0x0a, 0x0f,
+			},
+			Padding: []byte{
+				0x00, 0x00, 0x00, 0x00,
+				0x00, 0x00, 0x00, 0x00,
+				0x00, 0x00, 0x0b,
+			},
+		},
+		want: []byte{
+			0xa2, 0x06, 0x00, 0xa7,
+			0x00, 0x00, 0x00, 0xa0,
+			0x00, 0x00, 0x00, 0x0a,
+			0x01, 0x02, 0x03, 0x04,
+			0x05, 0x06, 0x07, 0x08,
+			0x00, 0x01, 0x07, 0xf0,
+			0x56, 0x37, 0x0a, 0x0f,
+			0x00, 0x00, 0x00, 0x00,
+			0x00, 0x00, 0x00, 0x00,
+			0x00, 0x00, 0x0b,
+		},
+	},
+	// With padding, CSRC and extension.
+	{
+		num: 4,
+		pkt: Pkt{
+			V:    2,
+			p:    true,
+			X:    true,
+			CC:   2,
+			M:    false,
+			PT:   6,
+			SN:   167,
+			TS:   160,
+			SSRC: 10,
+			CSRC: [][4]byte{
+				{0x01, 0x02, 0x03, 0x04},
+				{0x05, 0x06, 0x07, 0x08},
+			},
+			Extension: ExtensionHeader{
+				ID:  15,
+				Len: 2,
+				Header: [][4]byte{
+					{0x01, 0x02, 0x03, 0x04},
+					{0x05, 0x06, 0x07, 0x08},
+				},
+			},
+			Payload: []byte{
+				0x00, 0x01, 0x07, 0xf0,
+				0x56, 0x37, 0x0a, 0x0f,
+			},
+			Padding: []byte{
+				0x00, 0x00, 0x00, 0x00,
+				0x00, 0x00, 0x00, 0x00,
+				0x00, 0x00, 0x0b,
+			},
+		},
+		want: []byte{
+			0xb2, 0x06, 0x00, 0xa7,
+			0x00, 0x00, 0x00, 0xa0,
+			0x00, 0x00, 0x00, 0x0a,
+			0x01, 0x02, 0x03, 0x04,
+			0x05, 0x06, 0x07, 0x08,
+			0x00, 0x0f, 0x00, 0x02,
+			0x01, 0x02, 0x03, 0x04,
+			0x05, 0x06, 0x07, 0x08,
+			0x00, 0x01, 0x07, 0xf0,
+			0x56, 0x37, 0x0a, 0x0f,
+			0x00, 0x00, 0x00, 0x00,
+			0x00, 0x00, 0x00, 0x00,
+			0x00, 0x00, 0x0b,
+		},
+	},
 }
 
 func TestRtpPktToByteSlice(t *testing.T) {
 	for _, test := range rtpTests {
 		got := test.pkt.Bytes(nil)
 		if !reflect.DeepEqual(got, test.want) {
-			t.Errorf("unexpected error for test %v: got:%v want:%v", test.num, got,
+			t.Errorf("unexpected error for test %v: \ngot :%v\n want:%v\n", test.num, got,
 				test.want)
 		}
 	}