Merged in h264-rtp-lexer (pull request #202)

codec/h264: added RTPLexer

Approved-by: Alan Noble <anoble@gmail.com>

codec/h264/extract.go

@@ -0,0 +1,217 @@
+/*
+NAME
+  extract.go
+
+DESCRIPTION
+  extract.go provides an extracter to get access units from an RTP stream.
+
+AUTHOR
+  Saxon Nelson-Milton <saxon@ausocean.org>
+
+LICENSE
+  Copyright (C) 2017 the Australian Ocean Lab (AusOcean)
+
+  It is free software: you can redistribute it and/or modify them
+  under the terms of the GNU General Public License as published by the
+  Free Software Foundation, either version 3 of the License, or (at your
+  option) any later version.
+
+  It is distributed in the hope that it will be useful, but WITHOUT
+  ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+  FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License
+  for more details.
+
+  You should have received a copy of the GNU General Public License
+  along with revid in gpl.txt. If not, see http://www.gnu.org/licenses.
+*/
+package h264
+
+import (
+	"bytes"
+	"encoding/binary"
+	"fmt"
+	"io"
+	"time"
+
+	"bitbucket.org/ausocean/av/protocol/rtp"
+)
+
+// NAL types (from https://tools.ietf.org/html/rfc6184#page-13)
+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
+)
+
+// Min NAL lengths.
+const (
+	minSingleNALLen = 1
+	minSTAPALen     = 4
+	minFUALen       = 2
+)
+
+// Buffer sizes.
+const (
+	maxAUSize  = 100000 // Max access unit size in bytes.
+	maxRTPSize = 1500   // Max ethernet transmission unit in bytes.
+)
+
+// Extracter is an extracter for extracting H264 access units from RTP stream.
+type Extracter struct {
+	buf  *bytes.Buffer // Holds the current access unit.
+	frag bool          // Indicates if we're currently dealing with a fragmentation packet.
+}
+
+// NewExtracter returns a new Extracter.
+func NewExtracter() *Extracter {
+	return &Extracter{
+		buf: bytes.NewBuffer(make([]byte, 0, maxAUSize))}
+}
+
+// Extract extracts H264 access units from an RTP stream. This function
+// expects that each read from src will provide a single RTP packet.
+func (e *Extracter) Extract(dst io.Writer, src io.Reader, delay time.Duration) error {
+	buf := make([]byte, maxRTPSize)
+	for {
+		n, err := src.Read(buf)
+		switch err {
+		case nil: // Do nothing.
+		case io.EOF:
+			return nil
+		default:
+			return fmt.Errorf("source read error: %v\n", err)
+		}
+
+		// Get payload from RTP packet.
+		payload, err := rtp.Payload(buf[:n])
+		if err != nil {
+			return fmt.Errorf("could not get RTP payload, failed with err: %v\n", err)
+		}
+
+		nalType := payload[0] & 0x1f
+
+		// If not currently fragmented then we ignore current write.
+		if e.frag && nalType != typeFUA {
+			e.buf.Reset()
+			e.frag = false
+			continue
+		}
+
+		if typeSingleNALULowBound <= nalType && nalType <= typeSingleNALUHighBound {
+			// If len too small, ignore.
+			if len(payload) < minSingleNALLen {
+				continue
+			}
+			e.writeWithPrefix(payload)
+		} else {
+			switch nalType {
+			case typeSTAPA:
+				e.handleSTAPA(payload)
+			case typeFUA:
+				e.handleFUA(payload)
+			case typeSTAPB:
+				panic("STAP-B type unsupported")
+			case typeMTAP16:
+				panic("MTAP16 type unsupported")
+			case typeMTAP24:
+				panic("MTAP24 type unsupported")
+			case typeFUB:
+				panic("FU-B type unsupported")
+			default:
+				panic("unsupported type")
+			}
+		}
+
+		markerIsSet, err := rtp.Marker(buf[:n])
+		if err != nil {
+			return fmt.Errorf("could not get marker bit, failed with err: %v\n", err)
+		}
+
+		if markerIsSet {
+			e.buf.WriteTo(dst)
+			e.buf.Reset()
+		}
+	}
+	return nil
+}
+
+// handleSTAPA parses NAL units from an aggregation packet and writes
+// them to the Extracter's buffer buf.
+func (e *Extracter) handleSTAPA(d []byte) {
+	// If the length is too small, ignore.
+	if len(d) < minSTAPALen {
+		return
+	}
+
+	for i := 1; i < len(d); {
+		size := int(binary.BigEndian.Uint16(d[i:]))
+
+		// Skip over NAL unit size.
+		const sizeOfFieldLen = 2
+		i += sizeOfFieldLen
+
+		// Get the NALU.
+		nalu := d[i : i+size]
+		i += size
+		e.writeWithPrefix(nalu)
+	}
+}
+
+// handleFUA parses NAL units from fragmentation packets and writes
+// them to the Extracter's buf.
+func (e *Extracter) handleFUA(d []byte) {
+	// If length is too small, ignore.
+	if len(d) < minFUALen {
+		return
+	}
+
+	// Get start and end indiciators from FU header.
+	const FUHeadIdx = 1
+	start := d[FUHeadIdx]&0x80 != 0
+	end := d[FUHeadIdx]&0x40 != 0
+
+	// If start, form new header, skip FU indicator only and set first byte to
+	// new header. Otherwise, skip over both FU indicator and FU header.
+	if start {
+		newHead := (d[0] & 0xe0) | (d[1] & 0x1f)
+		d = d[1:]
+		d[0] = newHead
+		if end {
+			panic("bad fragmentation packet")
+		}
+		e.frag = true
+		e.writeWithPrefix(d)
+	} else {
+		d = d[2:]
+		if end {
+			e.frag = false
+		}
+		e.writeNoPrefix(d)
+	}
+}
+
+// write writes a NAL unit to the Extracter's buf in byte stream format using the
+// start code.
+func (e *Extracter) writeWithPrefix(d []byte) {
+	const prefix = "\x00\x00\x00\x01"
+	e.buf.Write([]byte(prefix))
+	e.buf.Write(d)
+}
+
+// writeNoPrefix writes data to the Extracter's buf. This is used for non start
+// fragmentations of a NALU.
+func (e *Extracter) writeNoPrefix(d []byte) {
+	e.buf.Write(d)
+}

