1package cc23import (4 "bytes"5 "io"6 "math"7 "slices"8)910const (11 tsPacket = 18812 tsPayload = 18413 tsSync = 0x4714 // maxHeld bounds what is held back waiting for a picture to end; past15 // it the picture goes out without a caption.16 maxHeld = 8 << 2017)1819// Writer adds CEA-608 captions to the H.264 picture in an MPEG-TS stream20// on its way to another writer, the way ATSC broadcasts carry them: each21// picture gets an SEI with the field 1 byte pair for its frame. Everything22// else passes through as it was. A picture is held back until the next23// begins, since only then is it known to be whole.24type Writer struct {25 w io.Writer26 fps int27 pair func(frame int64) (byte, byte)2829 part []byte // a packet split across writes30 out []byte31 pmts map[int]bool32 video int // the H.264 stream's PID, -1 until the PMT names it33 cc byte // its next continuity counter34 ccSet bool35 first int64 // the first picture's PTS: frame 036 started bool3738 // The picture being collected: its payload and its packets' adaptation39 // fields, and every packet since it began, nil standing for its own.40 // whole is set when pes is from its start; a picture too big to hold41 // goes out in parts.42 pes []byte43 slots []slot44 held [][]byte45 size int46 whole bool47 err error48}4950// slot is one of a picture's packets, minus its payload.51type slot struct {52 hdr [3]byte // the header after the sync byte53 af []byte // the adaptation field from its flags on, without stuffing54}5556// NewWriter returns a Writer to w for video at fps frames a second.57// pair gives the field 1 byte pair, with parity, for each frame, counted58// from the first picture by presentation time.59func NewWriter(w io.Writer, fps int, pair func(frame int64) (byte, byte)) *Writer {60 return &Writer{w: w, fps: fps, pair: pair, pmts: map[int]bool{}, video: -1}61}6263// Write takes MPEG-TS in pieces of any size.64func (x *Writer) Write(p []byte) (int, error) {65 if x.err != nil {66 return 0, x.err67 }68 n := len(p)69 if len(x.part) > 0 {70 k := min(tsPacket-len(x.part), len(p))71 x.part = append(x.part, p[:k]...)72 p = p[k:]73 if len(x.part) == tsPacket {74 x.packet(x.part)75 x.part = x.part[:0]76 }77 }78 for len(p) > 0 {79 if p[0] != tsSync {80 // Out of step: pass bytes through to the next sync byte.81 i := bytes.IndexByte(p, tsSync)82 if i < 0 {83 i = len(p)84 }85 x.other(p[:i])86 p = p[i:]87 continue88 }89 if len(p) < tsPacket {90 x.part = append(x.part, p...)91 break92 }93 x.packet(p[:tsPacket])94 p = p[tsPacket:]95 }96 if err := x.send(); err != nil {97 return 0, err98 }99 return n, nil100}101102// Flush sends the picture held back, for the end of the stream.103func (x *Writer) Flush() error {104 if x.err != nil {105 return x.err106 }107 x.finish(true)108 if len(x.part) > 0 {109 x.out = append(x.out, x.part...)110 x.part = x.part[:0]111 }112 return x.send()113}114115func (x *Writer) send() error {116 if len(x.out) == 0 {117 return nil118 }119 _, err := x.w.Write(x.out)120 x.out = x.out[:0]121 if err != nil {122 x.err = err123 }124 return err125}126127// other passes on bytes that aren't the picture's, in turn.128func (x *Writer) other(b []byte) {129 if len(x.slots) == 0 {130 x.out = append(x.out, b...)131 return132 }133 x.held = append(x.held, slices.Clone(b))134 x.size += len(b)135}136137func (x *Writer) packet(p []byte) {138 pid := int(p[1]&0x1f)<<8 | int(p[2])139 start := p[1]&0x40 != 0140 switch {141 case pid == 0 && start:142 x.readPAT(payload(p))143 case x.pmts[pid] && start:144 x.readPMT(payload(p))145 }146 scrambled := p[3]&0xc0 != 0147 if pid != x.video || scrambled || (!start && !x.ccSet) {148 x.other(p)149 return150 }151 if start {152 x.finish(true)153 x.whole = true154 }155 if !x.ccSet {156 x.cc, x.ccSet = p[3]&0x0f, true157 }158 s := slot{hdr: [3]byte{p[1], p[2], p[3]}}159 body := p[4:]160 if p[3]&0x20 != 0 {161 n := min(int(p[4]), tsPayload-1)162 s.af = adaptation(p[5 : 5+n])163 body = p[5+n:]164 }165 if p[3]&0x10 == 0 {166 body = nil167 }168 x.pes = append(x.pes, body...)169 x.slots = append(x.slots, s)170 x.held = append(x.held, nil)171 if x.size += tsPacket; x.size > maxHeld {172 x.finish(false)173 }174}175176// payload is a packet's payload, nil if it has none.177func payload(p []byte) []byte {178 if p[3]&0x10 == 0 {179 return nil180 }181 if p[3]&0x20 == 0 {182 return p[4:]183 }184 if int(p[4]) > tsPayload-1 {185 return nil186 }187 return p[5+int(p[4]):]188}189190// adaptation returns an adaptation field's contents (after its length)191// without the stuffing that pads it, nil when it holds nothing else.192func adaptation(af []byte) []byte {193 if len(af) == 0 || af[0] == 0 {194 return nil195 }196 n := 1197 flags := af[0]198 if flags&0x10 != 0 { // PCR199 n += 6200 }201 if flags&0x08 != 0 { // OPCR202 n += 6203 }204 if flags&0x04 != 0 { // splice countdown205 n++206 }207 if flags&0x02 != 0 && n < len(af) { // private data208 n += 1 + int(af[n])209 }210 if flags&0x01 != 0 && n < len(af) { // extension211 n += 1 + int(af[n])212 }213 return slices.Clone(af[:min(n, len(af))])214}215216// finish sends the picture collected so far, with its caption when caption217// is set, and the packets held back with it.218func (x *Writer) finish(caption bool) {219 if len(x.slots) == 0 {220 return221 }222 pes := x.pes223 if caption && x.whole {224 pes = x.caption(pes)225 }226 pkts := x.packetize(pes)227 i := 0228 for _, b := range x.held {229 if b == nil {230 b = pkts[i]231 i++232 }233 x.out = append(x.out, b...)234 }235 x.pes, x.slots, x.held, x.size, x.whole = x.pes[:0], x.slots[:0], x.held[:0], 0, false236}237238// packetize splits a picture's PES into packets in its slots, keeping each239// packet's adaptation field (the first's PCR and random access flag) where240// it was. What the caption added overflows into packets after the last.241func (x *Writer) packetize(pes []byte) [][]byte {242 out := make([][]byte, len(x.slots))243 for i, s := range x.slots {244 last := i == len(x.slots)-1245 for first := true; first || (last && len(pes) > 0); first = false {246 hdr, af := s.hdr, s.af247 if !first {248 hdr[0] &^= 0x40 // not a start249 af = nil250 }251 room := tsPayload252 if len(af) > 0 {253 room -= 1 + len(af)254 }255 n := min(room, len(pes))256 if n == 0 && len(af) == 0 {257 break258 }259 out[i] = append(out[i], x.build(hdr, af, pes[:n])...)260 pes = pes[n:]261 }262 }263 return out264}265266// build makes a packet, stuffing its adaptation field to fill it.267func (x *Writer) build(hdr [3]byte, af, data []byte) []byte {268 p := make([]byte, 4, tsPacket)269 p[0], p[1], p[2] = tsSync, hdr[0], hdr[1]270 ctl := hdr[2] & 0xc0 // scrambling271 if len(data) > 0 {272 ctl |= 0x10 | x.cc273 x.cc = (x.cc + 1) & 0x0f274 } else {275 ctl |= (x.cc - 1) & 0x0f // unchanged without a payload276 }277 if fill := tsPayload - len(data); fill > 0 {278 ctl |= 0x20279 p = append(p, byte(fill-1))280 if fill > 1 {281 if len(af) == 0 {282 af = []byte{0}283 }284 p = append(p, af...)285 for len(p) < 4+fill {286 p = append(p, 0xff)287 }288 }289 }290 p[3] = ctl291 return append(p, data...)292}293294// caption adds the SEI for its frame to a picture's PES, before the first295// slice, returning pes as it is when it isn't H.264 with a time.296func (x *Writer) caption(pes []byte) []byte {297 if len(pes) < 14 || pes[0] != 0 || pes[1] != 0 || pes[2] != 1 || pes[3]&0xf0 != 0xe0 || pes[7]&0x80 == 0 || pes[8] < 5 {298 return pes299 }300 es := 9 + int(pes[8])301 if es > len(pes) {302 return pes303 }304 at := firstSlice(pes[es:])305 if at < 0 {306 return pes307 }308 pts := int64(pes[9]>>1&7)<<30 | int64(pes[10])<<22 | int64(pes[11]>>1)<<15 | int64(pes[12])<<7 | int64(pes[13]>>1)309 if !x.started {310 x.first, x.started = pts, true311 }312 d := (pts - x.first) & (1<<33 - 1)313 if d >= 1<<32 {314 d -= 1 << 33 // before the first: B-frames315 }316 b1, b2 := x.pair(int64(math.Round(float64(d) * float64(x.fps) / 90000)))317 sei := captionSEI(b1, b2)318 out := make([]byte, 0, len(pes)+len(sei))319 out = append(append(append(out, pes[:es+at]...), sei...), pes[es+at:]...)320 if n := int(pes[4])<<8 | int(pes[5]); n != 0 {321 if n += len(sei); n > 0xffff {322 n = 0 // allowed for video323 }324 out[4], out[5] = byte(n>>8), byte(n)325 }326 return out327}328329// firstSlice finds the start code of the first coded slice in H.264, -1330// if there is none.331func firstSlice(es []byte) int {332 for i := 0; i+3 < len(es); i++ {333 if es[i] != 0 || es[i+1] != 0 || es[i+2] != 1 {334 continue335 }336 if t := es[i+3] & 0x1f; t >= 1 && t <= 5 {337 if i > 0 && es[i-1] == 0 {338 return i - 1339 }340 return i341 }342 i += 2343 }344 return -1345}346347// captionSEI is an SEI NAL unit carrying a pair of CEA-608 field 1 bytes348// (and a field 2 pair of padding) as ATSC A/53 cc_data in registered user349// data.350func captionSEI(b1, b2 byte) []byte {351 data := []byte{352 0xb5, 0x00, 0x31, // ITU-T T.35: United States, ATSC353 'G', 'A', '9', '4', 0x03, // user_identifier, user_data_type_code: cc_data354 0xc2, 0xff, // process_cc_data_flag, cc_count 2, em_data355 0xfc, b1, b2, // field 1 (CEA-608), valid356 0xfd, Padding, Padding, // field 2, padding357 0xff, // marker_bits358 }359 rbsp := append([]byte{4, byte(len(data))}, data...) // user_data_registered_itu_t_t35360 rbsp = append(rbsp, 0x80) // rbsp_trailing_bits361 return append([]byte{0, 0, 0, 1, 0x06}, escape(rbsp)...)362}363364// escape adds H.264's emulation prevention bytes, so the NAL unit's365// payload never looks like a start code.366func escape(b []byte) []byte {367 out := make([]byte, 0, len(b)+4)368 zeros := 0369 for _, c := range b {370 if zeros >= 2 && c <= 3 {371 out = append(out, 3)372 zeros = 0373 }374 out = append(out, c)375 if c == 0 {376 zeros++377 } else {378 zeros = 0379 }380 }381 return out382}383384// section returns the PSI section starting in a packet's payload, nil if385// it isn't all there.386func section(p []byte) []byte {387 if len(p) < 1 || 1+int(p[0])+3 > len(p) {388 return nil389 }390 s := p[1+int(p[0]):]391 n := int(s[1]&0x0f)<<8 | int(s[2])392 if n < 9 || 3+n > len(s) {393 return nil394 }395 return s[:3+n]396}397398// readPAT notes the programs' PMT PIDs.399func (x *Writer) readPAT(p []byte) {400 s := section(p)401 if s == nil || s[0] != 0 {402 return403 }404 body := s[8 : len(s)-4]405 for i := 0; i+4 <= len(body); i += 4 {406 if prog := int(body[i])<<8 | int(body[i+1]); prog != 0 {407 x.pmts[int(body[i+2]&0x1f)<<8|int(body[i+3])] = true408 }409 }410}411412// readPMT finds the H.264 stream.413func (x *Writer) readPMT(p []byte) {414 s := section(p)415 if s == nil || s[0] != 2 || len(s) < 16 || x.video >= 0 {416 return417 }418 body := s[12 : len(s)-4]419 info := int(s[10]&0x0f)<<8 | int(s[11])420 if info > len(body) {421 return422 }423 for body = body[info:]; len(body) >= 5; {424 if body[0] == 0x1b { // H.264425 x.video = int(body[1]&0x1f)<<8 | int(body[2])426 return427 }428 n := 5 + (int(body[3]&0x0f)<<8 | int(body[4]))429 if n > len(body) {430 return431 }432 body = body[n:]433 }434}