333 lines
8.7 KiB
Go
333 lines
8.7 KiB
Go
package main
|
|
|
|
import (
|
|
"context"
|
|
"flag"
|
|
"fmt"
|
|
"log"
|
|
"math"
|
|
"net"
|
|
"os"
|
|
"os/signal"
|
|
"strings"
|
|
"sync"
|
|
"syscall"
|
|
"time"
|
|
|
|
"frame"
|
|
"gocv.io/x/gocv"
|
|
"maunium.net/go/mautrix/id"
|
|
)
|
|
|
|
func main() {
|
|
var (
|
|
listenAddr string
|
|
passphrase string
|
|
authToken string
|
|
matrixHS string
|
|
matrixUser string
|
|
matrixTokenFile string
|
|
clipsRoom string
|
|
detectionsRoom string
|
|
viewRoom string
|
|
pbkdf2Iter int
|
|
)
|
|
|
|
flag.StringVar(&listenAddr, "listen", ":8082", "address to listen for clips from terp")
|
|
flag.StringVar(&passphrase, "passphrase", "", "decryption passphrase (empty = no decryption)")
|
|
flag.StringVar(&authToken, "auth-token", "", "shared pipeline auth token (or OKO_AUTH_TOKEN)")
|
|
flag.StringVar(&matrixHS, "matrix-homeserver", "", "Matrix homeserver URL (required)")
|
|
flag.StringVar(&matrixUser, "matrix-user", "", "Matrix user ID, e.g. @oko:example.org (required)")
|
|
flag.StringVar(&matrixTokenFile, "matrix-token-file", "", "file containing the Matrix access token (or OKO_MATRIX_TOKEN env)")
|
|
flag.StringVar(&clipsRoom, "clips-room", "", "room ID for all motion clips (required)")
|
|
flag.StringVar(&detectionsRoom, "detections-room", "", "room ID for clips with detections (required)")
|
|
flag.StringVar(&viewRoom, "view-room", "", "room ID for clips with minimal caption only, no JSON metadata (required)")
|
|
flag.IntVar(&pbkdf2Iter, "pbkdf2-iters", frame.DefaultPBKDF2Iter, "PBKDF2 iterations for clip encryption (1..1000000)")
|
|
flag.Parse()
|
|
|
|
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)
|
|
|
|
token := os.Getenv("OKO_MATRIX_TOKEN")
|
|
if matrixTokenFile != "" {
|
|
data, err := os.ReadFile(matrixTokenFile)
|
|
if err != nil {
|
|
log.Fatalf("Read matrix token file: %v", err)
|
|
}
|
|
token = strings.TrimSpace(string(data))
|
|
}
|
|
if matrixHS == "" || matrixUser == "" || token == "" {
|
|
log.Fatal("matrix-homeserver, matrix-user, and a token (OKO_MATRIX_TOKEN or -matrix-token-file) are required")
|
|
}
|
|
|
|
ms, err := newMatrixStore(matrixHS, matrixUser, token, clipsRoom, detectionsRoom, viewRoom)
|
|
if err != nil {
|
|
log.Fatalf("Initialize Matrix backend: %v", err)
|
|
}
|
|
log.Printf("Matrix backend ready")
|
|
|
|
ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
|
|
defer cancel()
|
|
|
|
ln, err := net.Listen("tcp", listenAddr)
|
|
if err != nil {
|
|
log.Fatalf("Listen %s: %v", listenAddr, err)
|
|
}
|
|
log.Printf("Listening on %s", listenAddr)
|
|
|
|
var wg sync.WaitGroup
|
|
rejected := frame.NewMalformedAttempts()
|
|
// Cap concurrent connections so a flood (or slowloris-style connection
|
|
// that parks on the auth handshake) cannot exhaust goroutines/FDs.
|
|
connSem := make(chan struct{}, 16)
|
|
// Cap concurrent Matrix ingest so a flood of clips cannot spawn unbounded
|
|
// encode+upload work. Deliberately acquired only inside the authenticated
|
|
// callback, so an unauthenticated connection cannot occupy an ingest slot.
|
|
sem := make(chan struct{}, 8)
|
|
go func() {
|
|
<-ctx.Done()
|
|
ln.Close()
|
|
}()
|
|
|
|
for {
|
|
conn, err := ln.Accept()
|
|
if err != nil {
|
|
if ctx.Err() != nil {
|
|
break
|
|
}
|
|
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
|
|
}
|
|
wg.Add(1)
|
|
go func(c net.Conn) {
|
|
defer func() { <-connSem }()
|
|
defer wg.Done()
|
|
frame.HandleClipConn(c, rejected, func(clp frame.Clip) {
|
|
sem <- struct{}{}
|
|
defer func() { <-sem }()
|
|
handleClip(ctx, ms, passphrase, clp)
|
|
})
|
|
}(conn)
|
|
}
|
|
|
|
if !frame.WaitGroupTimeout(&wg, 5*time.Second) {
|
|
log.Printf("In-flight handlers still working; exiting anyway")
|
|
}
|
|
if n := rejected.Total(); n > 0 {
|
|
log.Printf("Rejected %d malformed clips from the network", n)
|
|
}
|
|
log.Println("Exiting")
|
|
}
|
|
|
|
func handleClip(parent context.Context, ms *matrixStore, passphrase string, clp frame.Clip) {
|
|
targets := routeTargets(&clp, ms.clipsRoom, ms.detectionsRoom)
|
|
if len(targets) == 0 {
|
|
log.Printf("Dropped clip with no motion (source %s)", clp.SourceData)
|
|
return
|
|
}
|
|
|
|
if passphrase != "" {
|
|
if err := clp.Decrypt(passphrase); err != nil {
|
|
log.Printf("Decrypt error: %v", err)
|
|
return
|
|
}
|
|
}
|
|
|
|
frames := clp.Sublimate()
|
|
if len(frames) == 0 {
|
|
return
|
|
}
|
|
|
|
// Timebox encode + upload + posts together; this is the only archival hop,
|
|
// so a transient homeserver blip should be retried rather than dropping
|
|
// the footage.
|
|
ctx, cancel := context.WithTimeout(parent, 5*time.Minute)
|
|
defer cancel()
|
|
|
|
fps := computeFPS(clp.Timestamps)
|
|
mp4, err := encodeMP4(frames, int(clp.Width), int(clp.Height), fps)
|
|
if err != nil {
|
|
log.Printf("Encode clip error: %v", err)
|
|
return
|
|
}
|
|
|
|
filename := fmt.Sprintf("%x-%d.mp4", clp.Guids[0], clp.Timestamps[0])
|
|
var mxc id.ContentURI
|
|
if err := withRetry(ctx, 3, func() error {
|
|
mxc, err = ms.uploadClip(ctx, filename, mp4)
|
|
return err
|
|
}); err != nil {
|
|
log.Printf("Upload clip media: %v", err)
|
|
return
|
|
}
|
|
|
|
failed := 0
|
|
for _, roomID := range targets {
|
|
var eventID id.EventID
|
|
if err := withRetry(ctx, 2, func() error {
|
|
eventID, err = ms.sendClipToRoom(ctx, roomID, clp, filename, mxc, len(mp4), fps)
|
|
return err
|
|
}); err != nil {
|
|
failed++
|
|
log.Printf("Post clip to %s: %v", roomID, err)
|
|
continue
|
|
}
|
|
log.Printf("Stored clip (%d frames, motion=%d%%, %.1f fps) in %s as %s",
|
|
len(frames), clp.GetHighestMotion(), fps, roomID, eventID)
|
|
}
|
|
|
|
var viewID id.EventID
|
|
if err := withRetry(ctx, 2, func() error {
|
|
viewID, err = ms.sendClipView(ctx, ms.viewRoom, clp, mxc, len(mp4), fps)
|
|
return err
|
|
}); err != nil {
|
|
failed++
|
|
log.Printf("Post clip to view room: %v", err)
|
|
} else {
|
|
log.Printf("Stored clip (view, minimal) in %s as %s", ms.viewRoom, viewID)
|
|
}
|
|
|
|
if failed > 0 {
|
|
log.Printf("Failed to post clip to %d of %d target rooms", failed, len(targets)+1)
|
|
}
|
|
}
|
|
|
|
// withRetry runs fn up to attempts times with capped exponential backoff,
|
|
// returning early when ctx is done.
|
|
func withRetry(ctx context.Context, attempts int, fn func() error) error {
|
|
if attempts < 1 {
|
|
attempts = 1
|
|
}
|
|
var lastErr error
|
|
for i := 1; i <= attempts; i++ {
|
|
if err := fn(); err != nil {
|
|
lastErr = err
|
|
delay := 500 * time.Millisecond * time.Duration(1<<(i-1))
|
|
if delay > 5*time.Second {
|
|
delay = 5 * time.Second
|
|
}
|
|
select {
|
|
case <-ctx.Done():
|
|
return lastErr
|
|
case <-time.After(delay):
|
|
}
|
|
continue
|
|
}
|
|
return nil
|
|
}
|
|
return lastErr
|
|
}
|
|
|
|
func computeFPS(timestamps []uint64) float64 {
|
|
if len(timestamps) < 2 {
|
|
return 30.0
|
|
}
|
|
|
|
// Average of the consecutive inter-frame deltas, so a single irregular
|
|
// pair of timestamps cannot skew the reported rate.
|
|
var sum uint64
|
|
deltas := 0
|
|
for i := 1; i < len(timestamps); i++ {
|
|
if timestamps[i] <= timestamps[i-1] {
|
|
continue
|
|
}
|
|
sum += timestamps[i] - timestamps[i-1]
|
|
deltas++
|
|
}
|
|
if deltas == 0 {
|
|
return 30.0
|
|
}
|
|
|
|
avg := float64(sum) / float64(deltas)
|
|
fps := 1e9 / avg
|
|
if fps <= 0 || fps > 120 {
|
|
return 30.0
|
|
}
|
|
return math.Round(fps*100) / 100
|
|
}
|
|
|
|
// encodeMP4 renders the clip's frames to an H.264 MP4 in a temporary file and
|
|
// returns the bytes. The temp file is removed before returning.
|
|
func encodeMP4(frames []frame.Frame, width, height int, fps float64) ([]byte, error) {
|
|
tmp, err := os.CreateTemp("", "oko-clip-*.mp4")
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
path := tmp.Name()
|
|
tmp.Close()
|
|
defer os.Remove(path)
|
|
|
|
if err := writeVideo(path, frames, width, height, fps); err != nil {
|
|
return nil, err
|
|
}
|
|
return os.ReadFile(path)
|
|
}
|
|
|
|
func writeVideo(path string, frames []frame.Frame, width, height int, fps float64) error {
|
|
isColor := true
|
|
if len(frames) > 0 && frames[0].Channels == 1 {
|
|
isColor = false
|
|
}
|
|
|
|
var vw *gocv.VideoWriter
|
|
var err error
|
|
|
|
for _, codec := range []string{"avc1", "H264", "MJPG"} {
|
|
vw, err = gocv.VideoWriterFile(path, codec, fps, width, height, isColor)
|
|
if err == nil && vw.IsOpened() {
|
|
break
|
|
}
|
|
if vw != nil {
|
|
vw.Close()
|
|
}
|
|
vw = nil
|
|
}
|
|
|
|
if vw == nil {
|
|
return fmt.Errorf("failed to open video writer with any codec: %v", err)
|
|
}
|
|
defer vw.Close()
|
|
|
|
written := 0
|
|
for i := range frames {
|
|
mat, err := frames[i].ToMat()
|
|
if err != nil {
|
|
log.Printf("ToMat error frame %d: %v", i, err)
|
|
continue
|
|
}
|
|
if err := vw.Write(mat); err != nil {
|
|
mat.Close()
|
|
return fmt.Errorf("write frame %d: %w", i, err)
|
|
}
|
|
mat.Close()
|
|
written++
|
|
}
|
|
|
|
if written == 0 {
|
|
return fmt.Errorf("no frames could be encoded")
|
|
}
|
|
if written < len(frames) {
|
|
log.Printf("Encoded %d of %d frames (partial clip)", written, len(frames))
|
|
}
|
|
|
|
return nil
|
|
}
|