av/stream/mts/psi/psi.go

package psi

import (
	"errors"
	"hash/crc32"
	"math/bits"
)

// Some common lengths
const (
	ESSDDefLen = 5
	DescDefLen = 2
	PMTDefLen  = 4
	PATLen     = 4
	TSSDefLen  = 5
	PSIDefLen  = 3
)

// Table Type IDs
const (
	PATTableID = 0x00
	PMTTableID = 0x02
)

const (
	timestampDescTag = 234
)

// Program specific information
type PSI struct {
	Pf  byte   // Point field
	Pfb []byte // Pointer filler bytes
	Tid byte   // Table ID
	Ssi bool   // Section syntax indicator (1 for PAT, PMT, CAT)
	Pb  bool   // Private bit (0 for PAT, PMT, CAT)
	Sl  uint16 // Section length
	Tss *TSS   // Table syntax section (length defined by SL) if length 0 then nil
	Crc uint32 // crc32 of entire table excluding pointer field, pointer filler bytes and the trailing CRC32
}

// Table syntax section
type TSS struct {
	Tide uint16 // Table ID extension
	V    byte   // Version number
	Cni  bool   // Current/next indicator
	Sn   byte   // Section number
	Lsn  byte   // Last section number
	Sd   SD     // Specific data PAT/PMT
}

// Specific Data, (could be PAT or PMT)
type SD interface {
	Bytes() []byte
}

// Program association table, implements SD
type PAT struct {
	Pn    uint16 // Program Number
	Pmpid uint16 // Program map PID
}

// Program mapping table, implements SD
type PMT struct {
	Pcrpid uint16 // Program clock reference pid
	Pil    uint16 // Program info length
	Pd     []Desc // Program descriptors
	Essd   *ESSD  // Elementary stream specific data
}

// Elementary stream specific data
type ESSD struct {
	St   byte   // Stream type
	Epid uint16 // Elementary pid
	Esil uint16 // Elementary stream
	Esd  []Desc // Elementary stream desriptors
}

// Descriptor
type Desc struct {
	Dt byte   // Descriptor tag
	Dl byte   // Descriptor length
	Dd []byte // Descriptor data
}

// ReadPSI creates a PSI data structure from a given byte slice that represents a PSI
func ReadPSI(data []byte) *PSI {
	psi := PSI{}
	pos := 0
	psi.Pf = data[pos]
	if psi.Pf != 0 {
		psi.Pfb = make([]byte, 0, psi.Pf)
		pos++
		for i := 0; i < int(psi.Pf); i++ {
			psi.Pfb = append(psi.Pfb, data[pos])
			pos++
		}
	}
	psi.Tid = data[pos]
	pos++
	psi.Ssi = byteToBool(data[pos] & 0x80)
	psi.Pb = byteToBool(data[pos] & 0x40)
	psi.Sl = uint16(data[pos]&0x03)<<8 | uint16(data[pos+1])
	pos += 2
	psi.Tss = readTSS(data[pos:], &psi)
	return &psi
}

// ReadTSS creates a TSS data structure from a given byte slice that represents a TSS
func readTSS(data []byte, p *PSI) *TSS {
	tss := TSS{}
	pos := 0
	tss.Tide = uint16(data[pos])<<8 | uint16(data[pos+1])
	pos += 2
	tss.V = (data[pos] & 0x3e) >> 1
	tss.Cni = byteToBool(data[pos] & 0x01)
	pos++
	tss.Sn = data[pos]
	pos++
	tss.Lsn = data[pos]
	pos++
	switch p.Tid {
	case PATTableID:
		tss.Sd = readPAT(data[pos:], &tss)
	case PMTTableID:
		tss.Sd = readPMT(data[pos:], &tss)
	default:
		panic("Can't yet deal with tables that are not PAT or PMT")
	}
	return &tss
}

func readPAT(data []byte, p *TSS) *PAT {
	pat := PAT{}
	pos := 0
	pat.Pn = uint16(data[pos])<<8 | uint16(data[pos+1])
	pos += 2
	pat.Pmpid = uint16(data[pos]&0x1f)<<8 | uint16(data[pos+1])
	return &pat
}

func readPMT(data []byte, p *TSS) *PAT {
	pmt := PMT{}
	pos := 0
	pmt.Pcrpid = uint16(data[pos]&0x1f)<<8 | uint16(data[pos+1])
	pos += 2
	pmt.Pil = uint16(data[pos]&0x03)<<8 | uint16(data[pos+1])
	pos += 2
	if pmt.Pil != 0 {
		pmt.Pd = readDescs(data[pos:], int(pmt.Pil))
	}
	pos += int(pmt.Pil)
	// TODO Read ES stuff
	pmt.Essd = readEssd(data[pos:])
	return nil
}

