av/packets/packets_test.go

/*
NAME
  MpegTs.go - provides a data structure intended to encapsulate the properties
  of an MpegTs packet.

DESCRIPTION
  See Readme.md

AUTHOR
  Saxon Nelson-Milton <saxon.milton@gmail.com>

LICENSE
  MpegTs.go is Copyright (C) 2017 the Australian Ocean Lab (AusOcean)

  It is free software: you can redistribute it and/or modify them
  under the terms of the GNU General Public License as published by the
  Free Software Foundation, either version 3 of the License, or (at your
  option) any later version.

  It is distributed in the hope that it will be useful, but WITHOUT
  ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
  FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License
 for more details.

  You should have received a copy of the GNU General Public License
  along with revid in gpl.txt.  If not, see [GNU licenses](http://www.gnu.org/licenses).
*/

package packets

import (
	//"bytes"
	"fmt"
	"io"
	"log"
	"net"
	"reflect"
	"testing"
	"time"
	"math/rand"

	"github.com/beatgammit/rtsp"
)

/*******************************************************
Testing stuff related to connection i.e. rtsp, rtp, rtcp
********************************************************/
const (
	rtpPort  = 17300
	rtcpPort = 17319
	rtspUrl  = "rtsp://192.168.0.50:8554/CH002.sdp"
	rtpUrl   = "rtsp://192.168.0.50:8554/CH002.sdp/track1"
)

/* Let's see if we can connect to an rtsp device then read an rtp stream,
and then convert the rtp packets to mpegts packets and output. */
func TestRTSP(t *testing.T) {
	sess := rtsp.NewSession()
	res, err := sess.Options(rtspUrl)
	if err != nil {
		t.Errorf("Shouldn't have got error: %v\n", err)
	}

	res, err = sess.Describe(rtspUrl)
	if err != nil {
		log.Fatalln(err)
		t.Errorf("Shouldn't have got error: %v\n", err)
	}
	p, err := rtsp.ParseSdp(&io.LimitedReader{R: res.Body, N: res.ContentLength})
	if err != nil {
		t.Errorf("Shouldn't have got error: %v\n", err)
	}
	log.Printf("%+v", p)
	res, err = sess.Setup(rtpUrl, fmt.Sprintf("RTP/AVP;unicast;client_port=%d-%d", rtpPort, rtcpPort))
	if err != nil {
		t.Errorf("Shouldn't have got error: %v\n", err)
	}
	log.Println(res)
	res, err = sess.Play(rtspUrl, res.Header.Get("Session"))
	if err != nil {
		t.Errorf("Shouldn't have got error: %v\n", err)
	}
	log.Println(res)
}

func TestRTP(t *testing.T) {
	sess := rtsp.NewSession()
	res, err := sess.Options(rtspUrl)
	if err != nil {
		t.Errorf("Shouldn't have got error: %v\n", err)
	}
	res, err = sess.Describe(rtspUrl)
	if err != nil {
		log.Fatalln(err)
		t.Errorf("Shouldn't have got error: %v\n", err)
	}
	p, err := rtsp.ParseSdp(&io.LimitedReader{R: res.Body, N: res.ContentLength})
	if err != nil {
		t.Errorf("Shouldn't have got error: %v\n", err)
	}
	log.Printf("%+v", p)
	res, err = sess.Setup(rtpUrl, fmt.Sprintf("RTP/AVP;unicast;client_port=%d-%d", rtpPort, rtcpPort))
	if err != nil {
		t.Errorf("Shouldn't have got error: %v\n", err)
	}
	log.Println(res)
	res, err = sess.Play(rtspUrl, res.Header.Get("Session"))
	if err != nil {
		t.Errorf("Shouldn't have got error: %v\n", err)
	}
	log.Println(res)
	// create udp connection for rtp stuff
	rtpLaddr, err := net.ResolveUDPAddr("udp", "192.168.0.109:17300")
	if err != nil {
		t.Errorf("Local rtp addr not set! %v\n", err)
	}
	rtpAddr, err := net.ResolveUDPAddr("udp", "192.168.0.50:17300")
	if err != nil {
		t.Errorf("Resolving rtp address didn't work! %v\n", err)
	}
	rtpConn, err := net.DialUDP("udp", rtpLaddr, rtpAddr)
	if err != nil {
		t.Errorf("Conncection not established! %v\n", err)
	}
	// Create udp connection for rtcp stuff
	rtcpLaddr, err := net.ResolveUDPAddr("udp", "192.168.0.109:17319")
	if err != nil {
		t.Errorf("Local RTCP address not resolved! %v\n", err)
	}
	rtcpAddr, err := net.ResolveUDPAddr("udp", "192.168.0.50:17301")
	if err != nil {
		t.Errorf("Remote RTCP address not resolved! %v\n", err)
	}
	rtcpConn, err := net.DialUDP("udp", rtcpLaddr, rtcpAddr)
	if err != nil {
		t.Errorf("Connection not established! %v\n", err)
	}
	// let's create a session that will store useful stuff from the connections
	rtpSession := NewSession(rtpConn, rtcpConn)
	time.Sleep(2 * time.Second)
	select {
	default:
		t.Errorf("Should have got rtpPacket!")
	case rtpPacket := <-rtpSession.RtpChan:
		fmt.Printf("RTP packet: %v\n", rtpPacket)
	}
}


