oko_public/frame/network.go

package frame

import (
	"bufio"
	"bytes"
	"crypto/hmac"
	"encoding/binary"
	"encoding/gob"
	"errors"
	"fmt"
	"io"
	"log"
	"net"
	"strconv"
	"strings"
	"sync"
	"time"
)

// Network timeouts for a single TCP hop: the auth handshake and the clip
// transfer must both complete within their deadlines or the connection is
// dropped (fail-closed).
const (
	authTimeout      = 10 * time.Second
	ClipReadTimeout  = 30 * time.Second
	ClipWriteTimeout = 30 * time.Second
)

// maxAuthLine bounds the length of the auth token line, so a peer that never
// sends a newline cannot grow a buffer forever (the read deadline still
// applies).
const maxAuthLine = 4096

// Wire framing for clips: every gob-encoded Clip is sent as a 4-byte
// big-endian length followed by that many payload bytes. The receiving side
// refuses any length above MaxWireBytes before allocating, so a single
// crafted message cannot force a multi-gigabyte allocation on the decode
// path.
const (
	wireHeaderSize  = 4
	MaxClipWireSize = MaxClipBytes + (16 << 20) // 80 MiB: pixel data + guids/timestamps/detections + gob overhead
)

// authToken is the shared secret required on every hop. An empty token is
// never usable: receivers fail closed rather than accepting unauthenticated
// connections, and senders require a configured token to write the handshake.
var (
	authMu    sync.RWMutex
	authToken string
)

// SetAuthToken configures the shared token required by every send and receive
// on this process. The token is trimmed of surrounding whitespace so a
// config line with a stray newline keeps matching the handshake.
func SetAuthToken(tok string) {
	authMu.Lock()
	defer authMu.Unlock()
	authToken = strings.TrimSpace(tok)
}

func getAuthToken() string {
	authMu.RLock()
	defer authMu.RUnlock()
	return authToken
}

func isValidHostPort(s string) bool {
	host, port, err := net.SplitHostPort(s)
	if err != nil {
		return false
	}
	if host == "" {
		return false
	}
	p, err := strconv.Atoi(port)
	return err == nil && p >= 1 && p <= 65535
}

// writeAuthToken emits the auth line to conn. It returns an error when no
// token is configured so a misconfigured sender never silently transmits an
// unauthenticated clip.
func writeAuthToken(conn net.Conn) error {
	tok := getAuthToken()
	if tok == "" {
		return errors.New("auth: no token configured (fail-closed)")
	}
	if err := conn.SetWriteDeadline(time.Now().Add(authTimeout)); err != nil {
		return err
	}
	_, err := fmt.Fprintf(conn, "%s\n", tok)
	return err
}

// readAuthToken consumes the auth line from conn and returns a reader for the
// gob stream that follows. A missing, truncated, or mismatched token fails
// the connection (constant-time comparison), and an unconfigured token fails
// closed rather than accepting the peer silently.
func readAuthToken(conn net.Conn) (io.Reader, error) {
	tok := getAuthToken()
	if tok == "" {
		return nil, errors.New("auth: no token configured on receiver (fail-closed)")
	}
	if err := conn.SetReadDeadline(time.Now().Add(authTimeout)); err != nil {
		return nil, err
	}
	br := bufio.NewReader(conn)

	line, err := readAuthLine(br)
	if err != nil {
		return nil, fmt.Errorf("auth handshake read: %w", err)
	}
	if !hmac.Equal([]byte(strings.TrimSpace(line)), []byte(tok)) {
		return nil, errors.New("authenticate: token mismatch")
	}
	return io.MultiReader(br, conn), nil
}

// readAuthLine reads a single '\n'-terminated line, bounding how many bytes
// are consumed so a hostile peer cannot stream data without a newline.
func readAuthLine(br *bufio.Reader) (string, error) {
	var sb strings.Builder
	for {
		b, err := br.ReadByte()
		if err != nil {
			return "", err
		}
		if b == '\n' {
			return sb.String(), nil
		}
		if sb.Len() >= maxAuthLine {
			return "", errors.New("auth handshake line exceeds limit")
		}
		sb.WriteByte(b)
	}
}

