package-updater/internal/e2e/agent_master_test.go

// Package e2e runs a real agent connection against a real hub over TLS with
// a pinned key, in one process, with in-memory storage and fake Docker.
package e2e

import (
	"context"
	"crypto/tls"
	"io"
	"log"
	"log/slog"
	"net"
	"net/http"
	"net/http/httptest"
	"sort"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"github.com/shcizo/package-updater/internal/agent/conn"
	"github.com/shcizo/package-updater/internal/agent/updater"
	"github.com/shcizo/package-updater/internal/agenttoken"
	"github.com/shcizo/package-updater/internal/imageref"
	"github.com/shcizo/package-updater/internal/master/events"
	"github.com/shcizo/package-updater/internal/master/hub"
	"github.com/shcizo/package-updater/internal/master/model"
	"github.com/shcizo/package-updater/internal/pin"
	"github.com/shcizo/package-updater/internal/protocol"
)

type memStore struct {
	mu        sync.Mutex
	agents    map[string]model.Agent
	lookups   int
	snapshots int
	services  map[string]protocol.Service

	// Jobs: in-memory twins of the store methods in C5, with the same
	// ownership and "only active jobs change" rules.
	svcRows       map[int64]model.Service
	jobs          map[int64]*model.Job
	nextJob       int64
	finishApplied int // FinishJob calls that returned true
	versions      map[int64]model.ImageVersion
}

// seedVersion adds an image_versions row that GetImageVersion returns.
func (m *memStore) seedVersion(v model.ImageVersion) {
	m.mu.Lock()
	defer m.mu.Unlock()
	if m.versions == nil {
		m.versions = map[int64]model.ImageVersion{}
	}
	m.versions[v.ID] = v
}

func (m *memStore) GetImageVersion(_ context.Context, id int64) (model.ImageVersion, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	v, ok := m.versions[id]
	if !ok {
		return model.ImageVersion{}, model.ErrNotFound
	}
	return v, nil
}

func (m *memStore) ListServices(_ context.Context) ([]model.Service, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	var out []model.Service
	for _, svc := range m.svcRows {
		if svc.RemovedAt == nil {
			out = append(out, svc)
		}
	}
	sort.Slice(out, func(i, j int) bool { return out[i].ID < out[j].ID })
	return out, nil
}

func (m *memStore) SetPollError(_ context.Context, ids []int64, msg string) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	for _, id := range ids {
		svc := m.svcRows[id]
		svc.LastPollError = msg
		m.svcRows[id] = svc
	}
	return nil
}

func (m *memStore) SetLastTagDigest(_ context.Context, id int64, digest string) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	svc, ok := m.svcRows[id]
	if !ok {
		return model.ErrNotFound
	}
	svc.LastTagDigest = digest
	m.svcRows[id] = svc
	return nil
}

// seedService adds a service row that jobs can be triggered for.
func (m *memStore) seedService(svc model.Service) {
	m.mu.Lock()
	defer m.mu.Unlock()
	if m.svcRows == nil {
		m.svcRows = map[int64]model.Service{}
	}
	m.svcRows[svc.ID] = svc
}

func (m *memStore) GetService(_ context.Context, id int64) (model.Service, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	svc, ok := m.svcRows[id]
	if !ok {
		return model.Service{}, model.ErrNotFound
	}
	return svc, nil
}

// ServicesByRepo returns the seeded services whose image repo equals repo,
// ordered by ID.
func (m *memStore) ServicesByRepo(_ context.Context, repo string) ([]model.Service, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	var out []model.Service
	for _, svc := range m.svcRows {
		if svc.ImageRepo == repo && svc.RemovedAt == nil {
			out = append(out, svc)
		}
	}
	sort.Slice(out, func(i, j int) bool { return out[i].ID < out[j].ID })
	return out, nil
}

func (m *memStore) SetAvailable(_ context.Context, id int64, a model.Available) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	svc, ok := m.svcRows[id]
	if !ok {
		return model.ErrNotFound
	}
	svc.AvailableRef, svc.AvailableDigest, svc.AvailableSource, svc.AvailableAt = a.Ref, a.Digest, a.Source, &a.At
	m.svcRows[id] = svc
	return nil
}