/*******************************************************
Testing stuff related to the Nal.go file
********************************************************/
var parseInput = []byte{
	0x6C, // 3NalUnitBits = 101(5), Fragment type = 1100 (type = 12 )
	0x94, // starbit = 1, endbit = 0, Reservedbit = 0, 5NalUnitBits = 10100 (20)
	0x8E, // 10001110 random frame byte
	0x26, // 00100110 random frame byte
	0xD0, // 11010000 random frame byte
}

var expectedParsing = []interface{}{
	byte(3),
	byte(12),
	bool(true),
	bool(false),
	bool(false),
	byte(20),
	[]byte{0x8E, 0x26, 0xD0},
}

const (
	nalTestType = 12
)

func TestNalFragmentParsing(t *testing.T) {
	nalUnit := ParseNALFragment(parseInput)
	value := reflect.ValueOf(*nalUnit)
	length := value.NumField()
	fields := make([]interface{}, length)
	for ii := 0; ii < length; ii++ {
		fields[ii] = value.Field(ii).Interface()
	}
	for ii := range fields {
		if !reflect.DeepEqual(fields[ii], expectedParsing[ii]) {
			t.Errorf("Bad Parsing! Field: %v wanted: %v got: %v\n", ii, expectedParsing[ii],
				fields[ii])
		}
	}
}

func TestNalFragmentToByteSlice(t *testing.T) {
	nalUnit := ParseNALFragment(parseInput)
	output := nalUnit.ToByteSlice()
	for ii := range output {
		if output[ii] != parseInput[ii] {
			t.Errorf("Bad conversion to byte slice at %vth byte! wanted: %v got: %v",
				parseInput[ii], output[ii])
		}
	}
}

func TestNalFragmentType(t *testing.T) {
	nalUnit := ParseNALFragment(parseInput)
	nalType := nalUnit.GetType()
	if nalType != nalTestType {
		t.Errorf("Returned wrong type!")
	}
}

func TestNalSpsPpsParsing(t *testing.T) {
	nalSpsPps := ParseNALSpsPps(parseInput)
	for ii := range parseInput {
		if nalSpsPps.Data[ii] != parseInput[ii] {
			t.Errorf("Bad Parsing! Byte: %v wanted: %v got: %v\n", ii, parseInput[ii],
				nalSpsPps.Data[ii])
		}
	}
}

func TestNalSpsPpsToByteSlice(t *testing.T) {
	nalSpsPps := ParseNALSpsPps(parseInput)
	nalSpsPpsByteSlice := nalSpsPps.ToByteSlice()
	for ii := range parseInput {
		if nalSpsPpsByteSlice[ii] != parseInput[ii] {
			t.Errorf("Bad conversion to byte slice! Byte: %v wanted: %v got: %v\n", ii,
				parseInput[ii], nalSpsPpsByteSlice[ii])
		}
	}
}

func TestNalSpsPpsType(t *testing.T) {
	nalSpsPps := ParseNALSpsPps(parseInput)
	if nalSpsPps.GetType() != nalTestType {
		t.Errorf("Returned wrong type!")
	}
}

