silo-server/internal/streamenforcer/enforcer_test.go

package streamenforcer

import (
	"context"
	"errors"
	"sort"
	"testing"
	"time"

	"github.com/Silo-Server/silo-server/internal/nodesessions"
	"github.com/Silo-Server/silo-server/internal/playback"
	"github.com/Silo-Server/silo-server/internal/streammonitor"
)

type fakeSource struct {
	snap streammonitor.Snapshot
	err  error
}

func (f fakeSource) Snapshot(context.Context) (streammonitor.Snapshot, error) {
	return f.snap, f.err
}

type fakeRevoker struct {
	revoked []string
	err     error
	seen    map[string]struct{}
	ttls    []time.Duration
}

func (f *fakeRevoker) RevokeSessionForIfAbsent(
	_ context.Context,
	sessionID, _ string,
	ttl time.Duration,
) (bool, error) {
	if f.err != nil {
		return false, f.err
	}
	if f.seen == nil {
		f.seen = make(map[string]struct{})
	}
	if _, ok := f.seen[sessionID]; ok {
		return false, nil
	}
	f.seen[sessionID] = struct{}{}
	f.revoked = append(f.revoked, sessionID)
	f.ttls = append(f.ttls, ttl)
	return true, nil
}

// stream builds a LiveStream whose StartedAt and LastServedAt are both ts.
func stream(id string, user int, ts time.Time) streammonitor.LiveStream {
	return streammonitor.LiveStream{SessionID: id, UserID: user, StartedAt: ts, LastServedAt: ts}
}

// served builds a LiveStream started at `started` but last served at `lastServed`
// (to model a ghost: old activity) — StartedAt independent from real liveness.
func served(id string, user int, started, lastServed time.Time) streammonitor.LiveStream {
	return streammonitor.LiveStream{SessionID: id, UserID: user, StartedAt: started, LastServedAt: lastServed}
}

func TestEvaluateOnce(t *testing.T) {
	base := time.Date(2026, 7, 4, 12, 0, 0, 0, time.UTC)
	limit3 := func(context.Context, int) (int, error) { return 3, nil }

	tests := []struct {
		name       string
		streams    []streammonitor.LiveStream
		limits     LimitFunc
		wantRevoke []string
	}{
		{
			name:       "under cap: no revokes",
			streams:    []streammonitor.LiveStream{stream("a", 1, base), stream("b", 1, base.Add(time.Minute))},
			limits:     limit3,
			wantRevoke: nil,
		},
		{
			name: "over cap: least-recently-served trimmed, freshest kept",
			streams: []streammonitor.LiveStream{
				stream("stale", 1, base),
				stream("old", 1, base.Add(1*time.Minute)),
				stream("keep1", 1, base.Add(2*time.Minute)),
				stream("keep2", 1, base.Add(3*time.Minute)),
				stream("keep3", 1, base.Add(4*time.Minute)),
			},
			limits:     limit3,
			wantRevoke: []string{"stale", "old"},
		},
		{
			name: "ghost trimmed, fresh reconnect survives",
			// Two fresh (served just now) and two ghosts (served long ago), limit 2.
			streams: []streammonitor.LiveStream{
				served("ghost1", 1, base, base),
				served("ghost2", 1, base.Add(time.Second), base.Add(time.Second)),
				served("fresh1", 1, base.Add(10*time.Minute), base.Add(10*time.Minute)),
				served("fresh2", 1, base.Add(10*time.Minute), base.Add(11*time.Minute)),
			},
			limits:     func(context.Context, int) (int, error) { return 2, nil },
			wantRevoke: []string{"ghost1", "ghost2"},
		},
		{
			name:       "limit lookup error: fail open",
			streams:    []streammonitor.LiveStream{stream("a", 1, base), stream("b", 1, base), stream("c", 1, base), stream("d", 1, base)},
			limits:     func(context.Context, int) (int, error) { return 0, errors.New("db down") },
			wantRevoke: nil,
		},
		{
			name:       "unlimited (limit<=0): no revokes",
			streams:    []streammonitor.LiveStream{stream("a", 1, base), stream("b", 1, base), stream("c", 1, base)},
			limits:     func(context.Context, int) (int, error) { return 0, nil },
			wantRevoke: nil,
		},
		{
			name:       "user 0 never enforced",
			streams:    []streammonitor.LiveStream{stream("a", 0, base), stream("b", 0, base), stream("c", 0, base), stream("d", 0, base)},
			limits:     limit3,
			wantRevoke: nil,
		},
	}

	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			rev := &fakeRevoker{}
			e := New(fakeSource{snap: streammonitor.Snapshot{Streams: tc.streams}}, tc.limits, rev, 0)
			if err := e.EvaluateOnce(context.Background()); err != nil {
				t.Fatalf("EvaluateOnce: %v", err)
			}
			got := append([]string{}, rev.revoked...)
			sort.Strings(got)
			want := append([]string{}, tc.wantRevoke...)
			sort.Strings(want)
			if len(got) != len(want) {
				t.Fatalf("revoked = %v, want %v", got, want)
			}
			for i := range got {
				if got[i] != want[i] {
					t.Fatalf("revoked = %v, want %v", got, want)
				}
			}
		})
	}
}

