av/cmd/vidforward/main.go

/*
DESCRIPTION
  vidforward is a service for receiving video from cameras and then forwarding to
  youtube. By acting as the RTMP encoder (instead of the camera) vidforward can enable
  persistent streams by sending slate images during camera downtime.

AUTHORS
  Saxon A. Nelson-Milton <saxon@ausocean.org>

LICENSE
  Copyright (C) 2022 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 main

import (
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io"
	"net"
	"net/http"
	"strconv"
	"sync"
	"time"

	"bitbucket.org/ausocean/av/cmd/vidforward/global"
	"bitbucket.org/ausocean/av/codec/codecutil"
	"bitbucket.org/ausocean/av/container/mts"
	"bitbucket.org/ausocean/av/revid"
	"bitbucket.org/ausocean/av/revid/config"
	"bitbucket.org/ausocean/utils/logging"
	"gopkg.in/natefinch/lumberjack.v2"
)

// Server defaults.
const (
	defaultPort = "8080"
	defaultHost = ""
)

// Logging configuration.
const (
	logPath      = "/var/log/vidforward/vidforward.log"
	logMaxSize   = 500 // MB
	logMaxBackup = 10
	logMaxAge    = 28 // days
	logVerbosity = logging.Info
	logSuppress  = true
)

// recvErrorDelay is a delay used when there's recv issues. It is intended to
// prevent spamming from a single client.
const recvErrorDelay = 7 * time.Second

type MAC string

// The possible states for a broadcast.
const (
	statusActive = "active"
	statusSlate  = "slate"
	statusCreate = "create"
	statusPlay   = "play"
)

// Broadcast is representative of a broadcast to be forwarded.
type Broadcast struct {
	mac    MAC          // MAC address of the device from which the video is being received.
	url    string       // The destination youtube RTMP URL.
	status string       // The broadcast status i.e. active or slate.
	rv     *revid.Revid // The revid pipeline which will handle forwarding to youtube.
}

// broadcastManager manages a map of Broadcasts we expect to be forwarding video
// for. The broadcastManager is communicated with through a series of HTTP request
// handlers. There is a basic REST API through which we can add/delete broadcasts,
// and a recv handler which is invoked when a camera wishes to get its video
// forwarded to youtube.
type broadcastManager struct {
	broadcasts       map[MAC]Broadcast
	slateExitSignals map[MAC]chan struct{} // Used to signal to stop writing slate image.
	log              logging.Logger
	dogNotifier      *watchdogNotifier
	mu               sync.Mutex
}

// newBroadcastManager returns a new broadcastManager with the provided logger.
func newBroadcastManager(l logging.Logger) (*broadcastManager, error) {
	m := &broadcastManager{
		log:              l,
		broadcasts:       make(map[MAC]Broadcast),
		slateExitSignals: make(map[MAC]chan struct{}),
	}
	notifier, err := newWatchdogNotifier(l, terminationCallback(m))
	if err != nil {
		return nil, err
	}
	m.dogNotifier = notifier
	return m, nil
}

// terminationCallback provides a callback that saves the provided
// broadcastManagers state.
func terminationCallback(m *broadcastManager) func() {
	return func() {
		err := m.save()
		if err != nil {
			m.log.Error("could not save on notifier termination signal", "error", err)
			return
		}
		m.log.Info("successfully saved broadcast manager state on termination signal")
	}
}

// recvHandler handles recv requests for video forwarding. The MAC is firstly
// checked to ensure it is "active" i.e. should be sending data, and then the
// video is extracted from the request body and provided to the revid pipeline
// corresponding to said MAC.
// Clips of MPEG-TS h264 are the only accepted format and codec.
func (m *broadcastManager) recv(w http.ResponseWriter, r *http.Request) {
	done := m.dogNotifier.handlerInvoked("recv")
	defer done()

	m.log.Debug("recv handler")
	q := r.URL.Query()
	ma := MAC(q.Get("ma"))

	if !m.isActive(ma) {
		m.errorLogWrite(m.log, w, "forward request mac is not mapped, doing nothing", "mac", ma)
		time.Sleep(recvErrorDelay)
		return
	}

	// We can't receive video if we're in slate mode.
	if m.getStatus(ma) == statusSlate {
		m.errorLogWrite(m.log, w, "cannot receive video for this mac, status is slate", "mac", ma)
		time.Sleep(recvErrorDelay)
		return
	}

	const videoPin = "V0"
	sizeStr := q.Get(videoPin)
	size, err := strconv.Atoi(sizeStr)
	if err != nil || size <= 0 {
		m.errorLogWrite(m.log, w, "invalid video size", "error", err, "size str", sizeStr)
		return
	}

	// Prepare HTTP response with received video size and device mac.
	resp := map[string]interface{}{"ma": ma, "V0": size}

	mtsClip, err := io.ReadAll(r.Body)
	if err != nil {
		m.errorLogWrite(m.log, w, "could not read forward request body", "error", err)
		return
	}
	defer r.Body.Close()

	if len(mtsClip)%mts.PacketSize != 0 {
		m.errorLogWrite(m.log, w, "invalid clip length", "length", len(mtsClip))
		return
	}

	// Extract the pure h264 from the MPEG-TS clip.
	h264Clip, err := mts.Extract(mtsClip)
	if err != nil {
		m.errorLogWrite(m.log, w, "could not extract m.264 from the MPEG-TS clip", "error", err)
		return
	}

	rv := m.getPipeline(ma)
	if r == nil {
		panic(fmt.Sprintf("no revid pipeline for mac address: %s", ma))
	}

	for i, frame := range h264Clip.Frames() {
		_, err := rv.Write(frame.Media)
		if err != nil {
			m.errorLogWrite(m.log, w, "could not write frame", "no.", i, "error", err)
			return
		}
	}

	// Return response to client as JSON.
	jsn, err := json.Marshal(resp)
	if err != nil {
		m.errorLogWrite(m.log, w, "could not get json for response", "error", err)
		return
	}
	fmt.Fprint(w, string(jsn))
}

// control handles control API requests.
func (m *broadcastManager) control(w http.ResponseWriter, r *http.Request) {
	done := m.dogNotifier.handlerInvoked("control")
	defer done()

	m.log.Info("control request", "method", r.Method)
	switch r.Method {
	case http.MethodPut:
		m.processRequest(w, r, m.createOrUpdate)
	case http.MethodDelete:
		m.processRequest(w, r, m.delete)
	default:
		m.errorLogWrite(m.log, w, "unhandled http method", "method", r.Method)
	}
}

// processRequest unmarshals the broadcast data object from the request into
// a Broadcast value, and then performs the provided action with that value.
func (m *broadcastManager) processRequest(w http.ResponseWriter, r *http.Request, action func(Broadcast) error) {
	body, err := io.ReadAll(r.Body)
	if err != nil {
		m.errorLogWrite(m.log, w, "could not read request body", "body", r.Body)
		return
	}
	defer r.Body.Close()

	var broadcast Broadcast
	err = json.Unmarshal(body, &broadcast)
	if err != nil {
		m.errorLogWrite(m.log, w, "could not marshal data", "error", err)
		return
	}

	err = action(broadcast)
	if err != nil {
		m.errorLogWrite(m.log, w, "could not perform action", "method", r.Method, "error", err)
	}
}

// getPipeline gets the revid pipeline corresponding to a provided device MAC.
func (m *broadcastManager) getPipeline(ma MAC) *revid.Revid {
	m.mu.Lock()
	defer m.mu.Unlock()
	v, ok := m.broadcasts[ma]
	if !ok {
		return nil
	}
	return v.rv
}

// getStatus gets the broadcast's status corresponding to the provided MAC.
func (m *broadcastManager) getStatus(ma MAC) string {
	m.mu.Lock()
	defer m.mu.Unlock()
	v, ok := m.broadcasts[ma]
	if !ok {
		return ""
	}
	return v.status
}

// isActive returns true if a MAC is registered to the broadcast manager.
func (m *broadcastManager) isActive(ma MAC) bool {
	m.mu.Lock()
	defer m.mu.Unlock()
	_, ok := m.broadcasts[ma]
	return ok
}

// createOrUpdate creates or updates a Broadcast record. The revid pipeline
// corresponding to the broadcast MAC is firsty configured/re-configured, and
// the pipeline is "started", which will ready it for receiving video on its
// input. In the case that the status is "slate", we will spin up a routine to
// handle writing a slate image to the pipeline.
func (m *broadcastManager) createOrUpdate(broadcast Broadcast) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	cfg := config.Config{
		Logger:     m.log,
		Input:      config.InputManual,
		InputCodec: codecutil.H264_AU,
		Outputs:    []uint8{config.OutputRTMP},
		RTMPURL:    broadcast.url,
		LogLevel:   logging.Debug,
	}

	var err error
	broadcast.rv, err = revid.New(cfg, nil)
	if err != nil {
		return fmt.Errorf("could not initialise revid: %w", err)
	}

	m.broadcasts[broadcast.mac] = broadcast
	err = broadcast.rv.Start()
	if err != nil {
		return fmt.Errorf("could not start revid pipeline: %w", err)
	}

	switch broadcast.status {
	case statusCreate:
		fallthrough
	case statusActive, statusPlay:
		signal, ok := m.slateExitSignals[broadcast.mac]
		if ok {
			close(signal)
			delete(m.slateExitSignals, broadcast.mac)
		}
	case statusSlate:
		m.log.Debug("slate request")
		// If there's a signal channel it means that we're already writing the slate
		// image and theres nothing to do, so return.
		_, ok := m.slateExitSignals[broadcast.mac]
		if ok {
			m.log.Warning("already writing slate")
			return nil
		}

		// First create a signal that can be used to stop the slate writing routine.
		// This will be provided to the writeSlate routine below.
		signalCh := make(chan struct{})
		m.slateExitSignals[broadcast.mac] = signalCh

		err = writeSlateAndCheckErrors(broadcast.rv, signalCh, m.log)
		if err != nil {
			return err
		}
	default:
		return fmt.Errorf("unknown status string: %s", broadcast.status)
	}
	return nil
}

// delete removes a broadcast from the record.
func (m *broadcastManager) delete(broadcast Broadcast) error {
	m.mu.Lock()
	b, ok := m.broadcasts[broadcast.mac]
	if !ok {
		return errors.New("no broadcast by that mac in record")
	}
	b.rv.Stop()
	delete(m.broadcasts, broadcast.mac)
	m.mu.Unlock()
	return nil
}

func main() {
	host := flag.String("host", defaultHost, "Host IP to run video forwarder on.")
	port := flag.String("port", defaultPort, "Port to run video forwarder on.")
	flag.Parse()

	if *host == "" || net.ParseIP(*host) == nil {
		panic(fmt.Sprintf("invalid host, host: %s", *host))
	}

	// Create lumberjack logger to handle logging to file.
	fileLog := &lumberjack.Logger{
		Filename:   logPath,
		MaxSize:    logMaxSize,
		MaxBackups: logMaxBackup,
		MaxAge:     logMaxAge,
	}

	// Create logger that we call methods on to log, which in turn writes to the
	// lumberjack and netloggers.
	log := logging.New(logVerbosity, io.MultiWriter(fileLog), logSuppress)

	global.SetLogger(log)

	m, err := newBroadcastManager(log)
	if err != nil {
		log.Fatal("could not create new broadcast manager", "error", err)
	}
	http.HandleFunc("/recv", m.recv)
	http.HandleFunc("/control", m.control)

	go m.dogNotifier.notify()

	http.ListenAndServe(*host+":"+*port, nil)
}