/*******************************************************
Pes Packet testing!
********************************************************/
const (
	dataLength = 3 // bytes
)

func TestPesToByteSlice(t *testing.T) {
	pesPkt := PESPacket{
		byte(0xE0),  // StreamID
		uint16(6),  // Length
		byte(0),     // ScramblingControl
		bool(true),  // Priority
		bool(false), // DAI
		bool(false), // copyright
		bool(true),  // Original
		byte(0),     // PDI
		bool(false), // Escr
		bool(false), // ESRate
		bool(false), // DSMTrickMode
		bool(false), // ACI
		bool(false), // CRC
		bool(false), // Ext
		byte(0),     // header length
		[]byte{},
		[]byte{},
		[]byte{ // data
			0xEA,
			0x4B,
			0x12,
		},
	}
	pesExpectedOutput := []byte{
		0x00, // packet start code prefix byte 1
		0x00, // packet start code prefix byte 2
		0x01, // packet start code prefix byte 3
		0xE0, // stream ID
		0x00, // PES Packet length byte 1
		0x06, // PES packet length byte 2
		0x89, // Marker bits,ScramblingControl, Priority, DAI, Copyright, Original
		0x00, // PDI, ESCR, ESRate, DSMTrickMode, ACI, CRC, Ext
		0x00, // header length
		0xEA, // data byte 1
		0x4B, // data byte 2
		0x12, // data byte 3
	}
	pesPktAsByteSlice := pesPkt.ToByteSlice()
	for ii := range pesPktAsByteSlice {
		if pesPktAsByteSlice[ii] != pesExpectedOutput[ii] {
			t.Errorf("Conversion to byte slice bad! Byte: %v Wanted: %v Got: %v",
				ii, pesExpectedOutput[ii], pesPktAsByteSlice[ii])
		}
	}
}

/*******************************************************
Mpegts testing
********************************************************/

func TestMpegTsToByteSlice(t *testing.T){
	afRemainderLength := 180
	afField := make([]byte, afRemainderLength+1)
	afField[0] = byte(afRemainderLength)
	afField[1] = byte(0)
	for i := 2; i < len(afField); i++ {
		afField[i] = 0xFF
	}
	tsPkt := MpegTsPacket{
		byte(0x47), 	// sync byte
		bool(false),	// TEI
		bool(false),	// PUSI
		bool(false),	// Priority
		uint16(256),	// PID
		byte(0),			// TSC
		byte(3),			// AFC
		byte(6), 			// CC
		afField,			// AF
		[]byte{				// data
			0x67,
			0xB2,
			0xE3,
		},
	}
	expectedOutput := []byte{
		0x47,
		0x01,
		0x00,
		0x36,
		byte(afRemainderLength),
		0x00,
		// this is where stuffing is, expect that to be 0xFF
		0x67,
		0xB2,
		0xE3,
	}
	tsPktAsByteSlice := tsPkt.ToByteSlice()
	for ii := 0; ii < 6; ii++ {
		if tsPktAsByteSlice[ii] != expectedOutput[ii] {
			t.Errorf("Conversion to byte slice bad! Byte: %v Wanted: %v Got: %v",
				ii, expectedOutput[ii], tsPktAsByteSlice[ii])
		}
	}
	// Check that the stuffing is all there
	for ii := 6; ii < 185; ii++ {
		if tsPktAsByteSlice[ii] != 0xFF {
			t.Errorf("Conversion to byte slice bad! Byte: %v Wanted: %v Got: %v",
				ii, byte(0xFF), tsPktAsByteSlice[ii])
		}
	}
	for ii := 185; ii < 188; ii++ {
		if tsPktAsByteSlice[ii] != expectedOutput[ii-185+6] {
			t.Errorf("Conversion to byte slice bad! Byte: %v Wanted: %v Got: %v",
				ii, expectedOutput[ii], tsPktAsByteSlice[ii])
		}
	}
}

