Add AV1 OBU header parser and OBU types utilities

Implement AV1 bitstream format OBU header parser and utility functions
for OBU types.

Author: JoeTurki

codecs/av1/obu/errors.go

@@ -0,0 +1,14 @@
+// SPDX-FileCopyrightText: 2023 The Pion community <https://pion.ly>
+// SPDX-License-Identifier: MIT
+
+package obu
+
+import "errors"
+
+var (
+	// ErrInvalidOBUHeader is returned when an OBU header has forbidden bits set.
+	ErrInvalidOBUHeader = errors.New("invalid OBU header")
+	// ErrShortHeader is returned when an OBU header is not large enough.
+	// This can happen when an extension header is expected but not present.
+	ErrShortHeader = errors.New("OBU header is not large enough")
+)

codecs/av1/obu/obu.go

@@ -0,0 +1,219 @@
+// SPDX-FileCopyrightText: 2023 The Pion community <https://pion.ly>
+// SPDX-License-Identifier: MIT
+
+package obu
+
+import (
+	"fmt"
+)
+
+// Type represents the type of an AV1 OBU.
+type Type uint8
+
+// OBU types as defined in the AV1 specification.
+// 5.3.1: https://aomediacodec.github.io/av1-spec/av1-spec.pdf#page=39
+const (
+	// OBUSequenceHeader av1 sequence_header_obu.
+	OBUSequenceHeader = Type(1)
+	// OBUTemporalDelimiter av1 temporal_delimiter_obu.
+	OBUTemporalDelimiter = Type(2)
+	// OBUFrameHeader av1 frame_header_obu.
+	OBUFrameHeader = Type(3)
+	// OBUTileGroup av1 tile_group_obu.
+	OBUTileGroup = Type(4)
+	// OBUMetadata av1 metadata_obu.
+	OBUMetadata = Type(5)
+	// OBUFrame av1 frame_obu.
+	OBUFrame = Type(6)
+	// OBURedundantFrameHeader av1 redundant_frame_header_obu.
+	OBURedundantFrameHeader = Type(7)
+	// OBUTileList av1 tile_list_obu.
+	OBUTileList = Type(8)
+	// OBUPadding av1 padding_obu.
+	OBUPadding = Type(15)
+)
+
+//nolint:cyclop
+func (o Type) String() string {
+	switch o {
+	case OBUSequenceHeader:
+		return "OBU_SEQUENCE_HEADER"
+	case OBUTemporalDelimiter:
+		return "OBU_TEMPORAL_DELIMITER"
+	case OBUFrameHeader:
+		return "OBU_FRAME_HEADER"
+	case OBUTileGroup:
+		return "OBU_TILE_GROUP"
+	case OBUMetadata:
+		return "OBU_METADATA"
+	case OBUFrame:
+		return "OBU_FRAME"
+	case OBURedundantFrameHeader:
+		return "OBU_REDUNDANT_FRAME_HEADER"
+	case OBUTileList:
+		return "OBU_TILE_LIST"
+	case OBUPadding:
+		return "OBU_PADDING"
+	default:
+		return "OBU_RESERVED"
+	}
+}
+
+// Header represents the header of an OBU obu_header().
+// 5.3.2: https://aomediacodec.github.io/av1-spec/av1-spec.pdf#page=40
+type Header struct {
+	Type            Type
+	ExtensionHeader *ExtensionHeader
+	HasSizeField    bool
+	Reserved1Bit    bool
+}
+
+// ParseOBUHeader parses an OBU header from the given data.
+// 5.3.2: https://aomediacodec.github.io/av1-spec/av1-spec.pdf#page=40
+/*
+	obu_header() { Type
+		obu_forbidden_bit     f(1)
+		obu_type              f(4)
+		obu_extension_flag    f(1)
+		obu_has_size_field    f(1)
+		obu_reserved_1bit     f(1)
+		if ( obu_extension_flag == 1 )
+			obu_extension_header()
+		}
+	}
+*/
+func ParseOBUHeader(data []byte) (*Header, error) {
+	if len(data) < 1 {
+		return nil, fmt.Errorf("%w: data is too short", ErrShortHeader)
+	}
+
+	forbiddenBit := data[0] & 0x80
+	if forbiddenBit != 0 {
+		return nil, fmt.Errorf("%w: forbidden bit is set", ErrInvalidOBUHeader)
+	}
+
+	obuType := Type((data[0] & 0x78) >> 3)
+	obuExtensionFlag := (data[0] & 0x04) != 0
+	obuHasSizeField := (data[0] & 0x02) != 0
+	obuReserved1Bit := (data[0] & 0x01) != 0
+
+	header := &Header{
+		Type:         obuType,
+		HasSizeField: obuHasSizeField,
+		Reserved1Bit: obuReserved1Bit,
+	}
+
+	if obuExtensionFlag {
+		if len(data) < 2 {
+			return nil, fmt.Errorf("%w: Unexpected end of data, expected extension header", ErrShortHeader)
+		}
+
+		extensionHeader := ParseOBUExtensionHeader(data[1])
+		header.ExtensionHeader = &extensionHeader
+	}
+
+	return header, nil
+}
+
+// Marshal serializes the OBU header to a byte slice.
+// If the OBU has an extension header, the extension header is serialized as well.
+// 5.3.2: https://aomediacodec.github.io/av1-spec/av1-spec.pdf#page=40
+/*
+	obu_header() { Type
+		obu_forbidden_bit     f(1)
+		obu_type              f(4)
+		obu_extension_flag    f(1)
+		obu_has_size_field    f(1)
+		obu_reserved_1bit     f(1)
+		if ( obu_extension_flag == 1 )
+			obu_extension_header()
+		}
+	}
+*/
+func (o *Header) Marshal() []byte {
+	header := make([]byte, o.Size())
+
+	header[0] = (byte(o.Type) & 0x0f) << 3
+
+	if o.ExtensionHeader != nil {
+		header[0] |= 0x04
+		header[1] = o.ExtensionHeader.Marshal()
+	}
+
+	if o.HasSizeField {
+		header[0] |= 0x02
+	}
+
+	if o.Reserved1Bit {
+		header[0] |= 0x01
+	}
+
+	return header
+}
+
+// Size returns the size of the OBU header in bytes.
+func (o *Header) Size() int {
+	size := 1
+	if o.ExtensionHeader != nil {
+		size++
+	}
+
+	return size
+}
+
+// ExtensionHeader represents an OBU extension header obu_extension_header().
+// 5.3.3 https://aomediacodec.github.io/av1-spec/av1-spec.pdf#page=40
+type ExtensionHeader struct {
+	TemporalID    uint8
+	SpatialID     uint8
+	Reserved3Bits uint8
+}
+
+// ParseOBUExtensionHeader parses an OBU extension header from the given data.
+// 5.3.3 https://aomediacodec.github.io/av1-spec/av1-spec.pdf#page=40
+/*
+	obu_extension_header() { Type
+		temporal_id                      f(3)
+		spatial_id                       f(2)
+		extension_header_reserved_3bits  f(3)
+	}
+*/
+func ParseOBUExtensionHeader(headerData byte) ExtensionHeader {
+	return ExtensionHeader{
+		TemporalID:    headerData >> 5,
+		SpatialID:     (headerData >> 3) & 0x03,
+		Reserved3Bits: headerData & 0x07,
+	}
+}
+
+// Marshal serializes the OBU extension header to a byte slice.
+// 5.3.3 https://aomediacodec.github.io/av1-spec/av1-spec.pdf#page=40
+/*
+	obu_extension_header() { Type
+		temporal_id                      f(3)
+		spatial_id                       f(2)
+		extension_header_reserved_3bits  f(3)
+	}
+*/
+func (o *ExtensionHeader) Marshal() byte {
+	return (o.TemporalID << 5) | ((o.SpatialID & 0x3) << 3) | (o.Reserved3Bits & 0x07)
+}
+
+// OBU represents an AV1 OBU.
+// 5.1 https://aomediacodec.github.io/av1-spec/av1-spec.pdf#page=39
+type OBU struct {
+	Header  Header
+	Payload []byte
+}
+
+// Marshal serializes the OBU to low-overhead bitstream format.
+// 5.2 https://aomediacodec.github.io/av1-spec/av1-spec.pdf#page=40
+func (o *OBU) Marshal() []byte {
+	buffer := o.Header.Marshal()
+
+	if o.Header.HasSizeField {
+		buffer = append(buffer, WriteToLeb128(uint(len(o.Payload)))...)
+	}
+
+	return append(buffer, o.Payload...)
+}