// readDescs reads provides a slice of Descs given a byte slice that represents Descs
// and the no of bytes that the descs accumilate
func readDescs(data []byte, descLen int) (o []Desc) {
	pos := 0
	o = make([]Desc, 1)
	o[0].Dt = data[pos]
	pos++
	o[0].Dl = data[pos]
	pos++
	o[0].Dd = make([]byte, o[0].Dl)
	for i := 0; i < int(o[0].Dl); i++ {
		o[0].Dd[i] = data[pos]
		pos++
	}
	if 2+len(o[0].Dd) != descLen {
		panic("No support for reading more than one descriptor")
	}
	return
}

func readEssd(data []byte) *ESSD {
	essd := ESSD{}
	pos := 0
	essd.St = data[pos]
	pos++
	essd.Epid = uint16(data[pos]&0x1f)<<8 | uint16(data[pos+1])
	pos += 2
	essd.Esil = uint16(data[pos]&0x03)<<8 | uint16(data[pos+1])
	pos += 2
	essd.Esd = readDescs(data[pos:], int(essd.Esil))
	return &essd
}

// Bytes outputs a byte slice representation of the PSI
func (p *PSI) Bytes() []byte {
	l := 1 + len(p.Pfb)
	out := make([]byte, l+PSIDefLen)
	out[0] = p.Pf
	for i, b := range p.Pfb {
		out[1+i] = b
	}
	out[l] = p.Tid
	out[l+1] = 0x80 | 0x40 | 0x30 | (0x03 & byte(p.Sl>>8))
	out[l+2] = byte(p.Sl)
	out = append(out, p.Tss.Bytes()...)
	crc32 := crc32_Update(0xffffffff, crc32_MakeTable(bits.Reverse32(crc32.IEEE)), out[l:])
	out = append(out, make([]byte, 0, 4)...)
	out = append(out, byte(crc32>>24))
	out = append(out, byte(crc32>>16))
	out = append(out, byte(crc32>>8))
	out = append(out, byte(crc32))
	return out
}

// Bytes outputs a byte slice representation of the TSS
func (t *TSS) Bytes() []byte {
	out := make([]byte, TSSDefLen)
	out[0] = byte(t.Tide >> 8)
	out[1] = byte(t.Tide)
	out[2] = 0xc0 | (0x3e & (t.V << 1)) | (0x01 & boolToByte(t.Cni))
	out[3] = t.Sn
	out[4] = t.Lsn
	out = append(out, t.Sd.Bytes()...)
	return out
}

// Bytes outputs a byte slice representation of the PAT
func (p *PAT) Bytes() []byte {
	out := make([]byte, PATLen)
	out[0] = byte(p.Pn >> 8)
	out[1] = byte(p.Pn)
	out[2] = 0xe0 | (0x1f & byte(p.Pmpid>>8))
	out[3] = byte(p.Pmpid)
	return out
}

// Bytes outputs a byte slice representation of the PMT
func (p *PMT) Bytes() []byte {
	out := make([]byte, PMTDefLen)
	out[0] = 0xe0 | (0x1f & byte(p.Pcrpid>>8))
	out[1] = byte(p.Pcrpid)
	out[2] = 0xf0 | (0x03 & byte(p.Pil>>8))
	out[3] = byte(p.Pil)
	for _, d := range p.Pd {
		out = append(out, d.Bytes()...)
	}
	out = append(out, p.Essd.Bytes()...)
	return out
}

// Bytes outputs a byte slice representation of the Desc
func (d *Desc) Bytes() []byte {
	out := make([]byte, DescDefLen)
	out[0] = d.Dt
	out[1] = d.Dl
	out = append(out, d.Dd...)
	return out
}

// Bytes outputs a byte slice representation of the ESSD
func (e *ESSD) Bytes() []byte {
	out := make([]byte, ESSDDefLen)
	out[0] = e.St
	out[1] = 0xe0 | (0x1f & byte(e.Epid>>8))
	out[2] = byte(e.Epid)
	out[3] = 0xf0 | (0x03 & byte(e.Esil>>8))
	out[4] = byte(e.Esil)
	for _, d := range e.Esd {
		out = append(out, d.Bytes()...)
	}
	return out
}

func boolToByte(b bool) (o byte) {
	if b {
		o = 0x01
	}
	return
}

func byteToBool(b byte) (o bool) {
	if b != 0 {
		o = true
	}
	return
}

func crc32_MakeTable(poly uint32) *crc32.Table {
	var t crc32.Table
	for i := range t {
		crc := uint32(i) << 24
		for j := 0; j < 8; j++ {
			if crc&0x80000000 != 0 {
				crc = (crc << 1) ^ poly
			} else {
				crc <<= 1
			}
		}
		t[i] = crc
	}
	return &t
}

func crc32_Update(crc uint32, tab *crc32.Table, p []byte) uint32 {
	for _, v := range p {
		crc = tab[byte(crc>>24)^v] ^ (crc << 8)
	}
	return crc
}

// UpdateTimestamp
func UpdateTimestamp(data []byte, t int) error {
	psi := ReadPSI(data)
	if psi.Tid != PATTableID {
		return errors.New("Timestamp update failed because psi is not a PMT")
	}
	return nil
}