/* NAME extract.go DESCRIPTION extract.go provides an Extractor 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/codec/h264/h264dec" "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. ) // Bytes for an access unit delimeter. var aud = []byte{0x00, 0x00, 0x01, 0x09, 0xf0} // Extractor is an Extractor for extracting H264 access units from RTP stream. type Extractor struct { buf *bytes.Buffer // Holds the current access unit. frag bool // Indicates if we're currently dealing with a fragmentation packet. dst io.Writer // The destination we'll be writing extracted NALUs to. toWrite []byte // Holds the current NAL unit with start code to be written. } // NewExtractor returns a new Extractor. func NewExtractor() *Extractor { return &Extractor{ 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 *Extractor) Extract(dst io.Writer, src io.Reader, delay time.Duration) error { e.toWrite = []byte{0, 0, 0, 1} e.buf.Write(aud) e.dst = dst 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: %w\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: %w\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") } } } } // handleSTAPA parses NAL units from an aggregation packet and writes // them to the Extractor's buffer buf. func (e *Extractor) 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 Extractor's buf. func (e *Extractor) 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) } } // writeWithPrefix writes a NAL unit to the Extractor's buf in byte stream format // using the start code, and sends any ready prior access unit stored in the buf // to the destination. func (e *Extractor) writeWithPrefix(d []byte) { e.toWrite = append(e.toWrite, d...) curType, _ := NALType(e.toWrite) if e.buf.Len() != 0 && (curType == h264dec.NALTypeSPS || curType == h264dec.NALTypeIDR || curType == h264dec.NALTypeNonIDR) { e.buf.WriteTo(e.dst) e.buf.Reset() e.buf.Write(aud) } e.buf.Write(e.toWrite) e.toWrite = e.toWrite[:4] } // writeNoPrefix writes data to the Extractor's buf. This is used for non start // fragmentations of a NALU. func (e *Extractor) writeNoPrefix(d []byte) { e.buf.Write(d) }