// writeWireClip encodes clp as a length-prefixed gob stream to w. The length
// prefix lets the receiver bound allocation before decoding.
func writeWireClip(w io.Writer, clp *Clip) (int64, error) {
	var buf bytes.Buffer
	if err := gob.NewEncoder(&buf).Encode(clp); err != nil {
		return 0, err
	}
	if buf.Len() > MaxClipWireSize {
		return 0, fmt.Errorf("clip wire encoding of %d bytes exceeds limit of %d", buf.Len(), MaxClipWireSize)
	}
	var hdr [wireHeaderSize]byte
	binary.BigEndian.PutUint32(hdr[:], uint32(buf.Len()))
	if _, err := w.Write(hdr[:]); err != nil {
		return 0, err
	}
	n, err := w.Write(buf.Bytes())
	return int64(wireHeaderSize) + int64(n), err
}

// DecodeWireClip reads a length-prefixed Clip produced by writeWireClip. The
// declared length is validated against MaxClipWireSize before any allocation,
// so a hostile peer cannot force a huge single allocation. Validation of the
// Clip's structure is the caller's job (see Clip.Validate).
func DecodeWireClip(r io.Reader) (Clip, error) {
	var hdr [wireHeaderSize]byte
	if _, err := io.ReadFull(r, hdr[:]); err != nil {
		return Clip{}, err
	}
	n := binary.BigEndian.Uint32(hdr[:])
	if n > MaxClipWireSize {
		return Clip{}, fmt.Errorf("clip wire length %d exceeds limit of %d", n, MaxClipWireSize)
	}
	payload := make([]byte, n)
	if _, err := io.ReadFull(r, payload); err != nil {
		return Clip{}, err
	}
	var clp Clip
	if err := gob.NewDecoder(bytes.NewReader(payload)).Decode(&clp); err != nil {
		return Clip{}, err
	}
	return clp, nil
}

// Send delivers the clip to target: auth handshake, then a length-prefixed
// gob encode within a write deadline.
func (clp *Clip) Send(target string) error {
	if !isValidHostPort(target) {
		return fmt.Errorf("invalid send address: %s", target)
	}

	conn, err := net.DialTimeout("tcp", target, 10*time.Second)
	if err != nil {
		return fmt.Errorf("dial %s: %w", target, err)
	}
	defer conn.Close()

	if err := writeAuthToken(conn); err != nil {
		return fmt.Errorf("auth write to %s: %w", target, err)
	}
	if err := conn.SetWriteDeadline(time.Now().Add(ClipWriteTimeout)); err != nil {
		return err
	}
	if _, err := writeWireClip(conn, clp); err != nil {
		return fmt.Errorf("encode to %s: %w", target, err)
	}
	return nil
}

// SendClipRetry attempts to deliver clp to target up to attempts times, with a
// backoff that doubles from 100ms up to a 500ms cap. The leaf services use it
// so a transient downstream outage (restart, network blip) does not silently
// lose footage.
func (clp *Clip) SendClipRetry(target string, attempts int) error {
	if attempts < 1 {
		attempts = 1
	}

	var lastErr error
	for i := 1; i <= attempts; i++ {
		if err := clp.Send(target); err != nil {
			lastErr = err
			delay := 100 * time.Millisecond * (1 << (i - 1))
			if delay > 500*time.Millisecond {
				delay = 500 * time.Millisecond
			}
			log.Printf("send to %s failed (attempt %d/%d): %v; retrying in %s",
				target, i, attempts, err, delay)
			time.Sleep(delay)
			continue
		}
		return nil
	}
	return fmt.Errorf("send to %s failed after %d attempts: %w", target, attempts, lastErr)
}