oko_public/frame/receive.go

package frame

import (
	"errors"
	"fmt"
	"log"
	"net"
	"sync"
	"time"
)

// ErrMalformedClip marks a Clip that decoded but failed structural
// validation: the sign of a buggy sender or an attacker, as opposed to a
// truncated or corrupt stream.
var ErrMalformedClip = errors.New("malformed clip")

// Limits enforced when receiving a clip, so a single untrusted Clip cannot
// exhaust memory or fan out unbounded work downstream.
const (
	MaxClipFrames = 6000     // generous headroom over any realistic clip length
	MaxClipDim    = 8192     // pixels per edge; bounds the frame buffer size
	MaxClipBytes  = 64 << 20 // 64 MiB of pixel data (encrypted frames include their overhead)
)

// Validate checks a decoded Clip against its structural invariants so that
// downstream use (Sublimate, ToMat, encoding) can never panic or allocate
// wildly. Empty clips are considered valid and are simply ignored by
// consumers. Each frame's pixel buffer must match the declared dimensions
// (plaintext), or carry the fixed encryption overhead on top of that, so a
// sender cannot advertise one size but stream another.
func (c *Clip) Validate() error {
	if len(c.PixelMats) == 0 {
		return nil
	}
	if ok, err := c.CheckLenCorrelations(); !ok {
		return fmt.Errorf("%w: %v", ErrMalformedClip, err)
	}
	if n := len(c.PixelMats); n > MaxClipFrames {
		return fmt.Errorf("%w: %d frames exceeds the limit of %d", ErrMalformedClip, n, MaxClipFrames)
	}
	if c.Width == 0 || c.Height == 0 || c.Height > MaxClipDim || c.Width > MaxClipDim {
		return fmt.Errorf("%w: implausible frame dimensions %dx%d", ErrMalformedClip, c.Width, c.Height)
	}
	if c.Channels < 1 || c.Channels > 4 {
		return fmt.Errorf("%w: implausible channel count %d", ErrMalformedClip, c.Channels)
	}
	if !isCV8U(c.Types) {
		return fmt.Errorf("%w: unsupported MatType depth %d (only 8-bit unsigned is accepted)", ErrMalformedClip, c.Types)
	}
	if n := matTypeChannels(c.Types); n != c.Channels {
		return fmt.Errorf("%w: MatType declares %d channels but Clip.Channels is %d", ErrMalformedClip, n, c.Channels)
	}
	if c.PBKDF2Iter < 0 || c.PBKDF2Iter > MaxPBKDF2Iter {
		return fmt.Errorf("%w: implausible PBKDF2 iteration count %d", ErrMalformedClip, c.PBKDF2Iter)
	}

	expected := int(c.Width) * int(c.Height) * c.Channels
	clipEncrypted := expected + clipEncryptedOverhead
	perFrameEncrypted := expected + encryptedOverhead
	total := 0
	for i, px := range c.PixelMats {
		switch len(px) {
		case 0:
			continue
		case expected, clipEncrypted, perFrameEncrypted:
			total += len(px)
		default:
			return fmt.Errorf("%w: frame %d is %d bytes, want %d (or %d clip-level / %d per-frame encrypted)",
				ErrMalformedClip, i, len(px), expected, clipEncrypted, perFrameEncrypted)
		}
	}
	if total > MaxClipBytes {
		return fmt.Errorf("%w: clip pixel data totals %d bytes, exceeds limit of %d",
			ErrMalformedClip, total, MaxClipBytes)
	}
	return nil
}

// malformedAttemptSources caps the number of distinct source addresses the
// malformed-clip tracker retains, so a peer cycling many addresses cannot
// grow the map without bound.
const malformedAttemptSources = 256

// MalformedAttempts counts rejected malformed clip attempts per source
// address, so repeated bad senders are visible in the logs.
type MalformedAttempts struct {
	mu    sync.Mutex
	count map[string]int
	total int
}

// NewMalformedAttempts returns an empty malformed-clip attempt tracker.
func NewMalformedAttempts() *MalformedAttempts {
	return &MalformedAttempts{count: make(map[string]int)}
}

// Reject records one malformed clip attempt from src, logs it with the
// source and the running per-source count, and returns that count.
func (m *MalformedAttempts) Reject(src string) int {
	m.mu.Lock()
	defer m.mu.Unlock()
	if _, known := m.count[src]; !known && len(m.count) >= malformedAttemptSources {
		m.total++
		log.Printf("malformed clip rejected from %s (tracked-sources cap reached)", src)
		return 0
	}
	m.count[src]++
	n := m.count[src]
	m.total++
	log.Printf("malformed clip rejected from %s (attempt #%d)", src, n)
	return n
}

// Total returns the total number of rejected attempts across all sources.
func (m *MalformedAttempts) Total() int {
	m.mu.Lock()
	defer m.mu.Unlock()
	return m.total
}

// decodeClip reads and validates a single Clip from conn. The auth handshake
// runs first (writing nothing, fail-closed); corrupt or truncated streams
// yield a plain decode error; structurally invalid Clips yield an error
// wrapping ErrMalformedClip, counted and logged against rej when it is
// non-nil.
func decodeClip(conn net.Conn, rej *MalformedAttempts) (Clip, error) {
	var clp Clip
	reader, err := readAuthToken(conn)
	if err != nil {
		return clp, err
	}
	if err := conn.SetReadDeadline(time.Now().Add(ClipReadTimeout)); err != nil {
		return clp, err
	}
	clp, err = DecodeWireClip(reader)
	if err != nil {
		return clp, err
	}
	if err := clp.Validate(); err != nil {
		if rej != nil {
			rej.Reject(remoteAddr(conn))
		}
		return clp, err
	}
	return clp, nil
}

// HandleClipConn runs the standard single-clip ingest flow for conn: closes
// it, decodes and validates exactly one Clip, and hands it to fn. Panics and
// malformed input are contained and logged, so a rogue connection can never
// crash the process.
func HandleClipConn(conn net.Conn, rej *MalformedAttempts, fn func(Clip)) {
	defer conn.Close()
	defer func() {
		if r := recover(); r != nil {
			log.Printf("panic while ingesting clip from %s: %v", remoteAddr(conn), r)
		}
	}()
	clp, err := decodeClip(conn, rej)
	if err != nil {
		if !errors.Is(err, ErrMalformedClip) {
			log.Printf("decode error from %s: %v", remoteAddr(conn), err)
		}
		return
	}
	fn(clp)
}

func remoteAddr(conn net.Conn) string {
	if conn == nil {
		return "<unknown>"
	}
	return conn.RemoteAddr().String()
}