codec/h264/lex.go: wrote lexer for lexing h264 access units from RTP stream

codec/h264/lex.go

@@ -30,10 +30,14 @@ LICENSE
 package h264
 
 import (
+	"bytes"
+	"encoding/binary"
+	"fmt"
 	"io"
 	"time"
 
 	"bitbucket.org/ausocean/av/codec/codecutil"
+	"bitbucket.org/ausocean/av/protocol/rtp"
 )
 
 var noDelay = make(chan time.Time)
@@ -133,3 +137,166 @@ outer:
 	_, err := dst.Write(buf)
 	return err
 }
+
+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
+)
+
+// Buffer sizes.
+const (
+	maxAUSize  = 100000
+	maxRTPSize = 4096
+)
+
+// RTPLexer is a lexer for lexing H265 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))}
+}
+
+// LexFromRTP 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) LexFromRTP(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
+}
+
+// handleSTAP 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)
+	}
+}
+
+// handleFragmentation 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 {
+		const FUIndicatorIdx = 0
+		newHead := (d[FUIndicatorIdx] & 0xe0) | (d[FUHeadIdx] & 0x1f)
+		d = d[1:]
+		d[0] = newHead
+	} else {
+		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")
+	}
+}
+
+// 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)
+}