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"

	"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/iot/pi/netlogger"
	"bitbucket.org/ausocean/utils/logging"
	"gopkg.in/natefinch/lumberjack.v2"
)

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

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

type MAC string

// Broadcast is representative of a broadcast to be forwarded.
type Broadcast struct {
	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, inactive and 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
	log        logging.Logger
	mu         sync.Mutex
}

// 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) {
	m.log.Debug("recv handler")
	q := r.URL.Query()
	ma := MAC(q.Get("ma"))

	_, ok := m.broadcasts[ma]
	if !ok {
		m.errorLogWrite(w, "forward request mac is not mapped, doing nothing", "mac", ma)
		return
	}

	const videoPin = "V0"
	sizeStr := q.Get(videoPin)
	size, err := strconv.Atoi(sizeStr)
	if err != nil || size <= 0 {
		m.errorLogWrite(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(w, "could not read forward request body", "error", err)
		return
	}
	defer r.Body.Close()

	if len(mtsClip)%mts.PacketSize != 0 {
		m.errorLogWrite(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(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(w, "could not write frame", "no.", i)
			return
		}
	}

	// Return response to client as JSON.
	jsn, err := json.Marshal(resp)
	if err != nil {
		m.errorLogWrite(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) {
	switch r.Method {
	case http.MethodPut:
		m.processRequest(w, r, m.createOrUpdate)
	case http.MethodDelete:
		m.processRequest(w, r, m.delete)
	default:
		m.errorLogWrite(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(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(w, "could not marshal data", "error", err)
		return
	}

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

// writeError logs an error and writes to w in JSON format.
func (m *broadcastManager) errorLogWrite(w http.ResponseWriter, msg string, args ...interface{}) {
	m.log.Error(msg, args...)
	w.Header().Add("Content-Type", "application/json")
	fmt.Fprint(w, `{"er":"`+msg+`"}`)
}

// 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
}

// 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)
	}

	if broadcast.Status == "slate" {
		go m.writeSlate(broadcast.MAC, broadcast.RV)
	}
	return nil
}

// This is just a stub. Eventually this will handle writing of a slate image
// to the destination RTMP URL.
func (m *broadcastManager) writeSlate(ma MAC, rv *revid.Revid) {}

// 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 netlogger to handle logging to cloud.
	netLog := netlogger.New()

	// 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, netLog), logSuppress)

	bm := &broadcastManager{log: log, broadcasts: map[MAC]Broadcast{}}
	http.HandleFunc("/recv", bm.recv)
	http.HandleFunc("/control", bm.control)
	http.ListenAndServe(*host+":"+*port, nil)
}

func isMac(m string) bool {
	if len(m) != 17 || m == "00:00:00:00:00:00" {
		return false
	}

	for i := 0; i <= 15; i++ {
		if (i+1)%3 == 0 && m[i] != ':' {
			return false
		}

		if (3-i)%3 != 0 {
			continue
		}

		_, err := strconv.ParseUint(m[i:i+2], 16, 64)
		if err != nil {
			return false
		}
	}
	return true
}