streaming

Media streaming and broadcast systems in Go
Log | Files | Refs | README | LICENSE

codec.go (8080B)


      1 package mpegts
      2 
      3 import (
      4 	"encoding/binary"
      5 	"errors"
      6 	"fmt"
      7 	"io"
      8 )
      9 
     10 var ErrLongPacket = errors.New("long packet")
     11 var ErrShortPacket = errors.New("short packet")
     12 
     13 func Unmarshal(buf []byte, p *Packet) error {
     14 	if len(buf) != PacketSize {
     15 		return fmt.Errorf("need exactly %d bytes, have %d", PacketSize, len(buf))
     16 	}
     17 	if buf[0] != Sync {
     18 		return fmt.Errorf("expected sync byte, got %x", buf[0])
     19 	}
     20 	p.Error = (buf[1] & 0x80) > 0
     21 	p.PayloadStart = (buf[1] & 0x40) > 0
     22 	p.Priority = (buf[1] & 0x20) > 0
     23 	// Want next 13 bits. 5 from buf[1] and all of buf[2].
     24 	pid := binary.BigEndian.Uint16([]byte{buf[1] & 0x1f, buf[2]})
     25 	p.PID = PacketID(pid)
     26 
     27 	// next 2 bits
     28 	p.Scrambling = Scramble(buf[3] & 0xc0)
     29 	// skip next 2 bits until later when we need to decode adaptation or payload.
     30 	// now just get last 4 bits
     31 	p.Continuity = buf[3] & 0x0f
     32 
     33 	afc := buf[3] >> 4
     34 	switch afc {
     35 	case 0x01:
     36 		buf = buf[4:]
     37 	case 0x02, 0x03:
     38 		p.Adaptation = parseAdaptationField(buf[4:])
     39 		if p.Adaptation == nil {
     40 			p.emptyAdaptation = true
     41 		}
     42 		alen := int(buf[4])
     43 		buf = buf[4+1+alen:]
     44 	default:
     45 		return fmt.Errorf("neither adaptation field or payload present")
     46 	}
     47 
     48 	return unmarshalPayload(buf, p)
     49 }
     50 
     51 func Decode(r io.Reader) (*Packet, error) {
     52 	buf := make([]byte, PacketSize)
     53 	n, err := r.Read(buf)
     54 	if n != PacketSize {
     55 		if err != nil {
     56 			return nil, fmt.Errorf("short read (%d bytes): %w", n, err)
     57 		}
     58 		return nil, fmt.Errorf("short read (%d bytes)", n)
     59 	}
     60 	if err != nil {
     61 		return nil, err
     62 	}
     63 	var p Packet
     64 	if err := Unmarshal(buf, &p); err != nil {
     65 		return &p, fmt.Errorf("unmarshal packet: %w", err)
     66 	}
     67 	return &p, nil
     68 }
     69 
     70 func parseAdaptationField(buf []byte) *Adaptation {
     71 	length := int(buf[0])
     72 	if length == 0 {
     73 		return nil
     74 	}
     75 	buf = buf[1 : length+1]
     76 	var af Adaptation
     77 	flags := buf[0]
     78 	buf = buf[1:]
     79 	af.Discontinuous = flags&0x80 > 0
     80 	af.RandomAccess = flags&0x40 > 0
     81 	af.Priority = flags&0x20 > 0
     82 	if flags&0x10 > 0 {
     83 		var p [6]byte
     84 		copy(p[:], buf[:6])
     85 		pcr := parsePCR(p)
     86 		af.PCR = &pcr
     87 		buf = buf[6:]
     88 	}
     89 	if flags&0x08 > 0 {
     90 		var p [6]byte
     91 		copy(p[:], buf[:6])
     92 		pcr := parsePCR(p)
     93 		af.OPCR = &pcr
     94 		buf = buf[6:]
     95 	}
     96 	if flags&0x04 > 0 {
     97 		af.SpliceCountdownSet = true
     98 		af.SpliceCountdown = buf[0]
     99 		buf = buf[1:]
    100 	}
    101 	if flags&0x02 > 0 {
    102 		tlen := int(buf[0])
    103 		af.Private = buf[1:tlen]
    104 		buf = buf[tlen:]
    105 	}
    106 	if flags&0x01 > 0 {
    107 		extlen := int(buf[0])
    108 		af.Extension = buf[1:extlen]
    109 		buf = buf[extlen:]
    110 	}
    111 	if len(buf) > 0 {
    112 		af.Stuffing = buf
    113 	}
    114 	return &af
    115 }
    116 
    117 // parsePCR parses the encoded PCR from a.
    118 // The 33-bit base and the 9-bit extension
    119 // are stored in a 6 byte array with the following bit layout,
    120 // where "b" stands for "base", "r" for reserved bits, and "e" for extension.
    121 //
    122 //	0 bbbb bbbb
    123 //	1 bbbb bbbb
    124 //	2 bbbb bbbb
    125 //	3 bbbb bbbb
    126 //	4 brrr rrre
    127 //	5 eeee eeee
    128 func parsePCR(a [6]byte) PCR {
    129 	// we only want the left-most bit.
    130 	// 6 bits are reserved and the right-most bit is part of extension.
    131 	b := [8]byte{0, 0, 0, a[0], a[1], a[2], a[3], a[4] & 0x80}
    132 	base := binary.BigEndian.Uint64(b[:])
    133 	base = base >> 7 // trim masked reserved, extension bits
    134 	// next 6 bits of a[4] are reserved, so right-most bit in a[4]
    135 	// and all of a[5] have the extension.
    136 
    137 	ext := binary.BigEndian.Uint16([]byte{a[4] & 0x01, a[5]})
    138 	return PCR{base, ext}
    139 }
    140 
    141 func unmarshalPayload(payload []byte, p *Packet) error {
    142 	if isPESPayload(payload) && p.PayloadStart {
    143 		pes, err := decodePES(payload)
    144 		if err != nil {
    145 			return fmt.Errorf("decode PES packet: %w", err)
    146 		}
    147 		p.PES = pes
    148 	} else {
    149 		p.Payload = payload
    150 	}
    151 	return nil
    152 }
    153 
    154 func Encode(w io.Writer, p *Packet) error {
    155 	buf := make([]byte, 4)
    156 	buf[0] = Sync
    157 	if p.Error {
    158 		buf[1] |= 0x80
    159 	}
    160 	if p.PayloadStart {
    161 		buf[1] |= 0x40
    162 	}
    163 	if p.Priority {
    164 		buf[1] |= 0x20
    165 	}
    166 	if p.PID > PacketNull {
    167 		return fmt.Errorf("packet id %s greater than max %s", p.PID, PacketNull)
    168 	}
    169 	buf[1] |= byte(p.PID >> 8)
    170 	buf[2] = byte(p.PID)
    171 
    172 	buf[3] |= byte(p.Scrambling)
    173 	if p.Adaptation != nil || p.emptyAdaptation {
    174 		buf[3] |= 0x20
    175 	}
    176 	if p.Payload != nil || p.PES != nil {
    177 		buf[3] |= 0x10
    178 	}
    179 	if p.Continuity > 15 {
    180 		return fmt.Errorf("continuity %d larger than max 4-bit integer %d", p.Continuity, 15)
    181 	}
    182 	buf[3] |= p.Continuity
    183 
    184 	if p.Adaptation != nil {
    185 		alen := 1 // just flags
    186 		if p.Adaptation.PCR != nil {
    187 			alen += 6
    188 		}
    189 		if p.Adaptation.OPCR != nil {
    190 			alen += 6
    191 		}
    192 		if p.Adaptation.SpliceCountdownSet {
    193 			alen++ // single byte
    194 		}
    195 		if p.Adaptation.Private != nil {
    196 			alen++ // 1 byte to store length of private
    197 			alen += len(p.Adaptation.Private)
    198 		}
    199 		alen += len(p.Adaptation.Extension)
    200 		alen += len(p.Adaptation.Stuffing)
    201 		if alen > 255 {
    202 			return fmt.Errorf("adaptation field too long: have %d bytes, max %d", alen, 255)
    203 		}
    204 
    205 		abuf := make([]byte, 1+alen) // length + total
    206 		abuf[0] = uint8(alen)
    207 		var i int = 2 // cursor; after length and flags
    208 		if p.Adaptation.Discontinuous {
    209 			abuf[1] |= 0x80
    210 		}
    211 		if p.Adaptation.RandomAccess {
    212 			abuf[1] |= 0x40
    213 		}
    214 		if p.Adaptation.Priority {
    215 			abuf[1] |= 0x20
    216 		}
    217 		if p.Adaptation.PCR != nil {
    218 			abuf[1] |= 0x10
    219 			if err := putPCR(abuf[i:i+6], p.Adaptation.PCR); err != nil {
    220 				return fmt.Errorf("pack PCR: %w", err)
    221 			}
    222 			i += 6
    223 		}
    224 		if p.Adaptation.OPCR != nil {
    225 			abuf[1] |= 0x08
    226 			if err := putPCR(abuf[i:i+6], p.Adaptation.OPCR); err != nil {
    227 				return fmt.Errorf("pack OPCR: %w", err)
    228 			}
    229 			i += 6
    230 		}
    231 		if p.Adaptation.SpliceCountdownSet {
    232 			abuf[1] |= 0x04
    233 			abuf[i] = p.Adaptation.SpliceCountdown
    234 			i++
    235 		}
    236 		if p.Adaptation.Private != nil {
    237 			abuf[1] |= 0x02
    238 			if len(p.Adaptation.Private) > 255 {
    239 				return fmt.Errorf("private data length %d longer than max %d", len(p.Adaptation.Private), 255)
    240 			}
    241 			abuf[i] = byte(len(p.Adaptation.Private))
    242 			i++
    243 			copy(abuf[i:], p.Adaptation.Private)
    244 			i += len(p.Adaptation.Private)
    245 		}
    246 		if p.Adaptation.Extension != nil {
    247 			abuf[1] |= 0x01
    248 			copy(abuf[i:], p.Adaptation.Extension)
    249 			i += len(p.Adaptation.Extension)
    250 		}
    251 		if p.Adaptation.Stuffing != nil {
    252 			copy(abuf[i:], p.Adaptation.Stuffing)
    253 		}
    254 		buf = append(buf, abuf...)
    255 	} else if p.emptyAdaptation {
    256 		// no adaptation field to encode, but we need to store an adaptation field length of 0.
    257 		buf = append(buf, 0)
    258 	}
    259 	if p.PES != nil {
    260 		b, err := encodePESPacket(p.PES)
    261 		if err != nil {
    262 			return fmt.Errorf("encode PES packet: %w", err)
    263 		}
    264 		buf = append(buf, b...)
    265 	}
    266 	if p.Payload != nil {
    267 		buf = append(buf, p.Payload...)
    268 	}
    269 	if len(buf) > PacketSize {
    270 		return fmt.Errorf("%w: %d bytes", ErrLongPacket, len(buf))
    271 	} else if len(buf) < PacketSize {
    272 		return fmt.Errorf("%w: %d bytes", ErrShortPacket, len(buf))
    273 	}
    274 	_, err := w.Write(buf)
    275 	return err
    276 }
    277 
    278 const (
    279 	baseMax      = 8589934592 - 1 // max 33-bit uint
    280 	extensionMax = 512 - 1        // max 9-bit uint
    281 )
    282 
    283 func putPCR(b []byte, pcr *PCR) error {
    284 	if len(b) != 6 {
    285 		return fmt.Errorf("need %d bytes, got %d", 6, len(b))
    286 	}
    287 	if pcr.Base > baseMax {
    288 		return fmt.Errorf("base %d larger than max %d", pcr.Base, baseMax)
    289 	} else if pcr.Extension > extensionMax {
    290 		return fmt.Errorf("extension %d larger than max %d", pcr.Extension, extensionMax)
    291 	}
    292 
    293 	ubuf := make([]byte, 8, 8)
    294 	binary.BigEndian.PutUint64(ubuf, pcr.Base)
    295 	// we're only working with 33 bits, so slice off the first 3
    296 	// bytes to get 4 + 1 bytes (32+1 bits)
    297 	ubuf = ubuf[3:]
    298 
    299 	// now pack 33 bits from ubuf into b[:4].
    300 	// The 33rd bit of our 33-bit integer is in the first byte: 0b00000001.
    301 	// We're packing bits from left to right, so shift it left and assign to b[0].
    302 	b[0] = ubuf[0] << 7
    303 
    304 	// We have 7 bits free in b[0], so get 7 bits from ubuf[1] and pack it into b[0].
    305 	b[0] |= ubuf[1] >> 1
    306 
    307 	// 1 bit left in ubuf[1]; put it in the next dest byte.
    308 	// Rinse and repeat until we're out of bits.
    309 	b[1] = ubuf[1] << 7
    310 	b[1] |= ubuf[2] >> 1
    311 	b[2] = ubuf[2] << 7
    312 	b[2] |= ubuf[3] >> 1
    313 	b[3] = ubuf[3] << 7
    314 	b[3] |= ubuf[4] >> 1
    315 	b[4] = ubuf[4] << 7
    316 	// No more base bits to pack.
    317 
    318 	// Next, toggle the 6 reserved bits.
    319 	b[4] |= 0b01111110
    320 
    321 	var ext [2]byte
    322 	binary.BigEndian.PutUint16(ext[:], pcr.Extension)
    323 	b[4] |= ext[0]
    324 	b[5] = ext[1]
    325 	return nil
    326 }