github.com/pion/webrtc/v3@v3.2.24/pkg/media/ivfreader/ivfreader.go (about)

     1  // SPDX-FileCopyrightText: 2023 The Pion community <https://pion.ly>
     2  // SPDX-License-Identifier: MIT
     3  
     4  // Package ivfreader implements IVF media container reader
     5  package ivfreader
     6  
     7  import (
     8  	"encoding/binary"
     9  	"errors"
    10  	"fmt"
    11  	"io"
    12  )
    13  
    14  const (
    15  	ivfFileHeaderSignature = "DKIF"
    16  	ivfFileHeaderSize      = 32
    17  	ivfFrameHeaderSize     = 12
    18  )
    19  
    20  var (
    21  	errNilStream             = errors.New("stream is nil")
    22  	errIncompleteFrameHeader = errors.New("incomplete frame header")
    23  	errIncompleteFrameData   = errors.New("incomplete frame data")
    24  	errIncompleteFileHeader  = errors.New("incomplete file header")
    25  	errSignatureMismatch     = errors.New("IVF signature mismatch")
    26  	errUnknownIVFVersion     = errors.New("IVF version unknown, parser may not parse correctly")
    27  )
    28  
    29  // IVFFileHeader 32-byte header for IVF files
    30  // https://wiki.multimedia.cx/index.php/IVF
    31  type IVFFileHeader struct {
    32  	signature           string // 0-3
    33  	version             uint16 // 4-5
    34  	headerSize          uint16 // 6-7
    35  	FourCC              string // 8-11
    36  	Width               uint16 // 12-13
    37  	Height              uint16 // 14-15
    38  	TimebaseDenominator uint32 // 16-19
    39  	TimebaseNumerator   uint32 // 20-23
    40  	NumFrames           uint32 // 24-27
    41  	unused              uint32 // 28-31
    42  }
    43  
    44  // IVFFrameHeader 12-byte header for IVF frames
    45  // https://wiki.multimedia.cx/index.php/IVF
    46  type IVFFrameHeader struct {
    47  	FrameSize uint32 // 0-3
    48  	Timestamp uint64 // 4-11
    49  }
    50  
    51  // IVFReader is used to read IVF files and return frame payloads
    52  type IVFReader struct {
    53  	stream               io.Reader
    54  	bytesReadSuccesfully int64
    55  }
    56  
    57  // NewWith returns a new IVF reader and IVF file header
    58  // with an io.Reader input
    59  func NewWith(in io.Reader) (*IVFReader, *IVFFileHeader, error) {
    60  	if in == nil {
    61  		return nil, nil, errNilStream
    62  	}
    63  
    64  	reader := &IVFReader{
    65  		stream: in,
    66  	}
    67  
    68  	header, err := reader.parseFileHeader()
    69  	if err != nil {
    70  		return nil, nil, err
    71  	}
    72  
    73  	return reader, header, nil
    74  }
    75  
    76  // ResetReader resets the internal stream of IVFReader. This is useful
    77  // for live streams, where the end of the file might be read without the
    78  // data being finished.
    79  func (i *IVFReader) ResetReader(reset func(bytesRead int64) io.Reader) {
    80  	i.stream = reset(i.bytesReadSuccesfully)
    81  }
    82  
    83  // ParseNextFrame reads from stream and returns IVF frame payload, header,
    84  // and an error if there is incomplete frame data.
    85  // Returns all nil values when no more frames are available.
    86  func (i *IVFReader) ParseNextFrame() ([]byte, *IVFFrameHeader, error) {
    87  	buffer := make([]byte, ivfFrameHeaderSize)
    88  	var header *IVFFrameHeader
    89  
    90  	bytesRead, err := io.ReadFull(i.stream, buffer)
    91  	headerBytesRead := bytesRead
    92  	if errors.Is(err, io.ErrUnexpectedEOF) {
    93  		return nil, nil, errIncompleteFrameHeader
    94  	} else if err != nil {
    95  		return nil, nil, err
    96  	}
    97  
    98  	header = &IVFFrameHeader{
    99  		FrameSize: binary.LittleEndian.Uint32(buffer[:4]),
   100  		Timestamp: binary.LittleEndian.Uint64(buffer[4:12]),
   101  	}
   102  
   103  	payload := make([]byte, header.FrameSize)
   104  	bytesRead, err = io.ReadFull(i.stream, payload)
   105  	if errors.Is(err, io.ErrUnexpectedEOF) {
   106  		return nil, nil, errIncompleteFrameData
   107  	} else if err != nil {
   108  		return nil, nil, err
   109  	}
   110  
   111  	i.bytesReadSuccesfully += int64(headerBytesRead) + int64(bytesRead)
   112  	return payload, header, nil
   113  }
   114  
   115  // parseFileHeader reads 32 bytes from stream and returns
   116  // IVF file header. This is always called before ParseNextFrame()
   117  func (i *IVFReader) parseFileHeader() (*IVFFileHeader, error) {
   118  	buffer := make([]byte, ivfFileHeaderSize)
   119  
   120  	bytesRead, err := io.ReadFull(i.stream, buffer)
   121  	if errors.Is(err, io.ErrUnexpectedEOF) {
   122  		return nil, errIncompleteFileHeader
   123  	} else if err != nil {
   124  		return nil, err
   125  	}
   126  
   127  	header := &IVFFileHeader{
   128  		signature:           string(buffer[:4]),
   129  		version:             binary.LittleEndian.Uint16(buffer[4:6]),
   130  		headerSize:          binary.LittleEndian.Uint16(buffer[6:8]),
   131  		FourCC:              string(buffer[8:12]),
   132  		Width:               binary.LittleEndian.Uint16(buffer[12:14]),
   133  		Height:              binary.LittleEndian.Uint16(buffer[14:16]),
   134  		TimebaseDenominator: binary.LittleEndian.Uint32(buffer[16:20]),
   135  		TimebaseNumerator:   binary.LittleEndian.Uint32(buffer[20:24]),
   136  		NumFrames:           binary.LittleEndian.Uint32(buffer[24:28]),
   137  		unused:              binary.LittleEndian.Uint32(buffer[28:32]),
   138  	}
   139  
   140  	if header.signature != ivfFileHeaderSignature {
   141  		return nil, errSignatureMismatch
   142  	} else if header.version != uint16(0) {
   143  		return nil, fmt.Errorf("%w: expected(0) got(%d)", errUnknownIVFVersion, header.version)
   144  	}
   145  
   146  	i.bytesReadSuccesfully += int64(bytesRead)
   147  	return header, nil
   148  }