func (m *memStore) CreateJob(_ context.Context, nj model.NewJob) (model.Job, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	if m.jobs == nil {
		m.jobs = map[int64]*model.Job{}
	}
	for _, j := range m.jobs {
		if j.ServiceID == nj.ServiceID && j.Active() {
			return *j, model.ErrJobActive
		}
	}
	m.nextJob++
	j := &model.Job{ID: m.nextJob, ServiceID: nj.ServiceID, AgentID: nj.AgentID, FromRef: nj.FromRef, ToRef: nj.ToRef,
		Trigger: nj.Trigger, RequestedBy: nj.RequestedBy, Status: model.JobQueued, CreatedAt: nj.CreatedAt}
	m.jobs[j.ID] = j
	return *j, nil
}

func (m *memStore) FailJob(_ context.Context, id int64, msg string, at time.Time) (bool, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	j, ok := m.jobs[id]
	if !ok || !j.Active() {
		return false, nil
	}
	j.Status, j.Error, j.FinishedAt = model.JobFailed, msg, &at
	return true, nil
}

func (m *memStore) GetJob(_ context.Context, id int64) (model.Job, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	j, ok := m.jobs[id]
	if !ok {
		return model.Job{}, model.ErrNotFound
	}
	return *j, nil
}

func (m *memStore) FailLostJobs(_ context.Context, agentID int64, known []int64, at time.Time) ([]int64, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	var lost []int64
	for _, j := range m.jobs {
		if j.AgentID == nil || *j.AgentID != agentID || !j.Active() {
			continue
		}
		isKnown := false
		for _, k := range known {
			isKnown = isKnown || k == j.ID
		}
		if !isKnown {
			j.Status, j.Error, j.FinishedAt = model.JobFailed, "lost: the agent did not know this job when it reconnected", &at
			lost = append(lost, j.ID)
		}
	}
	return lost, nil
}

func (m *memStore) JobProgress(_ context.Context, agentID, jobID int64, chunk string, at time.Time) (bool, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	j, ok := m.jobs[jobID]
	if !ok || j.AgentID == nil || *j.AgentID != agentID || !j.Active() {
		return false, nil
	}
	j.Status, j.Output = model.JobRunning, j.Output+chunk
	if j.StartedAt == nil {
		j.StartedAt = &at
	}
	return true, nil
}

func (m *memStore) FinishJob(_ context.Context, agentID, jobID int64, f model.JobFinish) (bool, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	j, ok := m.jobs[jobID]
	if !ok || j.AgentID == nil || *j.AgentID != agentID || !j.Active() {
		return false, nil
	}
	j.Status, j.Error, j.Output, j.OutputTruncated, j.FinishedAt = f.Status, f.Error, f.Output, f.OutputTruncated, &f.FinishedAt
	requestedRef := j.ToRef
	if f.Ref != "" {
		j.ToRef = f.Ref
	}
	if j.StartedAt == nil {
		j.StartedAt = &f.StartedAt
	}
	if f.Status == model.JobSucceeded {
		svc := m.svcRows[j.ServiceID]
		svc.DesiredRef = f.Ref
		svc.AwaitingInventory = true
		if _, d, _ := strings.Cut(f.Ref, "@"); !strings.Contains(requestedRef, "@") && imageref.IsDigest(d) {
			svc.LastTagDigest = d
		}
		// Same rule as the SQL: the job was created for the suggested ref after it was reported.
		requested, digest, _ := strings.Cut(requestedRef, "@")
		if svc.AvailableAt != nil && svc.AvailableRef == requested && !svc.AvailableAt.After(j.CreatedAt) &&
			(digest == "" || svc.AvailableDigest == digest) {
			svc.AvailableRef, svc.AvailableDigest, svc.AvailableSource, svc.AvailableAt = "", "", "", nil
		}
		m.svcRows[j.ServiceID] = svc
	}
	m.finishApplied++
	return true, nil
}

func (m *memStore) AgentByTokenHash(_ context.Context, hash string) (model.Agent, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.lookups++
	a, ok := m.agents[hash]
	if !ok {
		return model.Agent{}, model.ErrNotFound
	}
	return a, nil
}

