github.com/pion/webrtc/v4@v4.0.1/pkg/media/h264writer/h264writer.go (about)

     1  // SPDX-FileCopyrightText: 2023 The Pion community <https://pion.ly>
     2  // SPDX-License-Identifier: MIT
     3  
     4  // Package h264writer implements H264 media container writer
     5  package h264writer
     6  
     7  import (
     8  	"bytes"
     9  	"encoding/binary"
    10  	"io"
    11  	"os"
    12  
    13  	"github.com/pion/rtp"
    14  	"github.com/pion/rtp/codecs"
    15  )
    16  
    17  type (
    18  	// H264Writer is used to take RTP packets, parse them and
    19  	// write the data to an io.Writer.
    20  	// Currently it only supports non-interleaved mode
    21  	// Therefore, only 1-23, 24 (STAP-A), 28 (FU-A) NAL types are allowed.
    22  	// https://tools.ietf.org/html/rfc6184#section-5.2
    23  	H264Writer struct {
    24  		writer       io.Writer
    25  		hasKeyFrame  bool
    26  		cachedPacket *codecs.H264Packet
    27  	}
    28  )
    29  
    30  // New builds a new H264 writer
    31  func New(filename string) (*H264Writer, error) {
    32  	f, err := os.Create(filename) //nolint:gosec
    33  	if err != nil {
    34  		return nil, err
    35  	}
    36  
    37  	return NewWith(f), nil
    38  }
    39  
    40  // NewWith initializes a new H264 writer with an io.Writer output
    41  func NewWith(w io.Writer) *H264Writer {
    42  	return &H264Writer{
    43  		writer: w,
    44  	}
    45  }
    46  
    47  // WriteRTP adds a new packet and writes the appropriate headers for it
    48  func (h *H264Writer) WriteRTP(packet *rtp.Packet) error {
    49  	if len(packet.Payload) == 0 {
    50  		return nil
    51  	}
    52  
    53  	if !h.hasKeyFrame {
    54  		if h.hasKeyFrame = isKeyFrame(packet.Payload); !h.hasKeyFrame {
    55  			// key frame not defined yet. discarding packet
    56  			return nil
    57  		}
    58  	}
    59  
    60  	if h.cachedPacket == nil {
    61  		h.cachedPacket = &codecs.H264Packet{}
    62  	}
    63  
    64  	data, err := h.cachedPacket.Unmarshal(packet.Payload)
    65  	if err != nil {
    66  		return err
    67  	}
    68  
    69  	_, err = h.writer.Write(data)
    70  
    71  	return err
    72  }
    73  
    74  // Close closes the underlying writer
    75  func (h *H264Writer) Close() error {
    76  	h.cachedPacket = nil
    77  	if h.writer != nil {
    78  		if closer, ok := h.writer.(io.Closer); ok {
    79  			return closer.Close()
    80  		}
    81  	}
    82  
    83  	return nil
    84  }
    85  
    86  func isKeyFrame(data []byte) bool {
    87  	const (
    88  		typeSTAPA       = 24
    89  		typeSPS         = 7
    90  		naluTypeBitmask = 0x1F
    91  	)
    92  
    93  	var word uint32
    94  
    95  	payload := bytes.NewReader(data)
    96  	if err := binary.Read(payload, binary.BigEndian, &word); err != nil {
    97  		return false
    98  	}
    99  
   100  	naluType := (word >> 24) & naluTypeBitmask
   101  	if naluType == typeSTAPA && word&naluTypeBitmask == typeSPS {
   102  		return true
   103  	} else if naluType == typeSPS {
   104  		return true
   105  	}
   106  
   107  	return false
   108  }