/* NAME lex.go DESCRIPTION lex.go provides a lexer for taking RTP HEVC (H265) and lexing into access units. AUTHORS Saxon A. Nelson-Milton <saxon@ausocean.org> LICENSE Copyright (C) 2019 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 in gpl.txt. If not, see http://www.gnu.org/licenses. */ // Package h265 provides an RTP h265 lexer that can extract h265 access units // from an RTP stream. package h265 import ( "bytes" "encoding/binary" "fmt" "io" "time" "bitbucket.org/ausocean/av/protocol/rtp" ) // NALU types. const ( typeAggregation = 48 typeFragmentation = 49 typePACI = 50 ) // Buffer sizes. const ( maxAUSize = 100000 maxRTPSize = 4096 ) // Lexer is an H265 lexer. type Lexer struct { donl bool // Indicates whether DONL and DOND will be used for the RTP stream. buf *bytes.Buffer // Holds the current access unit. frag bool // Indicates if we're currently dealing with a fragmentation packet. } // NewLexer returns a new Lexer. func NewLexer(donl bool) *Lexer { return &Lexer{ donl: donl, buf: bytes.NewBuffer(make([]byte, 0, maxAUSize)), } } // Lex continually reads RTP packets from the io.Reader src and lexes into // access units which are written to the io.Writer dst. Lex expects that for // each read from src, a single RTP packet is received. func (l *Lexer) 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: %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] >> 1) & 0x3f // If not currently fragmented then we ignore current write. if l.frag && nalType != typeFragmentation { l.buf.Reset() l.frag = false continue } switch nalType { case typeAggregation: l.handleAggregation(payload) case typeFragmentation: l.handleFragmentation(payload) case typePACI: l.handlePACI(payload) default: l.writeWithPrefix(payload) } markerIsSet, err := rtp.Marker(buf[:n]) if err != nil { return fmt.Errorf("could not get marker bit, failed with err: %w\n", err) } if markerIsSet { _, err := l.buf.WriteTo(dst) if err != nil { // TODO: work out what to do here. } l.buf.Reset() } } return nil } // handleAggregation parses NAL units from an aggregation packet and writes // them to the Lexers buffer buf. func (l *Lexer) handleAggregation(d []byte) { idx := 2 for idx < len(d) { if l.donl { switch idx { case 2: idx += 2 default: idx++ } } size := int(binary.BigEndian.Uint16(d[idx:])) idx += 2 nalu := d[idx : idx+size] idx += size l.writeWithPrefix(nalu) } } // handleFragmentation parses NAL units from fragmentation packets and writes // them to the Lexer's buf. func (l *Lexer) handleFragmentation(d []byte) { // Get start and end indiciators from FU header. start := d[2]&0x80 != 0 end := d[2]&0x40 != 0 b1 := (d[0] & 0x81) | ((d[2] & 0x3f) << 1) b2 := d[1] if start { d = d[1:] if l.donl { d = d[2:] } d[0] = b1 d[1] = b2 } else { d = d[3:] if l.donl { d = d[2:] } } switch { case start && !end: l.frag = true l.writeWithPrefix(d) case !start && end: l.frag = false fallthrough case !start && !end: l.writeNoPrefix(d) default: panic("bad fragmentation packet") } } // handlePACI will handl PACI packets // // TODO: complete this func (l *Lexer) handlePACI(d []byte) { panic("unsupported nal type") } // write writes a NAL unit to the Lexer's buf in byte stream format using the // start code. func (l *Lexer) 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 *Lexer) writeNoPrefix(d []byte) { l.buf.Write(d) }