func (m *memStore) RecordHello(context.Context, int64, protocol.Hello, time.Time) error { return nil }

func (m *memStore) ApplySnapshot(_ context.Context, _ int64, _ protocol.Hello, s []protocol.Service, _ time.Time) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.snapshots++
	m.services = map[string]protocol.Service{}
	for _, x := range s {
		m.services[x.Key()] = x
	}
	return nil
}

func (m *memStore) UpsertService(_ context.Context, _ int64, _ string, s protocol.Service, _ time.Time) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.services[s.Key()] = s
	return nil
}

func (m *memStore) RemoveService(_ context.Context, _ int64, _, key string, _ time.Time) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	delete(m.services, key)
	return nil
}

func (m *memStore) MarkSeen(context.Context, int64, time.Time) error { return nil }

func (m *memStore) keys() []string {
	m.mu.Lock()
	defer m.mu.Unlock()
	var out []string
	for k := range m.services {
		out = append(out, k)
	}
	return out
}

type staticSource struct {
	mu       sync.Mutex
	services []protocol.Service
}

func (s *staticSource) Hello(context.Context) (protocol.Hello, error) {
	return protocol.Hello{Hostname: "node-2"}, nil
}

func (s *staticSource) Snapshot(context.Context) ([]protocol.Service, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	return append([]protocol.Service(nil), s.services...), nil
}

func (s *staticSource) set(v ...protocol.Service) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.services = v
}

// testMaster is a running TLS test server with the hub mounted on /agent.
type testMaster struct {
	url         string
	fingerprint string
	hub         *currentHub   // the hub behind /agent; follows restart()
	attempts    *atomic.Int64 // TCP connections accepted by the server
	broker      *events.Broker
	restart     func() // replaces the hub and drops every connection
}

// currentHub forwards to whichever hub is live, so a master restart can swap
// the hub without the mux or the jobs service noticing.
type currentHub struct{ p atomic.Pointer[hub.Hub] }

func (c *currentHub) ServeHTTP(w http.ResponseWriter, r *http.Request) { c.p.Load().ServeHTTP(w, r) }
func (c *currentHub) Disconnect(id int64)                              { c.p.Load().Disconnect(id) }
func (c *currentHub) Route(svc model.Service, minProtocol int) (model.Route, error) {
	return c.p.Load().Route(svc, minProtocol)
}
func (c *currentHub) StartJob(r model.Route, id int64, s protocol.JobStart) error {
	return c.p.Load().StartJob(r, id, s)
}

func startMaster(t *testing.T, store *memStore) testMaster {
	t.Helper()
	key, err := pin.GenerateKey()
	require.NoError(t, err)
	cert, err := pin.Certificate(key)
	require.NoError(t, err)
	fingerprint, err := pin.Fingerprint(key.Public())
	require.NoError(t, err)

	quiet := slog.New(slog.NewTextHandler(io.Discard, nil))
	broker := events.NewBroker()
	h := &currentHub{}
	h.p.Store(hub.New(store, broker, quiet))
	mux := http.NewServeMux()
	mux.Handle("GET /agent", h)
	attempts := &atomic.Int64{}
	srv := httptest.NewUnstartedServer(mux)
	srv.Config.ErrorLog = log.New(io.Discard, "", 0)
	srv.Config.ConnState = func(_ net.Conn, s http.ConnState) {
		if s == http.StateNew {
			attempts.Add(1)
		}
	}
	srv.TLS = &tls.Config{Certificates: []tls.Certificate{cert}, MinVersion: tls.VersionTLS13}
	srv.StartTLS()
	t.Cleanup(srv.Close)
	restart := func() {
		// The hub holds no job state, so a new hub over the same store is a
		// faithful master restart.
		old := h.p.Swap(hub.New(store, broker, quiet))
		// CloseClientConnections skips hijacked (WebSocket) connections, so
		// cut each agent through the old hub.
		store.mu.Lock()
		defer store.mu.Unlock()
		for _, a := range store.agents {
			old.Disconnect(a.ID)
		}
	}
	return testMaster{url: "wss" + strings.TrimPrefix(srv.URL, "https") + "/agent", fingerprint: fingerprint,
		hub: h, attempts: attempts, broker: broker, restart: restart}
}