codec/h264/extract_test.go

@@ -0,0 +1,173 @@
+/*
+NAME
+  extract_test.go
+
+DESCRIPTION
+  extract_test.go provides tests for the extracter in extract.go
+
+AUTHOR
+  Saxon A. Nelson-Milton <saxon@ausocean.org>
+
+LICENSE
+  Copyright (C) 2017 the Australian Ocean Lab (AusOcean)
+
+  It is free software: you can redistribute it and/or modify them
+  under the terms of the GNU General Public License as published by the
+  Free Software Foundation, either version 3 of the License, or (at your
+  option) any later version.
+
+  It is distributed in the hope that it will be useful, but WITHOUT
+  ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+  FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License
+  for more details.
+
+  You should have received a copy of the GNU General Public License
+  along with revid in gpl.txt. If not, see http://www.gnu.org/licenses.
+*/
+
+package h264
+
+import (
+	"io"
+	"testing"
+)
+
+// rtpReader provides an io.Reader for reading the test 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) {
+	tmp := make([]byte, len(p))
+	copy(tmp, p)
+	*d = append(*d, tmp)
+	return len(p), nil
+}
+
+// TestLex checks that the Lexer can correctly extract H264 access units from
+// h264 RTP stream in RTP payload format.
+func TestRTPLex(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 := NewExtracter().Extract(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)
+				}
+			}
+		}
+	}
+}

codec/h264/lex.go

