av/packets/RtpToTsConverter.go

/*
NAME
  RtpToTsConverter.go - provides utilities for the conversion of Rtp packets
  to equivalent MpegTs packets.

DESCRIPTION
  See Readme.md

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

LICENSE
  RtpToTsConverter.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 (
	_ "fmt"
	"os"
)

type RtpToTsConverter interface {
	Convert()
}

type rtpToTsConverter struct {
	TsChan          <-chan *MpegTsPacket
	tsChan          chan<- *MpegTsPacket
	InputChan       chan<- RtpPacket
	inputChan       <-chan RtpPacket
	NalInputChan    chan<- []byte
	nalInputChan    <-chan []byte
	currentTsPacket *MpegTsPacket
	payloadByteChan chan byte
	currentCC       byte
}

func NewRtpToTsConverter() (c *rtpToTsConverter) {
	c = new(rtpToTsConverter)
	tsChan := make(chan *MpegTsPacket, 100)
	c.TsChan = tsChan
	c.tsChan = tsChan
	inputChan := make(chan RtpPacket, 100)
	c.InputChan = inputChan
	c.inputChan = inputChan
	nalInputChan := make(chan []byte, 10000)
	c.NalInputChan = nalInputChan
	c.nalInputChan = nalInputChan
	c.currentCC = 0
	return
}

func (c *rtpToTsConverter) Convert() {
	file, _ := os.Create("video")
	pesPktChan := make(chan []byte, 1000)
	pesDataChan := make(chan byte, 50000)
	payloadByteChan := make(chan byte, 100000)
	var rtpBuffer [](*RtpPacket)
	for {
		select {
		default:
		case rtpPacket := <-c.inputChan:
			rtpBuffer = append(rtpBuffer, &rtpPacket)
			if len(rtpBuffer) > 2 {
				// if there's something weird going on with sequence numbers then
				// insertion sort
				if rtpPacket.SequenceNumber < rtpBuffer[len(rtpBuffer)-2].SequenceNumber {
					for i := 1; i < len(rtpBuffer); i++ {
						for j := i; j > 0 && rtpBuffer[j].SequenceNumber < rtpBuffer[j-1].SequenceNumber; j-- {
							temp := rtpBuffer[j]
							rtpBuffer[j] = rtpBuffer[j-1]
							rtpBuffer[j-1] = temp
						}
					}
				}
			}
			if len(rtpBuffer) > 200 {
				// Discard everything before a type 7
				for GetOctectType(rtpBuffer[0]) != 7 {
					rtpBuffer = rtpBuffer[1:]
				}
				// get sps
				sps := make([]byte, len(rtpBuffer[0].Payload))
				copy(sps[:], rtpBuffer[0].Payload[:])
				rtpBuffer = rtpBuffer[1:]
				// get pps
				pps := make([]byte, len(rtpBuffer[0].Payload))
				copy(pps[:], rtpBuffer[0].Payload[:])
				rtpBuffer = rtpBuffer[1:]
				// get sei
				sei := make([]byte, len(rtpBuffer[0].Payload))
				copy(sei[:], rtpBuffer[0].Payload[:])
				rtpBuffer = rtpBuffer[1:]
				// while we haven't reached the next sps in the buffer
				for GetOctectType(rtpBuffer[0]) != 7 {
					switch GetOctectType(rtpBuffer[0]) {
					case 28:
						if GetStartBit(rtpBuffer[0]) == 1 {
							var buffer []byte
							buffer = append(buffer, []byte{0x00, 0x00, 0x01}...)
							buffer = append(buffer, []byte{0x09, 0x10}...)
							buffer = append(buffer, []byte{0x00, 0x00, 0x01}...)
							buffer = append(buffer, sps...)
							buffer = append(buffer, []byte{0x00, 0x00, 0x01}...)
							buffer = append(buffer, pps...)
							buffer = append(buffer, []byte{0x00, 0x00, 0x01}...)
							buffer = append(buffer, sei...)
							buffer = append(buffer, []byte{0x00, 0x00, 0x01}...)
							buffer = append(buffer, rtpBuffer[0].Payload[0]&0xE0|rtpBuffer[0].Payload[1]&0x1F)
							buffer = append(buffer, rtpBuffer[0].Payload[2:]...)
							rtpBuffer = rtpBuffer[1:]
							for {
								buffer = append(buffer, rtpBuffer[0].Payload[2:]...)
								if getEndBit(rtpBuffer[0]) == 1 {
									rtpBuffer = rtpBuffer[1:]
									c.NalInputChan <- buffer
									file.Write(buffer)
									break
								}
								rtpBuffer = rtpBuffer[1:]
							}
						}
					case 1:
						var buffer []byte
						buffer = append(buffer, []byte{0x00, 0x00, 0x01}...)
						buffer = append(buffer, []byte{0x09, 0x10}...)
						buffer = append(buffer, []byte{0x00, 0x00, 0x01}...)
						buffer = append(buffer, sps...)
						buffer = append(buffer, []byte{0x00, 0x00, 0x01}...)
						buffer = append(buffer, pps...)
						buffer = append(buffer, []byte{0x00, 0x00, 0x01}...)
						buffer = append(buffer, sei...)
						buffer = append(buffer, []byte{0x00, 0x00, 0x01}...)
						buffer = append(buffer, rtpBuffer[0].Payload[0]&0xE0|rtpBuffer[0].Payload[1]&0x1F)
						buffer = append(buffer, rtpBuffer[0].Payload[2:]...)
						rtpBuffer = rtpBuffer[1:]
						c.NalInputChan <- buffer
						file.Write(buffer)
					default:
					}

				}
			}
		case nalUnit := <-c.nalInputChan:
			for ii := range nalUnit {
				pesDataChan <- nalUnit[ii]
			}
			pesDataChanLen := len(nalUnit)
			pesPkt := new(PESPacket)
			pesPkt.StreamID = 0xE0
			pesPkt.Length = uint16(3 + pesDataChanLen)
			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, pesDataChanLen)
			for ii := 0; ii < pesDataChanLen; ii++ {
				pesPkt.Data[ii] = <-pesDataChan
			}
			pesPktChan <- pesPkt.ToByteSlice()
		case pesPkt := <-pesPktChan:
			for ii := range pesPkt {
				payloadByteChan <- pesPkt[ii]
			}
			firstPacket := true
			for len(payloadByteChan) > 0 {
				lengthOfByteChan := len(payloadByteChan)
				c.currentTsPacket = new(MpegTsPacket)
				c.currentTsPacket.SyncByte = 0x47
				c.currentTsPacket.TEI = false
				c.currentTsPacket.PUSI = false
				if firstPacket { // if it's the start of the payload
					c.currentTsPacket.PUSI = true
					firstPacket = false
				}
				c.currentTsPacket.Priority = false
				c.currentTsPacket.PID = 256
				c.currentTsPacket.TSC = 0
				c.currentTsPacket.CC = c.currentCC
				if c.currentCC++; c.currentCC > 15 {
					c.currentCC = 0
				}
				payloadLength := 182
				if lengthOfByteChan < 182 {
					payloadLength = lengthOfByteChan
				}
				c.currentTsPacket.AFC = 3
				stuffingLength := 182 - payloadLength
				c.currentTsPacket.AF = make([]byte, 2+stuffingLength) // adaptationfield flag length = 16
				c.currentTsPacket.AF[0] = byte(1 + stuffingLength)
				c.currentTsPacket.AF[1] = 0
				if c.currentTsPacket.PUSI {
					c.currentTsPacket.AF[1] = 0x00
				}
				for ii := 0; ii < stuffingLength; ii++ {
					c.currentTsPacket.AF[2+ii] = 0xFF
				}
				c.currentTsPacket.Payload = make([]byte, payloadLength)
				for ii := 0; ii < payloadLength; ii++ {
					c.currentTsPacket.Payload[ii] = <-payloadByteChan
				}
				c.tsChan <- c.currentTsPacket
			}
		}
	}
}