oko_public/mofin/main.go

package main

import (
	"context"
	"flag"
	"log"
	"net"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"

	"frame"
)

const (
	workerCount = 4

	// connLimit bounds concurrent inbound connections so a connection flood
	// (auth handshake parked, or clips queued behind busy workers) cannot
	// exhaust goroutines or file descriptors.
	connLimit = 16

	// clipChanCap gives in-flight ingest handlers somewhere to park a decoded
	// clip without holding the connection's socket open for the workers.
	clipChanCap = 32

	shutdownGrace = 5 * time.Second
)

func main() {
	var (
		listenAddr string
		terpAddr   string
		threshold  int
		passphrase string
		authToken  string
		pbkdf2Iter int
	)

	flag.StringVar(&listenAddr, "listen", ":8083", "address to listen for incoming clips")
	flag.StringVar(&terpAddr, "terp", "localhost:8081", "address to forward clips with motion to terp")
	flag.IntVar(&threshold, "threshold", 30, "pixel-value sensitivity (0-255) for motion detection comparison")
	flag.StringVar(&passphrase, "passphrase", "", "decryption passphrase (empty = no encryption)")
	flag.StringVar(&authToken, "auth-token", "", "shared pipeline auth token (or OKO_AUTH_TOKEN)")
	flag.IntVar(&pbkdf2Iter, "pbkdf2-iters", frame.DefaultPBKDF2Iter, "PBKDF2 iterations for clip encryption (1..1000000)")
	flag.Parse()

	if threshold < 0 || threshold > 255 {
		log.Fatalf("threshold must be in [0, 255], got %d", threshold)
	}
	if pbkdf2Iter < 1 || pbkdf2Iter > frame.MaxPBKDF2Iter {
		log.Fatalf("pbkdf2-iters must be in [1, %d], got %d", frame.MaxPBKDF2Iter, pbkdf2Iter)
	}
	frame.SetDefaultPBKDF2Iter(pbkdf2Iter)

	tok := authToken
	if tok == "" {
		tok = os.Getenv("OKO_AUTH_TOKEN")
	}
	if tok == "" {
		log.Fatal("authentication required: set -auth-token or OKO_AUTH_TOKEN")
	}
	frame.SetAuthToken(tok)

	ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer cancel()

	clipChan := make(chan frame.Clip, clipChanCap)
	connSem := make(chan struct{}, connLimit)

	ln, err := net.Listen("tcp", listenAddr)
	if err != nil {
		log.Fatalf("Listen %s: %v", listenAddr, err)
	}
	log.Printf("Listening on %s", listenAddr)

	var acceptWg sync.WaitGroup
	rejected := frame.NewMalformedAttempts()
	go func() {
		<-ctx.Done()
		ln.Close()
	}()
	go func() {
		for {
			conn, err := ln.Accept()
			if err != nil {
				if ctx.Err() != nil {
					return
				}
				log.Printf("Accept error: %v", err)
				continue
			}
			select {
			case connSem <- struct{}{}:
			default:
				log.Printf("Rejecting connection from %s: too many concurrent connections", conn.RemoteAddr())
				conn.Close()
				continue
			}
			acceptWg.Add(1)
			go func(c net.Conn) {
				defer func() { <-connSem }()
				defer acceptWg.Done()
				frame.HandleClipConn(c, rejected, func(clp frame.Clip) {
					clipChan <- clp
				})
			}(conn)
		}
	}()

	// Bounded worker pool, so a burst of clips cannot spin up unbounded
	// goroutines or saturate every core with motion passes.
	var procWg sync.WaitGroup
	procWg.Add(workerCount)
	for w := 0; w < workerCount; w++ {
		go func() {
			defer procWg.Done()
			for clp := range clipChan {
				processClip(clp, terpAddr, threshold, passphrase)
			}
		}()
	}

	<-ctx.Done()
	log.Println("Shutting down...")
	ln.Close()
	if !frame.WaitGroupTimeout(&acceptWg, shutdownGrace) {
		log.Println("Inbound handlers did not drain; closing channel anyway")
	}
	close(clipChan)
	if !frame.WaitGroupTimeout(&procWg, shutdownGrace) {
		log.Println("Workers still busy (downstream retry?); exiting anyway")
	}
	if n := rejected.Total(); n > 0 {
		log.Printf("Rejected %d malformed clips from the network", n)
	}
	log.Println("Exiting")
}

// processClip decrypts, drops clips with no motion, re-encrypts, and forwards
// survivors to terp.
func processClip(clp frame.Clip, terpAddr string, threshold int, passphrase string) {
	if passphrase != "" {
		if err := clp.Decrypt(passphrase); err != nil {
			log.Printf("Decrypt error: %v", err)
			return
		}
	}

	if err := clp.CountChangedPixels(uint8(threshold)); err != nil {
		log.Printf("CountChangedPixels: %v", err)
		return
	}
	if clp.GetHighestMotion() == 0 {
		log.Printf("Dropped no-motion clip (source %s)", clp.SourceData)
		return
	}

	if passphrase != "" {
		if err := clp.Encrypt(passphrase); err != nil {
			log.Printf("Encrypt error: %v", err)
			return
		}
	}

	if err := clp.SendClipRetry(terpAddr, 6); err != nil {
		log.Printf("Forward to terp failed, clip dropped: %v", err)
		return
	}
	log.Printf("Forwarded clip to terp (motion: %d%%)", clp.GetHighestMotion())
}