func runAgent(t *testing.T, url, fingerprint, token string, src conn.Source) *conn.Agent {
	t.Helper()
	a := conn.New(conn.Config{URL: url, Token: token, Fingerprint: fingerprint, Version: "e2e",
		MinBackoff: 10 * time.Millisecond, MaxBackoff: 50 * time.Millisecond,
	}, src, updater.NewRunner(nil, time.Minute), slog.New(slog.NewTextHandler(io.Discard, nil)))
	ctx, cancel := context.WithCancel(context.Background())
	done := make(chan struct{})
	go func() { _ = a.Run(ctx); close(done) }()
	t.Cleanup(func() { cancel(); <-done })
	return a
}

var (
	web   = protocol.Service{Kind: protocol.KindCompose, ComposeProject: "slusk", ComposeService: "web", Image: "slusk:1", Status: "running"}
	cache = protocol.Service{Kind: protocol.KindContainer, ContainerName: "cache", Image: "redis:alpine", Status: "running"}
)

func TestInventoryFlowsFromAgentToStore(t *testing.T) {
	store := &memStore{agents: map[string]model.Agent{agenttoken.Hash("pua_ok"): {ID: 2, Name: "node-2"}}}
	m := startMaster(t, store)
	src := &staticSource{}
	src.set(web)
	agent := runAgent(t, m.url, m.fingerprint, "pua_ok", src)

	require.Eventually(t, func() bool { return len(store.keys()) == 1 }, 2*time.Second, 10*time.Millisecond)

	src.set(cache)
	agent.Notify()
	require.Eventually(t, func() bool {
		k := store.keys()
		return len(k) == 1 && k[0] == "container|cache"
	}, 2*time.Second, 10*time.Millisecond)
}

func TestAgentWithWrongKeyNeverReachesHub(t *testing.T) {
	store := &memStore{agents: map[string]model.Agent{agenttoken.Hash("pua_ok"): {ID: 2, Name: "node-2"}}}
	m := startMaster(t, store)
	other, err := pin.GenerateKey()
	require.NoError(t, err)
	wrong, err := pin.Fingerprint(other.Public())
	require.NoError(t, err)

	runAgent(t, m.url, wrong, "pua_ok", &staticSource{})
	require.Eventually(t, func() bool { return m.attempts.Load() >= 1 }, 2*time.Second, 10*time.Millisecond, "the agent must actually try to connect")
	time.Sleep(200 * time.Millisecond)
	store.mu.Lock()
	defer store.mu.Unlock()
	assert.Zero(t, store.lookups, "TLS must fail before the token is ever presented")
}

func TestRevokedAgentIsRejected(t *testing.T) {
	store := &memStore{agents: map[string]model.Agent{
		agenttoken.Hash("pua_old"): {ID: 3, Name: "old", RevokedAt: func() *time.Time { n := time.Now(); return &n }()},
	}}
	m := startMaster(t, store)
	src := &staticSource{}
	src.set(web)
	runAgent(t, m.url, m.fingerprint, "pua_old", src)

	require.Eventually(t, func() bool {
		store.mu.Lock()
		defer store.mu.Unlock()
		return store.lookups >= 2
	}, 2*time.Second, 10*time.Millisecond, "the agent keeps retrying")
	assert.Empty(t, store.keys())
}

func TestAgentReconnectsAndResendsSnapshot(t *testing.T) {
	store := &memStore{agents: map[string]model.Agent{agenttoken.Hash("pua_ok"): {ID: 2, Name: "node-2"}}}
	m := startMaster(t, store)
	src := &staticSource{}
	src.set(web)
	runAgent(t, m.url, m.fingerprint, "pua_ok", src)

	snapshots := func() int {
		store.mu.Lock()
		defer store.mu.Unlock()
		return store.snapshots
	}
	require.Eventually(t, func() bool { return snapshots() >= 1 }, 2*time.Second, 10*time.Millisecond)

	m.hub.Disconnect(2)

	require.Eventually(t, func() bool { return snapshots() >= 2 }, 2*time.Second, 10*time.Millisecond, "agent reconnects and sends a fresh snapshot")
	assert.Len(t, store.keys(), 1)
}