1package cc23import (4 "bytes"5 "math/rand/v2"6 "slices"7 "testing"8)910// muxer makes MPEG-TS for tests, the way ffmpeg does: a PES per picture,11// the first packet's adaptation field holding the PCR, the last stuffed.12type muxer struct {13 out []byte14 cc map[int]byte15 ccs map[int][]byte // continuity counters sent, by PID16}1718func (m *muxer) packet(pid int, start bool, af, data []byte) {19 if m.cc == nil {20 m.cc, m.ccs = map[int]byte{}, map[int][]byte{}21 }22 p := []byte{tsSync, byte(pid >> 8 & 0x1f), byte(pid), 0x10 | m.cc[pid]}23 if start {24 p[1] |= 0x4025 }26 m.ccs[pid] = append(m.ccs[pid], m.cc[pid])27 m.cc[pid] = (m.cc[pid] + 1) & 0x0f28 if fill := tsPayload - len(data); fill > 0 {29 p[3] |= 0x2030 p = append(p, byte(fill-1))31 if fill > 1 {32 if af == nil {33 af = []byte{0}34 }35 p = append(p, af...)36 p = append(p, bytes.Repeat([]byte{0xff}, 4+fill-len(p))...)37 }38 }39 m.out = append(m.out, append(p, data...)...)40}4142// pes sends a PES, calling between after each packet but the last, to43// slip others in.44func (m *muxer) pes(pid int, pes, af []byte, between func()) {45 for first := true; len(pes) > 0; first = false {46 room := tsPayload47 var a []byte48 if first && af != nil {49 a = af50 room -= 1 + len(af)51 }52 n := min(room, len(pes))53 m.packet(pid, first, a, pes[:n])54 pes = pes[n:]55 if between != nil && len(pes) > 0 {56 between()57 }58 }59}6061func (m *muxer) psi(pid int, section []byte) {62 m.packet(pid, true, nil, append(append([]byte{0}, section...), bytes.Repeat([]byte{0xff}, tsPayload-1-len(section))...))63}6465const (66 videoPID = 0x10067 audioPID = 0x10168 pmtPID = 0x100069)7071func pat() []byte {72 return []byte{0x00, 0xb0, 13, 0, 1, 0xc1, 0, 0, 0, 1, 0xe0 | pmtPID>>8, pmtPID & 0xff, 0, 0, 0, 0}73}7475func pmt() []byte {76 return []byte{0x02, 0xb0, 23, 0, 1, 0xc1, 0, 0, 0xe1, 0x00, 0xf0, 0,77 0x1b, 0xe0 | videoPID>>8, videoPID & 0xff, 0xf0, 0,78 0x0f, 0xe0 | audioPID>>8, audioPID & 0xff, 0xf0, 0,79 0, 0, 0, 0}80}8182func pesHeader(pts int64, length int) []byte {83 h := []byte{0, 0, 1, 0xe0, byte(length >> 8), byte(length), 0x80, 0x80, 5,84 byte(pts>>29&0x0e | 0x21), byte(pts >> 22), byte(pts>>14 | 1), byte(pts >> 7), byte(pts<<1 | 1)}85 return h86}8788// picture is an access unit: an AUD, parameter sets for a keyframe, and a89// slice of n bytes.90func picture(rng *rand.Rand, key bool, n int) (es []byte, slice int) {91 es = []byte{0, 0, 0, 1, 0x09, 0xf0}92 if key {93 es = append(es, 0, 0, 0, 1, 0x67, 0x64, 0x00, 0x1f, 0xac, 0, 0, 0, 1, 0x68, 0xee, 0x3c, 0x80)94 }95 slice = len(es)96 nal := byte(0x41)97 if key {98 nal = 0x6599 }100 es = append(es, 0, 0, 1, nal)101 for range n {102 es = append(es, byte(0x10+rng.IntN(0xf0)))103 }104 return es, slice105}106107type sentPicture struct {108 pes []byte109 slice int // where the slice starts in pes110 frame int64111 pcr []byte112}113114// stream makes a TS of pictures in decode order (frames 0 3 1 2 6 4 5 ...)115// with audio between their packets.116func stream(t *testing.T, pictures int, sizes []int, lengths bool) ([]byte, []sentPicture, *muxer) {117 t.Helper()118 rng := rand.New(rand.NewPCG(7, 9))119 m := &muxer{}120 m.psi(0, pat())121 m.psi(pmtPID, pmt())122 audio := func() {123 if rng.IntN(3) == 0 {124 m.pes(audioPID, append([]byte{0, 0, 1, 0xc0, 0, 0, 0x80, 0x80, 5, 0x21, 0, 1, 0, 1}, byte(rng.IntN(256))), nil, nil)125 }126 }127 order := []int64{0}128 for g := int64(0); len(order) < pictures; g += 3 {129 order = append(order, g+3, g+1, g+2)130 }131 var sent []sentPicture132 for i, frame := range order[:pictures] {133 key := i%10 == 0134 es, slice := picture(rng, key, sizes[i%len(sizes)])135 length := 0136 if lengths {137 length = 8 + len(es)138 }139 pes := append(pesHeader(900000+frame*3000, length), es...)140 var af []byte141 if key {142 af = []byte{0x50, 0, 0, 0x2b, 0xf2, 0x7e, 0} // random access, PCR143 af[1] = byte(i)144 }145 sent = append(sent, sentPicture{pes: pes, slice: 14 + slice, frame: frame, pcr: af})146 m.pes(videoPID, pes, af, audio)147 audio()148 if i%7 == 0 {149 m.psi(0, pat())150 m.psi(pmtPID, pmt())151 }152 }153 return m.out, sent, m154}155156func pairFor(frame int64) (byte, byte) {157 return parity(byte(0x41 + frame%26)), parity(byte(0x61 + frame%26))158}159160type tsPacketInfo struct {161 pid int162 start bool163 cc byte164 af []byte // adaptation field contents, stuffing included165 data []byte166 raw []byte167}168169func readTS(t *testing.T, b []byte) []tsPacketInfo {170 t.Helper()171 if len(b)%tsPacket != 0 {172 t.Fatalf("%d bytes isn't whole packets", len(b))173 }174 var out []tsPacketInfo175 for ; len(b) > 0; b = b[tsPacket:] {176 p := b[:tsPacket]177 if p[0] != tsSync {178 t.Fatalf("packet %d lacks a sync byte", len(out))179 }180 info := tsPacketInfo{pid: int(p[1]&0x1f)<<8 | int(p[2]), start: p[1]&0x40 != 0, cc: p[3] & 0x0f, raw: p}181 rest := p[4:]182 if p[3]&0x20 != 0 {183 n := int(p[4])184 if n > 183 || (p[3]&0x10 == 0 && n != 183) {185 t.Fatalf("packet %d: adaptation field of %d", len(out), n)186 }187 info.af = p[5 : 5+n]188 rest = p[5+n:]189 }190 if p[3]&0x10 != 0 {191 info.data = rest192 } else if len(rest) > 0 {193 t.Fatalf("packet %d: %d bytes beyond its adaptation field", len(out), len(rest))194 }195 out = append(out, info)196 }197 return out198}199200func TestWriterAddsCaptions(t *testing.T) {201 // Slice sizes landing the PES exactly on packet boundaries, one byte202 // short of them, and in between.203 sizes := []int{2000, 140, 160, 161, 162, 165, 1, 345, 3, 5000}204 for _, lengths := range []bool{false, true} {205 in, sent, mux := stream(t, 40, sizes, lengths)206 var frames []int64207 var out bytes.Buffer208 w := NewWriter(&out, 30, func(f int64) (byte, byte) {209 frames = append(frames, f)210 return pairFor(f)211 })212 if _, err := w.Write(in); err != nil {213 t.Fatal(err)214 }215 if err := w.Flush(); err != nil {216 t.Fatal(err)217 }218 pkts := readTS(t, out.Bytes())219220 // Everything but the video is as it was, in the same order.221 var others, othersIn []byte222 for _, p := range pkts {223 if p.pid != videoPID {224 others = append(others, p.raw...)225 }226 }227 for _, p := range readTS(t, in) {228 if p.pid != videoPID {229 othersIn = append(othersIn, p.raw...)230 }231 }232 if !bytes.Equal(others, othersIn) {233 t.Error("other streams changed")234 }235236 // The video counts on unbroken from where it started.237 var want byte = mux.ccs[videoPID][0]238 var pes [][]byte239 var afs [][]byte240 for i, p := range pkts {241 if p.pid != videoPID {242 continue243 }244 if p.cc != want {245 t.Fatalf("packet %d: continuity counter %d, want %d", i, p.cc, want)246 }247 want = (want + 1) & 0x0f248 if p.start {249 pes = append(pes, nil)250 afs = append(afs, p.af)251 } else if len(p.af) > 0 && p.af[0]&0x10 != 0 {252 t.Errorf("packet %d: a PCR moved mid-picture", i)253 }254 pes[len(pes)-1] = append(pes[len(pes)-1], p.data...)255 }256 if len(pes) != len(sent) {257 t.Fatalf("%d pictures, want %d", len(pes), len(sent))258 }259 var wantFrames []int64260 for i, s := range sent {261 wantFrames = append(wantFrames, s.frame)262 sei := captionSEI(pairFor(s.frame))263 exp := slices.Concat(s.pes[:s.slice], sei, s.pes[s.slice:])264 if lengths {265 n := len(exp) - 6266 exp[4], exp[5] = byte(n>>8), byte(n)267 }268 if !bytes.Equal(pes[i], exp) {269 t.Errorf("picture %d (frame %d): PES of %d bytes differs from the %d expected", i, s.frame, len(pes[i]), len(exp))270 }271 if s.pcr != nil && !bytes.Equal(afs[i][:len(s.pcr)], s.pcr) {272 t.Errorf("picture %d: adaptation field % x, want % x", i, afs[i][:len(s.pcr)], s.pcr)273 }274 }275 if !slices.Equal(frames, wantFrames) {276 t.Errorf("frames asked for %v, want %v", frames, wantFrames)277 }278 }279}280281func TestWriterTakesAnyPieces(t *testing.T) {282 in, _, _ := stream(t, 25, []int{700, 90, 3000}, false)283 var whole bytes.Buffer284 w := NewWriter(&whole, 30, pairFor)285 w.Write(in)286 w.Flush()287 rng := rand.New(rand.NewPCG(3, 4))288 for _, size := range []int{1, 7, 187, 188, 189, 4096, -1} {289 var out bytes.Buffer290 w := NewWriter(&out, 30, pairFor)291 for p := in; len(p) > 0; {292 n := size293 if n < 0 {294 n = 1 + rng.IntN(600)295 }296 n = min(n, len(p))297 if k, err := w.Write(p[:n]); err != nil || k != n {298 t.Fatalf("Write = %d, %v", k, err)299 }300 p = p[n:]301 }302 w.Flush()303 if !bytes.Equal(out.Bytes(), whole.Bytes()) {304 t.Errorf("written %d bytes at a time: output differs", size)305 }306 }307}308309func TestWriterPassesOtherStreams(t *testing.T) {310 // Without a PAT and PMT naming H.264 video, nothing changes.311 m := &muxer{}312 rng := rand.New(rand.NewPCG(5, 6))313 es, _ := picture(rng, true, 900)314 m.pes(videoPID, append(pesHeader(1000, 0), es...), nil, nil)315 m.pes(audioPID, []byte{0, 0, 1, 0xc0, 0, 3, 0x80, 0, 0}, nil, nil)316 var out bytes.Buffer317 w := NewWriter(&out, 30, pairFor)318 w.Write(m.out)319 w.Flush()320 if !bytes.Equal(out.Bytes(), m.out) {321 t.Error("a stream without a PMT changed")322 }323 // Bytes out of step with packets pass through too.324 out.Reset()325 w = NewWriter(&out, 30, pairFor)326 junk := append([]byte("junk"), m.out...)327 w.Write(junk)328 w.Flush()329 if !bytes.Equal(out.Bytes(), junk) {330 t.Error("junk changed")331 }332}333334func TestEscape(t *testing.T) {335 got := escape([]byte{0, 0, 0, 0, 0, 1, 0, 0, 3, 0, 0, 4})336 if want := []byte{0, 0, 3, 0, 0, 3, 0, 1, 0, 0, 3, 3, 0, 0, 4}; !bytes.Equal(got, want) {337 t.Errorf("escape = % x, want % x", got, want)338 }339}340341func TestCaptionSEI(t *testing.T) {342 got := captionSEI(0x94, 0x2c)343 want := []byte{0, 0, 0, 1, 0x06, 0x04, 17, 0xb5, 0x00, 0x31, 'G', 'A', '9', '4', 0x03, 0xc2, 0xff,344 0xfc, 0x94, 0x2c, 0xfd, 0x80, 0x80, 0xff, 0x80}345 if !bytes.Equal(got, want) {346 t.Errorf("SEI\n% x, want\n% x", got, want)347 }348}