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 = 33
	timestampFreq  = 90000 // Hz
	mtsSize        = 188
	bufferSize     = 1000
	sendSize       = 7 * 188
)

// Encoder implements io writer and provides functionality to wrap data into
// rtp packets
type Encoder struct {
	dst           io.Writer
	ssrc          uint32
	seqNo         uint16
	clock         time.Duration
	frameInterval time.Duration
	fps           int
	buffer        []byte
	pktSpace      [defPktSize]byte
}

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

// Write provides an interface between a prior encoder and this rtp encoder,
// so that multiple layers of packetization can occur.
func (e *Encoder) Write(data []byte) (int, error) {
	e.buffer = append(e.buffer, data...)
	if len(e.buffer) < sendSize {
		return len(data), nil
	}
	buf := e.buffer
	for len(buf) != 0 {
		l := min(sendSize, len(buf))
		err := e.Encode(buf[:l])
		if err != nil {
			return len(data), err
		}
		buf = buf[l:]
	}
	e.buffer = e.buffer[:0]
	return len(data), nil
}

func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}

// 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(payload []byte) error {
	pkt := Pkt{
		V:       rtpVer,           // version
		X:       no,               // header extension
		CC:      no,               // CSRC count
		M:       no,               // NOTE: need to check if this works (decoders should ignore this)
		PT:      defaultPktType,   // 33 for mpegts
		SN:      e.nxtSeqNo(),     // sequence number
		TS:      e.nxtTimestamp(), // timestamp
		SSRC:    e.ssrc,           // source identifier
		Payload: payload,
		Padding: no,
	}
	_, err := e.dst.Write(pkt.Bytes(e.pktSpace[:defPktSize]))
	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++
	return e.seqNo - 1
}