From 722ddc6ccb04b82c181ae24931969ce30943049d Mon Sep 17 00:00:00 2001 From: Saxon Date: Wed, 22 May 2019 14:49:25 +0930 Subject: [PATCH] codec/h264/lex_test.go: added test for RTPLexer.Lex(...) and also fixed some build errors --- codec/h264/lex.go | 39 +++++----- codec/h264/lex_test.go | 143 +++++++++++++++++++++++++++++++++++++ protocol/rtmp/rtmp_test.go | 4 +- revid/revid.go | 2 +- 4 files changed, 166 insertions(+), 22 deletions(-) diff --git a/codec/h264/lex.go b/codec/h264/lex.go index 14e0f63a..5f100c82 100644 --- a/codec/h264/lex.go +++ b/codec/h264/lex.go @@ -40,6 +40,25 @@ import ( "bitbucket.org/ausocean/av/protocol/rtp" ) +// NAL types. +const ( + // Single nal units bounds. + typeSingleNALULowBound = 1 + typeSingleNALUHighBound = 23 + + // Single-time aggregation packets. + typeSTAPA = 24 + typeSTAPB = 25 + + // Multi-time aggregation packets. + typeMTAP16 = 26 + typeMTAP24 = 27 + + // Fragmentation packets. + typeFUA = 28 + typeFUB = 29 +) + var noDelay = make(chan time.Time) func init() { @@ -138,24 +157,6 @@ outer: return err } -const ( - // Single nal units bounds. - typeSingleNALULowBound = 1 - typeSingleNALUHighBound = 23 - - // Single-time aggregation packets. - typeSTAPA = 24 - typeSTAPB = 25 - - // Multi-time aggregation packets. - typeMTAP16 = 26 - typeMTAP24 = 27 - - // Fragmentation packets. - typeFUA = 28 - typeFUB = 29 -) - // Buffer sizes. const ( maxAUSize = 100000 @@ -176,7 +177,7 @@ func NewRTPLexer() *RTPLexer { // LexFromRTP extracts H264 access units from an RTP stream. This function // expects that each read from src will provide a single RTP packet. -func (l *RTPLexer) LexFromRTP(dst io.Writer, src io.Reader, delay time.Duration) error { +func (l *RTPLexer) Lex(dst io.Writer, src io.Reader, delay time.Duration) error { buf := make([]byte, maxRTPSize) for { n, err := src.Read(buf) diff --git a/codec/h264/lex_test.go b/codec/h264/lex_test.go index d2eeae2a..3632268b 100644 --- a/codec/h264/lex_test.go +++ b/codec/h264/lex_test.go @@ -7,6 +7,7 @@ DESCRIPTION AUTHOR Dan Kortschak + Saxon A. Nelson-Milton LICENSE lex_test.go is Copyright (C) 2017 the Australian Ocean Lab (AusOcean) @@ -30,6 +31,8 @@ LICENSE package h264 import ( + "io" + "testing" "time" ) @@ -220,3 +223,143 @@ func TestH264(t *testing.T) { } } */ + +// rtpReader provides the RTP stream. +type rtpReader struct { + packets [][]byte + idx int +} + +// Read implements io.Reader. +func (r *rtpReader) Read(p []byte) (int, error) { + if r.idx == len(r.packets) { + return 0, io.EOF + } + b := r.packets[r.idx] + n := copy(p, b) + if n < len(r.packets[r.idx]) { + r.packets[r.idx] = r.packets[r.idx][n:] + } else { + r.idx++ + } + return n, nil +} + +// destination holds the access units extracted during the lexing process. +type destination [][]byte + +// Write implements io.Writer. +func (d *destination) Write(p []byte) (int, error) { + t := make([]byte, len(p)) + copy(t, p) + *d = append([][]byte(*d), t) + return len(p), nil +} + +// TestLex checks that the Lexer can correctly extract H264 access units from +// h264 RTP stream in RTP payload format. +func TestLex(t *testing.T) { + const rtpVer = 2 + + tests := []struct { + packets [][]byte + expect [][]byte + }{ + { + packets: [][]byte{ + { // Single NAL unit. + 0x80, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // RTP header. + typeSingleNALULowBound, // NAL header. + 0x01, 0x02, 0x03, 0x04, // NAL Data. + }, + { // Fragmentation (start packet). + 0x80, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // RTP header. + typeFUA, // FU indicator. + 0x80 | typeSingleNALULowBound, // FU header. + 0x01, 0x02, 0x03, // FU payload. + }, + { // Fragmentation (middle packet) + 0x80, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // RTP header. + typeFUA, // NAL indicator. + typeSingleNALULowBound, // FU header. + 0x04, 0x05, 0x06, // FU payload. + }, + { // Fragmentation (end packet) + 0x80, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // RTP header. + typeFUA, // NAL indicator. + 0x40 | typeSingleNALULowBound, // FU header. + 0x07, 0x08, 0x09, // FU payload + }, + + { // Aggregation. Make last packet of access unit => marker bit true. + 0x80, 0x80, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // RTP header. + typeSTAPA, // NAL header. + 0x00, 0x04, // NAL 1 size. + 0x01, 0x02, 0x03, 0x04, // NAL 1 data. + 0x00, 0x04, // NAL 2 size. + 0x01, 0x02, 0x03, 0x04, // NAL 2 data. + }, + // Second access unit. + { // Single NAL unit. + 0x80, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // RTP header. + typeSingleNALULowBound, // NAL header. + 0x01, 0x02, 0x03, 0x04, // NAL Data. + }, + { // Single NAL. Make last packet of access unit => marker bit true. + 0x80, 0x80, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // RTP header. + typeSingleNALULowBound, // NAL header. + 0x01, 0x02, 0x03, 0x04, // NAL data. + }, + }, + expect: [][]byte{ + // First access unit. + { + // NAL 1 + 0x00, 0x00, 0x00, 0x01, // Start code. + typeSingleNALULowBound, // NAL header. + 0x01, 0x02, 0x03, 0x04, // NAL data. + // NAL 2 + 0x00, 0x00, 0x00, 0x01, // Start code. + typeSingleNALULowBound, + 0x01, 0x02, 0x03, // FU payload. + 0x04, 0x05, 0x06, // FU payload. + 0x07, 0x08, 0x09, // FU payload. + // NAL 3 + 0x00, 0x00, 0x00, 0x01, // Start code. + 0x01, 0x02, 0x03, 0x04, // NAL data. + // NAL 4 + 0x00, 0x00, 0x00, 0x01, // Start code. + 0x01, 0x02, 0x03, 0x04, // NAL 2 data + }, + // Second access unit. + { + // NAL 1 + 0x00, 0x00, 0x00, 0x01, // Start code. + typeSingleNALULowBound, // NAL header. + 0x01, 0x02, 0x03, 0x04, // Data. + // NAL 2 + 0x00, 0x00, 0x00, 0x01, // Start code. + typeSingleNALULowBound, // NAL header. + 0x01, 0x02, 0x03, 0x04, // Data. + }, + }, + }, + } + + for testNum, test := range tests { + r := &rtpReader{packets: test.packets} + d := &destination{} + err := NewRTPLexer().Lex(d, r, 0) + if err != nil { + t.Fatalf("error lexing: %v\n", err) + } + + for i, accessUnit := range test.expect { + for j, part := range accessUnit { + if part != [][]byte(*d)[i][j] { + t.Fatalf("did not get expected data for test: %v.\nGot: %v\nWant: %v\n", testNum, d, test.expect) + } + } + } + } +} diff --git a/protocol/rtmp/rtmp_test.go b/protocol/rtmp/rtmp_test.go index e1e79796..c7c38bd4 100644 --- a/protocol/rtmp/rtmp_test.go +++ b/protocol/rtmp/rtmp_test.go @@ -199,7 +199,7 @@ func TestFromFrame(t *testing.T) { if err != nil { t.Errorf("Failed to create flv encoder with error: %v", err) } - err = h264.Lex(flvEncoder, bytes.NewReader(videoData), time.Second/time.Duration(frameRate)) + err = h264.LexFromBytestream(flvEncoder, bytes.NewReader(videoData), time.Second/time.Duration(frameRate)) if err != nil { t.Errorf("Lexing failed with error: %v", err) } @@ -251,7 +251,7 @@ func TestFromFile(t *testing.T) { if err != nil { t.Fatalf("failed to create encoder: %v", err) } - err = h264.Lex(flvEncoder, f, time.Second/time.Duration(25)) + err = h264.LexFromBytestream(flvEncoder, f, time.Second/time.Duration(25)) if err != nil { t.Errorf("Lexing and encoding failed with error: %v", err) } diff --git a/revid/revid.go b/revid/revid.go index 911d095f..c5d92c5a 100644 --- a/revid/revid.go +++ b/revid/revid.go @@ -272,7 +272,7 @@ func (r *Revid) setupPipeline(mtsEnc, flvEnc func(dst io.WriteCloser, rate int) switch r.config.InputCodec { case H264: r.config.Logger.Log(logger.Info, pkg+"using H264 lexer") - r.lexTo = h264.Lex + r.lexTo = h264.LexFromBytestream case Mjpeg: r.config.Logger.Log(logger.Info, pkg+"using MJPEG lexer") r.lexTo = mjpeg.Lex