@@ -30,10 +30,40 @@ LICENSE
 package h264
 
 import (
+	"bytes"
+	"encoding/binary"
+	"fmt"
 	"io"
 	"time"
 
 	"bitbucket.org/ausocean/av/codec/codecutil"
+	"bitbucket.org/ausocean/av/protocol/rtp"
+)
+
+// NAL types (from https://tools.ietf.org/html/rfc6184#page-13)
+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
+)
+
+// Min NAL lengths.
+const (
+	minSingleNALLen = 1
+	minSTAPALen     = 4
+	minFUALen       = 2
 )
 
 var noDelay = make(chan time.Time)
@@ -44,11 +74,11 @@ func init() {
 
 var h264Prefix = [...]byte{0x00, 0x00, 0x01, 0x09, 0xf0}
 
-// Lex lexes H.264 NAL units read from src into separate writes to dst with
+// LexFromBytestream lexes H.264 NAL units read from src into separate writes
-// successive writes being performed not earlier than the specified delay.
+// to dst with successive writes being performed not earlier than the specified
-// NAL units are split after type 1 (Coded slice of a non-IDR picture), 5
+// delay. NAL units are split after type 1 (Coded slice of a non-IDR picture), 5
 // (Coded slice of a IDR picture) and 8 (Picture parameter set).
-func Lex(dst io.Writer, src io.Reader, delay time.Duration) error {
+func LexFromBytestream(dst io.Writer, src io.Reader, delay time.Duration) error {
 	var tick <-chan time.Time
 	if delay == 0 {
 		tick = noDelay
@@ -133,3 +163,157 @@ outer:
 	_, err := dst.Write(buf)
 	return err
 }
+
+// Buffer sizes.
+const (
+	maxAUSize  = 100000 // Max access unit size in bytes.
+	maxRTPSize = 1500   // Max ethernet transmission unit in bytes.
+)
+
+// RTPLexer is a lexer for lexing H264 from RTP packets.
+type RTPLexer struct {
+	buf  *bytes.Buffer // Holds the current access unit.
+	frag bool          // Indicates if we're currently dealing with a fragmentation packet.
+}
+
+// NewRTPLexer returns a new RTPLexer.
+func NewRTPLexer() *RTPLexer {
+	return &RTPLexer{
+		buf: bytes.NewBuffer(make([]byte, 0, maxAUSize))}
+}
+
+// Lex 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) Lex(dst io.Writer, src io.Reader, delay time.Duration) error {
+	buf := make([]byte, maxRTPSize)
+	for {
+		n, err := src.Read(buf)
+		switch err {
+		case nil: // Do nothing.
+		case io.EOF:
+			return nil
+		default:
+			return fmt.Errorf("source read error: %v\n", err)
+		}
+
+		// Get payload from RTP packet.
+		payload, err := rtp.Payload(buf[:n])
+		if err != nil {
+			return fmt.Errorf("could not get RTP payload, failed with err: %v\n", err)
+		}
+
+		nalType := payload[0] & 0x1f
+
+		// If not currently fragmented then we ignore current write.
+		if l.frag && nalType != typeFUA {
+			l.buf.Reset()
+			l.frag = false
+			continue
+		}
+
+		if typeSingleNALULowBound <= nalType && nalType <= typeSingleNALUHighBound {
+			// If len too small, ignore.
+			if len(payload) < minSingleNALLen {
+				continue
+			}
+			l.writeWithPrefix(payload)
+		} else {
+			switch nalType {
+			case typeSTAPA:
+				l.handleSTAPA(payload)
+			case typeFUA:
+				l.handleFUA(payload)
+			case typeSTAPB:
+				panic("STAP-B type unsupported")
+			case typeMTAP16:
+				panic("MTAP16 type unsupported")
+			case typeMTAP24:
+				panic("MTAP24 type unsupported")
+			case typeFUB:
+				panic("FU-B type unsupported")
+			default:
+				panic("unsupported type")
+			}
+		}
+
+		markerIsSet, err := rtp.Marker(buf[:n])
+		if err != nil {
+			return fmt.Errorf("could not get marker bit, failed with err: %v\n", err)
+		}
+
+		if markerIsSet {
+			l.buf.WriteTo(dst)
+			l.buf.Reset()
+		}
+	}
+	return nil
+}
+
+// handleSTAPA parses NAL units from an aggregation packet and writes
+// them to the Lexers buffer buf.
+func (l *RTPLexer) handleSTAPA(d []byte) {
+	// If the length is too small, ignore.
+	if len(d) < minSTAPALen {
+		return
+	}
+
+	for i := 1; i < len(d); {
+		size := int(binary.BigEndian.Uint16(d[i:]))
+
+		// Skip over NAL unit size.
+		const sizeOfFieldLen = 2
+		i += sizeOfFieldLen
+
+		// Get the NALU.
+		nalu := d[i : i+size]
+		i += size
+		l.writeWithPrefix(nalu)
+	}
+}
+
+// handleFUA parses NAL units from fragmentation packets and writes
+// them to the Lexer's buf.
+func (l *RTPLexer) handleFUA(d []byte) {
+	// If length is too small, ignore.
+	if len(d) < minFUALen {
+		return
+	}
+
+	// Get start and end indiciators from FU header.
+	const FUHeadIdx = 1
+	start := d[FUHeadIdx]&0x80 != 0
+	end := d[FUHeadIdx]&0x40 != 0
+
+	// If start, form new header, skip FU indicator only and set first byte to
+	// new header. Otherwise, skip over both FU indicator and FU header.
+	if start {
+		newHead := (d[0] & 0xe0) | (d[1] & 0x1f)
+		d = d[1:]
+		d[0] = newHead
+		if end {
+			panic("bad fragmentation packet")
+		}
+		l.frag = true
+		l.writeWithPrefix(d)
+	} else {
+		d = d[2:]
+		if end {
+			l.frag = false
+		}
+		l.writeNoPrefix(d)
+	}
+}
+
+// write writes a NAL unit to the Lexer's buf in byte stream format using the
+// start code.
+func (l *RTPLexer) writeWithPrefix(d []byte) {
+	const prefix = "\x00\x00\x00\x01"
+	l.buf.Write([]byte(prefix))
+	l.buf.Write(d)
+}
+
+// writeNoPrefix writes data to the Lexer's buf. This is used for non start
+// fragmentations of a NALU.
+func (l *RTPLexer) writeNoPrefix(d []byte) {
+	l.buf.Write(d)
+}

codec/h264/lex_test.go

@@ -7,6 +7,7 @@ DESCRIPTION
 
 AUTHOR
   Dan Kortschak <dan@ausocean.org>
+  Saxon A. Nelson-Milton <saxon@ausocean.org>
 
 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 an io.Reader for reading the test 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) {
+	tmp := make([]byte, len(p))
+	copy(tmp, p)
+	*d = append(*d, tmp)
+	return len(p), nil
+}
+
+// TestLex checks that the Lexer can correctly extract H264 access units from
+// h264 RTP stream in RTP payload format.
+func TestRTPLex(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)
+				}
+			}
+		}
+	}
+}

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)
 	}

revid/revid.go

@@ -280,13 +280,12 @@ func (r *Revid) setupPipeline(mtsEnc, flvEnc func(dst io.WriteCloser, rate int)
 	switch r.config.Input {
 	case Raspivid:
 		r.setupInput = r.startRaspivid
-		r.lexTo = h264.Lex
+		r.lexTo = h264.LexFromBytestream
 	case V4L:
 		r.setupInput = r.startV4L
-		r.lexTo = h264.Lex
+		r.lexTo = h264.LexFromBytestream
 	case File:
 		r.setupInput = r.setupInputForFile
-		r.lexTo = h264.Lex
 	case RTSP:
 		r.setupInput = r.startRTSPCamera
 		r.lexTo = h265.NewLexer(false).Lex