av/stream/rtp/encoder.go

/*
NAME
  encoder.go

DESCRIPTION
  See Readme.md

AUTHOR
  Saxon Nelson-Milton (saxon@ausocean.org)

LICENSE
  encoder.go is Copyright (C) 2018 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
  along with revid in gpl.txt. If not, see http://www.gnu.org/licenses.
*/

package rtp

import (
	"io"
	"math/rand"
	"time"
)

const (
	yes            = 1
	no             = 0
	defaultPktType = 1
	timestampFreq  = 90000 // Hz
)

type Encoder struct {
	dst           io.Writer
	ssrc          uint32
	seqNo         uint16
	clock         time.Duration
	frameInterval time.Duration
}

// NewEncoder returns a new Encoder type given an io.Writer - the destination
// after encoding and the desired fps
func NewEncoder(dst io.Writer, fps float64) *Encoder {
	return &Encoder{
		dst:           dst,
		ssrc:          rand.Uint32(),
		frameInterval: time.Duration(float64(time.Second) / fps),
	}
}

// Encode takes a nalu unit and encodes it into an rtp packet and
// writes to the io.Writer given in NewEncoder
func (e *Encoder) Encode(nalu []byte) error {
	pkt := Pkt{
		V:       rtpVer,           // version
		P:       no,               // padding
		X:       no,               // header extension
		CC:      no,               // CSRC count
		M:       no,               // NOTE: need to check if this works (decoders should ignore this)
		PT:      defaultPktType,   // NOTE: 1-23 according to rtp-h264 specs (don't think we need this)
		SN:      e.nxtSeqNo(),     // sequence number
		TS:      e.nxtTimestamp(), // timestamp
		SSRC:    e.ssrc,           // source identifier
		Payload: nalu,
		Padding: no,
	}

	_, err := e.dst.Write(pkt.Bytes())
	if err != nil {
		return err
	}

	e.tick()

	return nil
}

// tick advances the clock one frame interval.
func (e *Encoder) tick() {
	e.clock += e.frameInterval
}

// nxtTimestamp gets the next timestamp
func (e *Encoder) nxtTimestamp() uint32 {
	return uint32(e.clock.Seconds() * timestampFreq)
}

// nxtSeqNo gets the next rtp packet sequence number
func (e *Encoder) nxtSeqNo() uint16 {
	e.seqNo += 1
	return e.seqNo - 1
}