/*******************************************************
RtpToTsConverter testing
********************************************************/
/*
func TestRtpToTsConverter(t *testing.T){
	converter := NewRtpToTsConverter()
	go converter.Convert()
	// Create first rtp packet
	rtpPacket1 := new(RtpPacket)
	rtpPacket1.Version = 2
	rtpPacket1.Padding = false
	rtpPacket1.Ext = false
	rtpPacket1.CC = 0
	rtpPacket1.Marker = true
	rtpPacket1.PayloadType = 0xE3
	rtpPacket1.SequenceNumber = 1
	rtpPacket1.Timestamp = 200
	rtpPacket1.SyncSource = 0
	rtpPacket1.CSRC = nil
	rtpPacket1.ExtHeader = 0
	rtpPacket1.ExtData = nil
	nalFragment := new(NALFragment)
	nalFragment.ThreeNUBs = 0x02
	nalFragment.FragmentType = byte(28)
	nalFragment.Start = true
	nalFragment.End = false
	nalFragment.Reserved = true
	nalFragment.FiveNUBs = 0x03
	nalFragment.Data = make([]byte,98)
	rand.Seed(int64(time.Now().Nanosecond()))
	for i := range nalFragment.Data {
		nalFragment.Data[i] = byte(rand.Intn(255))
		fmt.Printf(" %v ", nalFragment.Data[i])
	}
	rtpPacket1.Payload = make([]byte,100)
	copy(rtpPacket1.Payload[:], nalFragment.ToByteSlice())
	fmt.Println(rtpPacket1.Payload)
	converter.InputChan<-(*rtpPacket1)
	// Create second rtp packet
	rtpPacket2 := new(RtpPacket)
	rtpPacket2.Version = 2
	rtpPacket2.Padding = false
	rtpPacket2.Ext = false
	rtpPacket2.CC = 0
	rtpPacket2.Marker = false
	rtpPacket2.PayloadType = 0xE3
	rtpPacket2.SequenceNumber = 2
	rtpPacket2.Timestamp = 300
	rtpPacket2.SyncSource = 0
	rtpPacket2.CSRC = nil
	rtpPacket2.ExtHeader = 0
	rtpPacket2.ExtData = nil
	nalFragment = new(NALFragment)
	nalFragment.ThreeNUBs = 0x02
	nalFragment.FragmentType = byte(28)
	nalFragment.Start = false
	nalFragment.End = true
	nalFragment.Reserved = true
	nalFragment.FiveNUBs = 0x03
	nalFragment.Data = make([]byte,198)
	for i := range nalFragment.Data {
		nalFragment.Data[i] = byte(rand.Intn(255))
	}
	rtpPacket2.Payload = make([]byte,200)
	copy(rtpPacket2.Payload[:], nalFragment.ToByteSlice())
	converter.InputChan<-(*rtpPacket2)

	// Create first expected tsPacket
	afField := make([]byte, 2)
	afField[0] = byte(1)
	afField[1] = byte(0)
	pesPkt := new(PESPacket)
	pesPkt.StreamID = 0xE0
	pesPkt.Length = uint16( 3 + 300 )
	pesPkt.ScramblingControl = 0
	pesPkt.Priority = true
	pesPkt.DAI = false
	pesPkt.Copyright = false
	pesPkt.Original = true
	pesPkt.PDI = 0
	pesPkt.ESCR = false
	pesPkt.ESRate = false
	pesPkt.DSMTrickMode = false
	pesPkt.ACI = false
	pesPkt.CRC = false
	pesPkt.Ext = false
	pesPkt.HeaderLength = 0
	pesPkt.Data = make([]byte,300)
	for ii:=0; ii<100; ii++ {
		pesPkt.Data[ii] = rtpPacket1.Payload[ii]
	}
	for ii:=100; ii <300; ii++ {
		pesPkt.Data[ii] = rtpPacket2.Payload[ii-100]
	}

	pesPacketAsByteSlice := pesPkt.ToByteSlice()

	data := make([]byte, 182)

	copy(data[:],pesPacketAsByteSlice[:182])

	expectedPkt1 := MpegTsPacket{
		byte(0x47), 	// sync byte
		bool(false),	// TEI
		bool(true),	// PUSI
		bool(false),	// Priority
		uint16(256),	// PID
		byte(0),			// TSC
		byte(3),			// AFC
		byte(0), 			// CC
		afField,			// AF
		data,
	}

	data = make([]byte, len(pesPacketAsByteSlice)-182)

	copy(data[:], pesPacketAsByteSlice[182:])

	afField = make([]byte, 2+(182-(len(pesPacketAsByteSlice)-182)))
	afField[0] = byte(1+(182-(len(pesPacketAsByteSlice)-182)))
	afField[1] = byte(0)
	for ii := 2; ii < len(afField); ii++ {
		afField[ii] = 0xFF
	}

	expectedPkt2 := MpegTsPacket{
		byte(0x47), 	// sync byte
		bool(false),	// TEI
		bool(false),	// PUSI
		bool(false),	// Priority
		uint16(256),	// PID
		byte(0),			// TSC
		byte(3),			// AFC
		byte(1), 			// CC
		afField,			// AF
		data,
	}
	// Now let's get our two Ts packets from the converter and see if they're  G
	tsPacket := <-converter.TsChan
	expectedPkt1AsByteSlice := expectedPkt1.ToByteSlice()
	tsPacketAsByteSlice := tsPacket.ToByteSlice()
	for ii := range expectedPkt1AsByteSlice {
		if expectedPkt1AsByteSlice[ii] != tsPacketAsByteSlice[ii] {
			t.Errorf("Not equal! Byte: %v Exptected: %v Got: %v\n",ii,
				expectedPkt1AsByteSlice[ii],tsPacketAsByteSlice[ii])
		}
	}
	fmt.Printf("Expected packet: %v\n", expectedPkt1.ToByteSlice())
	fmt.Printf("Got packet: %v\n", tsPacket.ToByteSlice())
	tsPacket = <-converter.TsChan
	expectedPkt2AsByteSlice := expectedPkt2.ToByteSlice()
	tsPacketAsByteSlice = tsPacket.ToByteSlice()
	for ii := range expectedPkt2AsByteSlice {
		if expectedPkt2AsByteSlice[ii] != tsPacketAsByteSlice[ii] {
			t.Errorf("Not equal! Byte: %v Exptected: %v Got: %v\n",ii,
				expectedPkt2AsByteSlice[ii],tsPacketAsByteSlice[ii])
		}
	}
	fmt.Printf("Expected packet: %v\n", expectedPkt2.ToByteSlice())
	fmt.Printf("Got packet: %v\n", tsPacket.ToByteSlice())
}
*/

