codec/h264: renamed RTPLexer to Extracter

Renamed the RTPLexer to Extracter, renamed NewRTPLexer to NewExtracter and renamed Lex to Extract. Put Extracter and accompanying methods in file extract.go.
Put tests relating to Extracter in extract_test.go.
This commit is contained in:
Saxon 2019-06-07 01:27:31 +09:30
parent aad0c54039
commit fa128d1809
2 changed files with 390 additions and 0 deletions

217
codec/h264/extract.go Normal file
View File

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

173
codec/h264/extract_test.go Normal file
View File

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