diff --git a/internal/e2e/agent_master_test.go b/internal/e2e/agent_master_test.go new file mode 100644 index 0000000..36a071b --- /dev/null +++ b/internal/e2e/agent_master_test.go @@ -0,0 +1,230 @@ +// 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" + "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/agenttoken" + "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 +} + +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 *hub.Hub + attempts *atomic.Int64 // TCP connections accepted by the server +} + +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) + + h := hub.New(store, events.NewBroker(), slog.New(slog.NewTextHandler(io.Discard, nil))) + 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) + return testMaster{url: "wss" + strings.TrimPrefix(srv.URL, "https") + "/agent", fingerprint: fingerprint, hub: h, attempts: attempts} +} + +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, 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) +}