oko_public/frame/receive_test.go

package frame

import (
	"errors"
	"net"
	"strings"
	"testing"
	"time"

	"gocv.io/x/gocv"
)

func validTestClip() Clip {
	return Clip{
		PixelMats:  [][]byte{{1, 2, 3}, {4, 5, 6}},
		Guids:      [][]byte{{1}, {2}},
		Timestamps: []uint64{100, 200},
		Width:      1,
		Height:     1,
		Types:      gocv.MatTypeCV8UC3,
		Channels:   3,
	}
}

func TestValidate_ValidClip(t *testing.T) {
	clp := validTestClip()
	if err := clp.Validate(); err != nil {
		t.Errorf("valid clip rejected: %v", err)
	}
}

func TestValidate_EmptyClipIsValid(t *testing.T) {
	clp := Clip{}
	if err := clp.Validate(); err != nil {
		t.Errorf("empty clip rejected: %v", err)
	}
}

func TestValidate_MissingGuids(t *testing.T) {
	clp := validTestClip()
	clp.Guids = nil
	if err := clp.Validate(); !errors.Is(err, ErrMalformedClip) {
		t.Errorf("expected ErrMalformedClip, got %v", err)
	}
}

func TestValidate_MissingTimestamps(t *testing.T) {
	clp := validTestClip()
	clp.Timestamps = nil
	if err := clp.Validate(); !errors.Is(err, ErrMalformedClip) {
		t.Errorf("expected ErrMalformedClip, got %v", err)
	}
}

func TestValidate_ZeroDimensions(t *testing.T) {
	clp := validTestClip()
	clp.Width = 0
	if err := clp.Validate(); !errors.Is(err, ErrMalformedClip) {
		t.Errorf("expected ErrMalformedClip for zero width, got %v", err)
	}
}

func TestValidate_HugeDimensions(t *testing.T) {
	clp := validTestClip()
	clp.Width = MaxClipDim + 1
	if err := clp.Validate(); !errors.Is(err, ErrMalformedClip) {
		t.Errorf("expected ErrMalformedClip for oversized dimensions, got %v", err)
	}
}

func TestValidate_TooManyFrames(t *testing.T) {
	clp := Clip{
		PixelMats:  make([][]byte, MaxClipFrames+1),
		Guids:      make([][]byte, MaxClipFrames+1),
		Timestamps: make([]uint64, MaxClipFrames+1),
		Width:      1,
		Height:     1,
		Channels:   1,
	}
	if err := clp.Validate(); !errors.Is(err, ErrMalformedClip) {
		t.Errorf("expected ErrMalformedClip for excessive frames, got %v", err)
	}
}

func TestValidate_InvalidChannels(t *testing.T) {
	clp := validTestClip()
	clp.Channels = 5
	if err := clp.Validate(); !errors.Is(err, ErrMalformedClip) {
		t.Errorf("expected ErrMalformedClip for channel count 5, got %v", err)
	}
}

func TestValidate_EncryptedClip(t *testing.T) {
	clp := validTestClip()
	if err := clp.Encrypt("secret"); err != nil {
		t.Fatalf("Encrypt() error: %v", err)
	}
	if err := clp.Validate(); err != nil {
		t.Errorf("valid encrypted clip rejected: %v", err)
	}
}

func TestValidate_WrongFrameSize(t *testing.T) {
	clp := validTestClip()
	clp.PixelMats[0] = []byte{1, 2, 3, 4} // 1x1x3 declared, 4 bytes sent
	if err := clp.Validate(); !errors.Is(err, ErrMalformedClip) {
		t.Errorf("expected ErrMalformedClip for undersized frame, got %v", err)
	}
}

func TestValidate_OverBudgetTotalBytes(t *testing.T) {
	// 22 frames of a legitimate 1024x1024x3 frame (3 MiB each) exceed the
	// 64 MiB total cap even though every frame is individually well-formed.
	expected := 1024 * 1024 * 3
	frames := make([][]byte, 22)
	for i := range frames {
		frames[i] = make([]byte, expected)
	}
	clp := Clip{
		PixelMats:  frames,
		Guids:      make([][]byte, 22),
		Timestamps: make([]uint64, 22),
		Width:      1024,
		Height:     1024,
		Types:      gocv.MatTypeCV8UC3,
		Channels:   3,
	}
	err := clp.Validate()
	if !errors.Is(err, ErrMalformedClip) {
		t.Errorf("expected ErrMalformedClip for over-budget clip, got %v", err)
	}
	if !strings.Contains(err.Error(), "exceeds limit") {
		t.Errorf("expected the total-byte error, got: %v", err)
	}
}