func TestEvaluateOnceProgressOnlySessionIsVictimBeforeServingStreams(t *testing.T) {
	sm := playback.NewSessionManager(0, 0)
	phantom, err := sm.StartSession(7, "profile", 1, playback.PlayDirect, false)
	if err != nil {
		t.Fatalf("StartSession(phantom): %v", err)
	}

	for i := 0; i < 2; i++ {
		session, startErr := sm.StartSession(7, "profile", 10+i, playback.PlayDirect, false)
		if startErr != nil {
			t.Fatalf("StartSession(serving %d): %v", i, startErr)
		}
		if beginErr := sm.BeginTransport(session.ID); beginErr != nil {
			t.Fatalf("BeginTransport(serving %d): %v", i, beginErr)
		}
		if bytesErr := sm.AddServedBytes(session.ID, 1024); bytesErr != nil {
			t.Fatalf("AddServedBytes(serving %d): %v", i, bytesErr)
		}
		if endErr := sm.EndTransport(session.ID); endErr != nil {
			t.Fatalf("EndTransport(serving %d): %v", i, endErr)
		}
	}

	time.Sleep(time.Until(time.Now().Truncate(time.Second).Add(time.Second)) + 10*time.Millisecond)
	if err := sm.UpdateProgress(phantom.ID, 30, false); err != nil {
		t.Fatalf("UpdateProgress(phantom): %v", err)
	}

	source := streammonitor.NewFuncSource(func(context.Context) ([]nodesessions.SessionInfo, error) {
		return streammonitor.LiveLocalSessions(sm, "local"), nil
	})
	rev := &fakeRevoker{}
	e := New(source, func(context.Context, int) (int, error) { return 2, nil }, rev, 0)
	if err := e.EvaluateOnce(context.Background()); err != nil {
		t.Fatalf("EvaluateOnce: %v", err)
	}
	if len(rev.revoked) != 1 || rev.revoked[0] != phantom.ID {
		t.Fatalf("revoked = %v, want progress-only phantom %q", rev.revoked, phantom.ID)
	}
}

func TestEvaluateOnceSourceError(t *testing.T) {
	rev := &fakeRevoker{}
	e := New(fakeSource{err: errors.New("scan failed")}, func(context.Context, int) (int, error) { return 1, nil }, rev, 0)
	if err := e.EvaluateOnce(context.Background()); err == nil {
		t.Fatal("expected error from source")
	}
	if len(rev.revoked) != 0 {
		t.Fatalf("expected no revokes on source error, got %v", rev.revoked)
	}
}

func TestEvaluateOnceRevokesCanonicalLogicalSessionID(t *testing.T) {
	source := streammonitor.NewFuncSource(func(context.Context) ([]nodesessions.SessionInfo, error) {
		return []nodesessions.SessionInfo{
			{SessionID: "transport-a", LogicalSessionID: "logical-a", AuthUserID: 7, LastServedAt: "2026-07-30T01:00:00Z"},
			{SessionID: "transport-b", LogicalSessionID: "logical-b", AuthUserID: 7, LastServedAt: "2026-07-30T02:00:00Z"},
		}, nil
	})
	rev := &fakeRevoker{}
	e := New(source, func(context.Context, int) (int, error) { return 1, nil }, rev, 0)
	if err := e.EvaluateOnce(context.Background()); err != nil {
		t.Fatalf("EvaluateOnce: %v", err)
	}
	if len(rev.revoked) != 1 || rev.revoked[0] != "logical-a" {
		t.Fatalf("revoked = %v, want logical-a", rev.revoked)
	}
}

func TestEvaluateOnceDoesNotExtendExistingAutomatedRevocation(t *testing.T) {
	base := time.Now()
	source := fakeSource{snap: streammonitor.Snapshot{Streams: []streammonitor.LiveStream{
		stream("keep", 7, base.Add(time.Minute)),
		stream("victim", 7, base),
	}}}
	rev := &fakeRevoker{}
	e := New(source, func(context.Context, int) (int, error) { return 1, nil }, rev, 0)
	if err := e.EvaluateOnce(context.Background()); err != nil {
		t.Fatal(err)
	}
	if err := e.EvaluateOnce(context.Background()); err != nil {
		t.Fatal(err)
	}
	if len(rev.revoked) != 1 {
		t.Fatalf("revocation writes = %d, want 1", len(rev.revoked))
	}
	if len(rev.ttls) != 1 || rev.ttls[0] != playback.MaxTokenTTL {
		t.Fatalf("revocation TTLs = %v, want [%v]", rev.ttls, playback.MaxTokenTTL)
	}
}

func TestParseRevocationTTL(t *testing.T) {
	if got, err := ParseRevocationTTL(""); err != nil || got != playback.MaxTokenTTL {
		t.Fatalf("default = %v, %v; want %v", got, err, playback.MaxTokenTTL)
	}
	if got, err := ParseRevocationTTL("6h"); err != nil || got != 6*time.Hour {
		t.Fatalf("valid = %v, %v; want 6h", got, err)
	}
	for _, raw := range []string{"bad", "4m59s", "24h1s"} {
		if _, err := ParseRevocationTTL(raw); err == nil {
			t.Fatalf("ParseRevocationTTL(%q) accepted invalid value", raw)
		}
	}
}