test: agent to master inventory over pinned TLS end to end
This commit is contained in:
@@ -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