func TestMalformedAttempts_CountsPerSource(t *testing.T) {
	m := NewMalformedAttempts()
	if n := m.Reject("10.0.0.1:5000"); n != 1 {
		t.Errorf("first reject: got %d, want 1", n)
	}
	if n := m.Reject("10.0.0.1:5000"); n != 2 {
		t.Errorf("second reject from same source: got %d, want 2", n)
	}
	if n := m.Reject("10.0.0.2:5000"); n != 1 {
		t.Errorf("reject from new source: got %d, want 1", n)
	}
	if total := m.Total(); total != 3 {
		t.Errorf("Total: got %d, want 3", total)
	}
}

func TestHandleClipConn_ValidClipDelivered(t *testing.T) {
	SetAuthToken("test-token")
	t.Cleanup(func() { SetAuthToken("") })

	server, client := net.Pipe()
	defer server.Close()
	defer client.Close()

	delivered := make(chan Clip, 1)
	go HandleClipConn(server, NewMalformedAttempts(), func(clp Clip) {
		delivered <- clp
	})

	clp := validTestClip()
	if err := writeAuthToken(client); err != nil {
		t.Fatalf("auth: %v", err)
	}
	if _, err := writeWireClip(client, &clp); err != nil {
		t.Fatalf("write: %v", err)
	}

	select {
	case got := <-delivered:
		if len(got.PixelMats) != 2 {
			t.Errorf("delivered clip has %d frames, want 2", len(got.PixelMats))
		}
	case <-time.After(3 * time.Second):
		t.Fatal("timed out waiting for valid clip")
	}
}

func TestHandleClipConn_MalformedClipRejectedAndCounted(t *testing.T) {
	// The exact shape that used to crash the daemons: a clip whose
	// PixelMats outlives Guids/Timestamps. Sublimate(indexed) panics on it.
	SetAuthToken("test-token")
	t.Cleanup(func() { SetAuthToken("") })

	server, client := net.Pipe()
	defer server.Close()
	defer client.Close()

	rej := NewMalformedAttempts()
	done := make(chan struct{})
	go func() {
		HandleClipConn(server, rej, func(Clip) {
			t.Error("callback must not run for a malformed clip")
		})
		close(done)
	}()

	bad := Clip{
		PixelMats:  [][]byte{{1, 2, 3}},
		Guids:      nil,
		Timestamps: []uint64{100},
		Width:      1,
		Height:     1,
		Types:      gocv.MatTypeCV8UC3,
		Channels:   3,
	}
	if err := writeAuthToken(client); err != nil {
		t.Fatalf("auth: %v", err)
	}
	if _, err := writeWireClip(client, &bad); err != nil {
		t.Fatalf("write: %v", err)
	}

	select {
	case <-done:
	case <-time.After(3 * time.Second):
		t.Fatal("HandleClipConn did not return for malformed clip")
	}
	if rej.Total() != 1 {
		t.Errorf("rejected attempts: got %d, want 1", rej.Total())
	}
}

func TestHandleClipConn_GarbageDoesNotPanic(t *testing.T) {
	SetAuthToken("test-token")
	t.Cleanup(func() { SetAuthToken("") })

	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	defer ln.Close()

	done := make(chan struct{})
	go func() {
		conn, err := ln.Accept()
		if err != nil {
			t.Errorf("accept: %v", err)
			return
		}
		HandleClipConn(conn, NewMalformedAttempts(), func(Clip) {
			t.Error("callback must not run for garbage")
		})
		close(done)
	}()

	conn, err := net.Dial("tcp", ln.Addr().String())
	if err != nil {
		t.Fatalf("dial: %v", err)
	}
	if _, err := conn.Write([]byte("test-token" + "\n")); err != nil {
		t.Fatalf("write auth: %v", err)
	}
	if _, err := conn.Write([]byte("this is not a gob stream at all")); err != nil {
		t.Fatalf("write: %v", err)
	}
	conn.Close()

	select {
	case <-done:
	case <-time.After(3 * time.Second):
		t.Fatal("HandleClipConn did not return for garbage input")
	}
}