feat: master/agent phase 1 — live inventory (#6) #19
@@ -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)
|
||||
}
|
||||
Reference in New Issue
Block a user