mirror of https://bitbucket.org/ausocean/av.git
299 lines
6.8 KiB
Go
299 lines
6.8 KiB
Go
/*
|
|
NAME
|
|
lex.go
|
|
|
|
DESCRIPTION
|
|
lex.go provides a lexer to lex h264 bytestream into access units.
|
|
|
|
AUTHOR
|
|
Dan Kortschak <dan@ausocean.org>
|
|
|
|
LICENSE
|
|
lex.go is 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.
|
|
*/
|
|
|
|
// lex.go provides a lexer to lex h264 bytestream into access units.
|
|
|
|
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
|
|
)
|
|
|
|
var noDelay = make(chan time.Time)
|
|
|
|
func init() {
|
|
close(noDelay)
|
|
}
|
|
|
|
var h264Prefix = [...]byte{0x00, 0x00, 0x01, 0x09, 0xf0}
|
|
|
|
// LexFromBytestream lexes H.264 NAL units read from src into separate writes
|
|
// to dst with successive writes being performed not earlier than the specified
|
|
// 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 LexFromBytestream(dst io.Writer, src io.Reader, delay time.Duration) error {
|
|
var tick <-chan time.Time
|
|
if delay == 0 {
|
|
tick = noDelay
|
|
} else {
|
|
ticker := time.NewTicker(delay)
|
|
defer ticker.Stop()
|
|
tick = ticker.C
|
|
}
|
|
|
|
const bufSize = 8 << 10
|
|
|
|
c := codecutil.NewByteScanner(src, make([]byte, 4<<10)) // Standard file buffer size.
|
|
|
|
buf := make([]byte, len(h264Prefix), bufSize)
|
|
copy(buf, h264Prefix[:])
|
|
writeOut := false
|
|
outer:
|
|
for {
|
|
var b byte
|
|
var err error
|
|
buf, b, err = c.ScanUntil(buf, 0x00)
|
|
if err != nil {
|
|
if err != io.EOF {
|
|
return err
|
|
}
|
|
break
|
|
}
|
|
|
|
for n := 1; b == 0x0 && n < 4; n++ {
|
|
b, err = c.ReadByte()
|
|
if err != nil {
|
|
if err != io.EOF {
|
|
return err
|
|
}
|
|
break outer
|
|
}
|
|
buf = append(buf, b)
|
|
|
|
if b != 0x1 || (n != 2 && n != 3) {
|
|
continue
|
|
}
|
|
|
|
if writeOut {
|
|
<-tick
|
|
_, err := dst.Write(buf[:len(buf)-(n+1)])
|
|
if err != nil {
|
|
return err
|
|
}
|
|
buf = make([]byte, len(h264Prefix)+n, bufSize)
|
|
copy(buf, h264Prefix[:])
|
|
buf = append(buf, 1)
|
|
writeOut = false
|
|
}
|
|
|
|
b, err = c.ReadByte()
|
|
if err != nil {
|
|
if err != io.EOF {
|
|
return err
|
|
}
|
|
break outer
|
|
}
|
|
buf = append(buf, b)
|
|
|
|
// http://www.itu.int/rec/dologin_pub.asp?lang=e&id=T-REC-H.264-200305-S!!PDF-E&type=items
|
|
// Table 7-1 NAL unit type codes
|
|
const (
|
|
nonIdrPic = 1
|
|
idrPic = 5
|
|
suppEnhInfo = 6
|
|
paramSet = 8
|
|
)
|
|
switch nalTyp := b & 0x1f; nalTyp {
|
|
case nonIdrPic, idrPic, paramSet, suppEnhInfo:
|
|
writeOut = true
|
|
}
|
|
}
|
|
}
|
|
if len(buf) == len(h264Prefix) {
|
|
return nil
|
|
}
|
|
<-tick
|
|
_, 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 nalType >= typeSingleNALULowBound && nalType <= typeSingleNALUHighBound {
|
|
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) {
|
|
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) {
|
|
// 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 {
|
|
if end {
|
|
panic("bad fragmentation packet")
|
|
}
|
|
newHead := (d[0] & 0xe0) | (d[0] & 0x1f)
|
|
d = d[1:]
|
|
d[0] = newHead
|
|
l.frag = true
|
|
l.writeWithPrefix(d)
|
|
} else {
|
|
if end {
|
|
l.frag = false
|
|
}
|
|
d = d[2:]
|
|
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)
|
|
}
|