func TestH264Parsing(t *testing.T) {
	// Using file
	/*
	file, err := os.Open(fileName)
	if err != nil {
		panic("Could not open file!")
		return
	}
	stats, err := file.Stat()
	if err != nil {
		panic("Could not get file stats!")
	}
	buffer := make([]byte, stats.Size())
	_, err = file.Read(buffer)
	if err != nil {
		panic("Could not read file!")
	}
	*/
	// straight from buffer
	someData := []byte{
		0,0,1,7,59,100,45,82,93,0,0,1,8,23,78,65,0,0,1,6,45,34,23,3,2,0,0,1,5,3,4,5,
		56,76,4,234,78,65,34,34,43,0,0,1,7,67,10,45,8,93,0,0,1,8,23,7,5,0,0,1,6,
		4,34,2,3,2,0,0,1,1,3,4,5,5,76,4,234,78,65,34,34,43,45,
	}
	nalAccess1 := []byte{
		0,0,1,7,59,100,45,82,93,0,0,1,8,23,78,65,0,0,1,6,45,34,23,3,2,0,0,1,5,3,4,5,
		56,76,4,234,78,65,34,34,43,
	}
	nalAccess2 := []byte{
		0,0,1,7,67,10,45,8,93,0,0,1,8,23,7,5,0,0,1,6,
		4,34,2,3,2,0,0,1,1,3,4,5,5,76,4,234,78,65,34,34,43,45,
	}

	aChannel := make(chan []byte, 10)
	var nalAccessChan chan<- []byte
	nalAccessChan = aChannel
	go parseH264Buffer(someData,nalAccessChan)
	anAccessUnit := <-nalAccessChan
	for i := range anAccessUnit {
		if anAccessUnit[i] != nalAccess1[i] {
			t.Errorf("Should have been equal!")
		}
	}
	anAccessUnit = <-nalAccessChan
	for i := range anAccessUnit {
		if anAccessUnit[i] != nalAccess2[i] {
			t.Errorf("Should have been equal!")
		}
	}
}