416 lines
12 KiB
Go
416 lines
12 KiB
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/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
|
|
}
|
|
|
|
// 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.AvailableSource, svc.AvailableAt = a.Ref, 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
|
|
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
|
|
// Same rule as the SQL: a job created after the suggestion deployed it.
|
|
if svc.AvailableAt != nil && !svc.AvailableAt.After(j.CreatedAt) {
|
|
svc.AvailableRef, 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) (model.Route, error) { return c.p.Load().Route(svc) }
|
|
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 := ¤tHub{}
|
|
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)
|
|
}
|