fix(cluster): harden the worker tunnel, and stop starting a database per spec

Review follow-up. Twelve findings, none blocking, grouped here by what they
protect.

Panics. The handler now recovers between the WebSocket upgrade and the
hand-off, the way the peer link next door already did: net/http recovers the
panic but leaves the hijacked socket open, so without this a worker keeps a
session this replica has no entry for and will never detach. The claim gate in
Attach and reclaimOne is now released with defer, so a panic under Claim cannot
wedge one node's gate for the life of the process. SetTunnels gained the
nil-receiver guard its sibling Stop has.

Operability. A deployment with no registration token stores an empty token_hash
on every worker, so every tunnel dial 401s forever on a frontend that looks
correctly configured. That now warns at startup, logs its own line rather than
sharing the "wrong token" one, and is stated in the docs together with the fact
that setting the token later needs the workers to register again.

Authorization. A node still awaiting admin approval is refused with 403. The
rest of /api/node/ gates on nothing, but the two places that hand a node
something durable, its API key and its NATS credential, both refuse a pending
one, and a tunnel is that kind of grant. Draining and unhealthy nodes keep
their tunnels on purpose.

Comments that claimed more than the code. The global auth middleware does run
on this path and then declines to reject; the future per-node secret only lands
without a change here if it lands in TokenHash; the empty-hash guard is
defensive rather than deciding; ClusterPathPrefix is no longer only
replica-to-replica; the docs no longer say a reaped replica re-claims
unconditionally.

And the test harness. SetupTestDB started a PostgreSQL container per BeforeEach
with a readiness deadline it asserted on, which is one chance per spec to fail
one spec inside its setup, anywhere, never twice in the same place: the shape of
the flake seen twice here and never reproduced. It now starts one container per
process and creates a database per call, which is the pattern tests/e2e already
proved. Isolation is unchanged and is now asserted for the first time. All 69
call sites are untouched; the eleven consumer packages run 1404 specs green, and
jobs went from 34.3s to 3.3s, agents from 13.8s to 1.9s, cluster from 97.4s to
37.5s.

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
Assisted-by: Claude Opus 5 [claude-code]
This commit is contained in:
Ettore Di Giacinto committed 2026-09-01 10:36:10 +00:00
1 parent 6e55092a4b
commit 48ece89c63
10 files changed
+559 -68

No files matched your search

+13
View File
@@ -602,6 +602,19 @@ func API(application *application.Application) (*echo.Echo, error) {
var tunnels *clustersvc.TunnelRegistry
if d := application.Distributed(); d != nil {
tunnels = d.Tunnels
if distCfg.RegistrationToken == "" {
// A separate warning from the peer link's, for the same missing
// knob, because the broken thing is different and an operator has
// to be told both. A worker registering against a frontend with no
// registration token sends no token, so RegisterNodeEndpoint stores
// an empty token_hash, and the tunnel authenticates against exactly
// that column: every dial 401s, forever, on a deployment that looks
// correctly configured. Fixing it needs the token set AND the
// workers re-registered, which is why it is worth saying at boot
// rather than leaving in the 401s.
xlog.Warn("Worker tunnels will refuse every dial: no registration token is configured, so no worker has a stored token to authenticate against",
"route", clustersvc.ConnectPath, "knob", "LOCALAI_REGISTRATION_TOKEN")
}
}
routes.RegisterWorkerTunnelRoute(e, registry, tunnels)
+8 -3
View File
@@ -74,9 +74,14 @@ func isPublicRoute(method, path string) bool {
return false
}
// ClusterPathPrefix is the replica-to-replica namespace. Its handlers check the
// cluster token in the Authorization header themselves, so the check below lets
// them through the global session middleware.
// ClusterPathPrefix is the machine-to-machine cluster namespace. It carries two
// different trust relationships, on two different credentials: the
// replica-to-replica peer link, which checks the shared cluster token, and the
// worker-to-frontend tunnel, which checks the dialing node's own stored token
// hash. What they have in common is the only thing this prefix asserts, that
// each handler checks its own Authorization header, so the check below lets them
// through the global session middleware rather than rejecting a caller that has
// no session and no user.
//
// The cluster routes do NOT derive their paths from this constant: they are
// registered from core/services/cluster's own literal, because that package
+78 -9
View File
@@ -37,10 +37,16 @@ import (
// panic; see the 503 below.
//
// It is deliberately absent from auth.RouteFeatureRegistry. That registry gates
// a route on the FEATURES OF AN AUTHENTICATED USER, and there is no user here:
// the dialer is a worker process holding a machine credential, and the global
// auth middleware never runs on this path because it sits under
// auth.ClusterPathPrefix.
// a route on the FEATURES OF AN AUTHENTICATED USER, resolved from auth.GetUser,
// and there is no user here: the dialer is a worker process holding a machine
// credential.
//
// The global auth middleware does RUN on this path; what it does not do is
// reject. It attempts session, bearer and legacy-key authentication first, so it
// may even have set auth_user from a worker token that happens to match an API
// key, and then core/http/auth/middleware.go:90 lets the request through
// because usesAlternativeAuthentication reports the path as one whose
// credentials its own route checks. Nothing here reads what it set.
func ConnectHandler(registry *nodes.NodeRegistry, tunnels *clustersvc.TunnelRegistry) echo.HandlerFunc {
// gorilla's default CheckOrigin restricts a browser to same-origin and lets
// a header-less client (which every worker is) through, so the zero value
@@ -62,6 +68,14 @@ func ConnectHandler(registry *nodes.NodeRegistry, tunnels *clustersvc.TunnelRegi
// answering "unauthorized" would send the operator hunting a token
// problem that does not exist. It is checked after the header so that
// an anonymous dial still gets the 401 the coverage test requires.
//
// Only the registry half is covered by a spec. The two are read from one
// application.Distributed() in core/http/app.go and initDistributed
// returns an error rather than a partial struct, so a non-nil registry
// beside a nil tunnel registry is unreachable and no spec constructs it;
// the second half is defence against a future wiring that splits them,
// where the cost would be a nil dereference in Attach after the
// connection is already hijacked.
if registry == nil || tunnels == nil {
return echo.NewHTTPError(http.StatusServiceUnavailable, "distributed mode not enabled")
}
@@ -89,11 +103,42 @@ func ConnectHandler(registry *nodes.NodeRegistry, tunnels *clustersvc.TunnelRegi
return echo.NewHTTPError(http.StatusInternalServerError, "node lookup failed")
}
// Split from the mismatch below because they are different operator
// problems with different fixes. An empty stored hash means the node
// registered against a frontend with no LOCALAI_REGISTRATION_TOKEN, so
// no worker on that deployment can ever tunnel until one is set and the
// workers re-register; a mismatch means one worker holds the wrong
// secret. One log line for both leaves an operator reading "wrong token"
// while every worker fails identically.
if node.TokenHash == "" {
xlog.Warn("Refusing a worker tunnel: this node has no stored token, which means it registered with no registration token configured",
"node", nodeID, "knob", "LOCALAI_REGISTRATION_TOKEN")
return echo.NewHTTPError(http.StatusUnauthorized, "unauthorized")
}
if !authorizedWorker(token, node.TokenHash) {
xlog.Debug("worker tunnel dial presented the wrong token", "node", nodeID)
return echo.NewHTTPError(http.StatusUnauthorized, "unauthorized")
}
// Authenticated but not authorised, so 403 rather than 401: the fix is an
// admin approving the node, not a different credential, and answering
// 401 would send an operator looking at tokens.
//
// Only StatusPending is refused. The rest of /api/node/ self-service
// gates on nothing at all, but the two places that hand a node something
// DURABLE both refuse a pending one: the agent worker's API key
// (core/http/endpoints/localai/nodes.go:224) and its NATS credential
// (nodes.go:293). A tunnel is that kind of grant, not a heartbeat: it is
// a standing pipe into the worker recorded in node_connections and
// relayed to by every other replica. Draining and unhealthy nodes keep
// their tunnels on purpose; draining means finish what you have, and a
// node marked unhealthy for missed heartbeats needs the pipe to recover
// through.
if node.Status == nodes.StatusPending {
xlog.Warn("Refusing a worker tunnel: this node is awaiting admin approval", "node", nodeID)
return echo.NewHTTPError(http.StatusForbidden, "node is pending approval")
}
ws, err := upgrader.Upgrade(c.Response(), c.Request(), nil)
if err != nil {
// Upgrade has already written its own failure to the client.
@@ -108,6 +153,23 @@ func ConnectHandler(registry *nodes.NodeRegistry, tunnels *clustersvc.TunnelRegi
return nil
}
// The same guard PeerHandler carries, for the same reason and one more.
// net/http recovers a panic from this goroutine but does not close the
// hijacked connection, and middleware.Recover does not either, so a
// panic below would leave the worker holding a live session this replica
// has no entry for and will never detach. The extra reason here is that
// Attach does database work: a panic inside it, with the session left
// open, is a tunnel nothing can reach and nothing will clean up.
//
// It re-panics rather than swallowing. Whatever it caught is a bug, and
// the recovery middleware above is what should report it.
defer func() {
if r := recover(); r != nil {
_ = sess.Close()
panic(r)
}
}()
// From here the connection is hijacked, so no status can reach the
// worker any more: a failure is a closed socket, which is what its
// reconnect loop reads.
@@ -176,12 +238,19 @@ func bearerToken(r *http.Request) (string, bool) {
// attacker knows. What this does rule out is the shortcut of comparing against
// the configured token itself, which is what would have to be unpicked later;
// the day a worker is issued its own secret at registration, this check starts
// isolating workers from each other with no change here.
// isolating workers from each other with no change here, PROVIDED the secret
// lands in TokenHash. The one per-node secret LocalAI mints today, the agent
// worker's api_token, does not: provisionAgentWorkerKey writes an auth.User and
// an auth.APIKey referenced by node.AuthUserID / node.APIKeyID and never touches
// this column. If the next task follows that precedent instead, this comparison
// is what has to change.
//
// An empty stored hash authorizes nobody. Such a row exists whenever a worker
// registered against a frontend with no registration token configured, and the
// alternative, letting any token through for those, would make the check depend
// on a setting the dialer can neither see nor be blamed for.
// The empty-hash guard is defensive rather than deciding: a stored hash is
// hex-encoded SHA-256, so 64 bytes or nothing, and ConstantTimeCompare already
// returns 0 on a length mismatch (crypto/internal/fips140/subtle/constant_time.go:17-20
// returns 0 outright when the lengths differ). It is kept because a reader should not have to
// derive "an unregistered node authorizes nobody" from a length rule, and
// because the caller logs that case separately.
func authorizedWorker(token, storedHash string) bool {
if storedHash == "" {
return false
@@ -194,6 +194,45 @@ var _ = Describe("Worker tunnel handler", func() {
"the claim outlived the socket, so this replica keeps being named the owner of a worker it no longer holds")
})
It("refuses a node whose row carries no stored token", func() {
// A deployment with no registration token configured produces exactly
// these rows: the worker sends no token, so registration stores none,
// and there is nothing here to authenticate against.
Expect(db.Exec(`UPDATE backend_nodes SET token_hash = '' WHERE id = ?`, nodeID).Error).To(Succeed())
_, resp, err := websocket.DefaultDialer.Dial(wsConnectURL(srv, nodeID), bearer(workerToken))
Expect(err).To(HaveOccurred())
Expect(resp).ToNot(BeNil())
Expect(resp.StatusCode).To(Equal(http.StatusUnauthorized))
Expect(tun.Held()).To(BeEmpty())
})
It("refuses a node that is still awaiting admin approval", func() {
// Approval is what gates a node's participation, and a tunnel is a
// standing pipe recorded in node_connections, not a heartbeat. 403 and
// not 401: the credential is right, the authorisation is missing, and
// the fix is an admin rather than a different token.
Expect(db.Exec(`UPDATE backend_nodes SET status = ? WHERE id = ?`, nodes.StatusPending, nodeID).Error).To(Succeed())
_, resp, err := websocket.DefaultDialer.Dial(wsConnectURL(srv, nodeID), bearer(workerToken))
Expect(err).To(HaveOccurred())
Expect(resp).ToNot(BeNil())
Expect(resp.StatusCode).To(Equal(http.StatusForbidden))
Expect(tun.Held()).To(BeEmpty())
})
It("still admits a draining node, which has work to finish", func() {
// Only pending is refused. Draining means "start nothing new", not
// "lose the pipe your in-flight requests travel on", and a node marked
// unhealthy for missed heartbeats needs the tunnel to recover through.
Expect(db.Exec(`UPDATE backend_nodes SET status = ? WHERE id = ?`, nodes.StatusDraining, nodeID).Error).To(Succeed())
conn, _, err := websocket.DefaultDialer.Dial(wsConnectURL(srv, nodeID), bearer(workerToken))
Expect(err).ToNot(HaveOccurred())
DeferCleanup(func() { _ = conn.Close() })
Eventually(tun.Held, "10s").Should(ConsistOf(nodeID))
})
It("reports a lookup failure as a failure, not as a refusal", func() {
// ErrNotOwner, 401 and 404 are all ANSWERS. A database that cannot be
// read is none of them: telling a worker its credentials are wrong when
@@ -208,6 +247,56 @@ var _ = Describe("Worker tunnel handler", func() {
})
})
var _ = Describe("Worker tunnel handler when the attach panics", func() {
// net/http recovers a panic from the request goroutine but does NOT close a
// hijacked connection, so without the handler's own recover the worker keeps
// a live session this replica has no entry for and will never detach: its
// opens fill yamux's 256-deep backlog and then hang with no error. The panic
// is injected through a real path rather than a fake one, a registry built
// over no database at all, which is what Attach's first database call
// dereferences.
var (
srv *httptest.Server
db *gorm.DB
nodeID string
)
BeforeEach(func() {
ctx := context.Background()
db = testutil.SetupTestDB()
nodeReg, err := nodes.NewNodeRegistry(db)
Expect(err).ToNot(HaveOccurred())
node := &nodes.BackendNode{Name: "worker-1", Address: "10.0.0.9:50051", TokenHash: tokenHash(workerToken)}
Expect(nodeReg.Register(ctx, node, true)).To(Succeed())
nodeID = node.ID
e := echo.New()
routes.RegisterWorkerTunnelRoute(e, nodeReg, clustersvc.NewTunnelRegistry(nil, "me"))
srv = httptest.NewServer(e)
DeferCleanup(func() {
// A hijacked connection the handler never closed would park Close
// forever, turning the assertion below into a suite hang.
srv.CloseClientConnections()
srv.Close()
})
})
It("closes the worker's session instead of stranding it", func() {
conn, _, err := websocket.DefaultDialer.Dial(wsConnectURL(srv, nodeID), bearer(workerToken))
Expect(err).ToNot(HaveOccurred())
DeferCleanup(func() { _ = conn.Close() })
workerSess, err := yamux.Client(clustersvc.WebsocketConn(conn), nil, nil)
Expect(err).ToNot(HaveOccurred())
DeferCleanup(func() { _ = workerSess.Close() })
// Asserting on OpenStream would hang rather than fail: yamux only
// acknowledges a stream once the peer accepts it, and the leak this
// pins is precisely that nobody ever will.
Eventually(workerSess.IsClosed, "10s").Should(BeTrue())
})
})
var _ = Describe("Worker tunnel handler without distributed mode", func() {
// The route is registered in every deployment so that the route-coverage
// test sees it, which is what pins the reject-before-upgrade rule. With no
+8
View File
@@ -81,7 +81,15 @@ func NewMembership(reg *Registry, id, addr, version string) *Membership {
// It is a setter rather than a constructor argument because the tunnel registry
// is what the tunnel endpoint is built on, and that is wired after membership
// is already running.
//
// Safe on a nil receiver, like Stop. This package deliberately produces a nil
// *Membership (core/application/distributed.go leaves it nil when no
// peer-reachable address can be derived), so a setter that panicked on one
// would be a trap for the next caller rather than an impossibility.
func (m *Membership) SetTunnels(t *TunnelRegistry) {
if m == nil {
return
}
m.mu.Lock()
defer m.mu.Unlock()
m.tunnels = t
+62 -37
View File
@@ -182,18 +182,33 @@ func (t *TunnelRegistry) Attach(ctx context.Context, nodeID string, sess *yamux.
return 0, fmt.Errorf("attaching tunnel for node %q: %w", nodeID, err)
}
epoch, err := t.reg.Claim(ctx, nodeID, t.selfID)
// The gated part is a closure so its release can be DEFERRED while the
// session close below still happens outside the gate. Releasing on each
// return path instead leaves one way out uncovered: a panic. Claim does
// database work, and a panic anywhere under it would leave this node's gate
// closed for the life of the process, so every later Attach or Reclaim for
// that worker would block in enterClaim until its own context expired. The
// caller's recover would report the panic and the worker would look
// permanently unable to reconnect, with nothing linking the two.
var previous *heldTunnel
epoch, err := func() (int64, error) {
defer t.leaveClaim(nodeID)
epoch, err := t.reg.Claim(ctx, nodeID, t.selfID)
if err != nil {
return 0, err
}
t.mu.Lock()
previous = t.tunnels[nodeID]
t.tunnels[nodeID] = &heldTunnel{sess: sess, token: epoch, claim: epoch}
t.mu.Unlock()
return epoch, nil
}()
if err != nil {
t.leaveClaim(nodeID)
return 0, err
}
t.mu.Lock()
previous := t.tunnels[nodeID]
t.tunnels[nodeID] = &heldTunnel{sess: sess, token: epoch, claim: epoch}
t.mu.Unlock()
t.leaveClaim(nodeID)
// Closed after the gate is released, not under it. The gate is justified by
// being held for one claim round trip, and closing a session is not that:
// yamux closes the underlying conn and then waits for both its send and
@@ -381,39 +396,49 @@ func (t *TunnelRegistry) reclaimOne(ctx context.Context, nodeID string) error {
return fmt.Errorf("re-claiming node %q: %w", nodeID, err)
}
t.mu.Lock()
tunnel, ok := t.tunnels[nodeID]
t.mu.Unlock()
if !ok {
// Detached between Reclaim listing the nodes and this gate. Nothing was
// claimed, so there is nothing to undo.
t.leaveClaim(nodeID)
return errTunnelNotReclaimed
}
if tunnel.sess.IsClosed() {
xlog.Debug("skipping re-claim of a worker tunnel whose session is closed", "node", nodeID)
t.leaveClaim(nodeID)
return errTunnelNotReclaimed
}
// Gated part in a closure so the release is DEFERRED, for the reason Attach
// gives: a panic under Claim would otherwise wedge this node's gate for the
// life of the process. The trailing Release still runs outside the gate.
var epoch int64
var installed bool
if err := func() error {
defer t.leaveClaim(nodeID)
epoch, err := t.reg.Claim(ctx, nodeID, t.selfID)
if err != nil {
t.leaveClaim(nodeID)
t.mu.Lock()
tunnel, ok := t.tunnels[nodeID]
t.mu.Unlock()
if !ok {
// Detached between Reclaim listing the nodes and this gate. Nothing
// was claimed, so there is nothing to undo.
return errTunnelNotReclaimed
}
if tunnel.sess.IsClosed() {
xlog.Debug("skipping re-claim of a worker tunnel whose session is closed", "node", nodeID)
return errTunnelNotReclaimed
}
var err error
epoch, err = t.reg.Claim(ctx, nodeID, t.selfID)
if err != nil {
return err
}
t.mu.Lock()
current, present := t.tunnels[nodeID]
installed = present
if installed {
// current is necessarily the entry read above: the gate is still
// held, and Attach and reclaimOne are the only writers that install
// one. The identity is therefore not re-checked; the case that IS
// reachable is the entry being gone, because Detach is not gated.
current.claim = epoch
}
t.mu.Unlock()
return nil
}(); err != nil {
return err
}
t.mu.Lock()
current, installed := t.tunnels[nodeID]
if installed {
// current is necessarily the entry read above: the gate is still held,
// and Attach and reclaimOne are the only writers that install one. The
// identity is therefore not re-checked; the case that IS reachable is
// the entry being gone, because Detach is not gated.
current.claim = epoch
}
t.mu.Unlock()
t.leaveClaim(nodeID)
if installed {
return nil
}
+75
View File
@@ -716,3 +716,78 @@ var _ = Describe("Re-claiming tunnels after this replica's rows were reaped", fu
Consistently(stored, 2*cluster.InstanceHeartbeat, time.Second).Should(Equal(epoch))
})
})
var _ = Describe("The worker tunnel registry's claim gate", func() {
// These specs need no database. They pin what happens when the database
// call under the gate does not return normally, which is the one exit a
// release-on-every-return-path cannot cover.
//
// The panic is produced by the production code itself: a registry with no
// *Registry behind it dereferences nothing at the gate and then panics
// inside Claim, which is exactly where a real one does its work.
var tun *cluster.TunnelRegistry
BeforeEach(func() {
tun = cluster.NewTunnelRegistry(nil, "me")
})
// attachPanics runs one Attach that is expected to panic, swallowing the
// panic so the spec can go on to ask what state it left behind.
attachPanics := func(ctx context.Context, nodeID string) {
defer GinkgoRecover()
defer func() { _ = recover() }()
frontend, _ := workerTunnel()
_, _ = tun.Attach(ctx, nodeID, frontend)
Fail("Attach was expected to panic inside Claim, so this spec is no longer testing what it claims")
}
It("frees the node's gate when the claim panics", func() {
attachPanics(context.Background(), "w1")
// A wedged gate is indistinguishable from a slow one except by waiting,
// so the second attempt is given a deadline. Reaching Claim means
// panicking again; returning a context error means it never got past
// enterClaim and this worker could never reconnect to this replica.
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
reached := make(chan any, 1)
go func() {
defer GinkgoRecover()
defer func() { reached <- recover() }()
frontend, _ := workerTunnel()
_, err := tun.Attach(ctx, "w1", frontend)
Expect(err).To(MatchError(context.DeadlineExceeded),
"the gate for this node was never released, so every later dial from it blocks until its own context expires")
}()
Eventually(reached, "5s").Should(Receive(Not(BeNil())),
"the second Attach did not reach Claim, so the panicking one left the gate closed")
})
It("leaves another node's gate alone", func() {
// The gate is per node so that one wedged worker cannot stop the rest;
// this holds that property against the panic path too.
attachPanics(context.Background(), "w1")
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
reached := make(chan any, 1)
go func() {
defer GinkgoRecover()
defer func() { reached <- recover() }()
frontend, _ := workerTunnel()
_, _ = tun.Attach(ctx, "w2", frontend)
}()
Eventually(reached, "5s").Should(Receive(Not(BeNil())))
})
})
var _ = Describe("Membership.SetTunnels", func() {
It("is safe on a nil receiver, like Stop", func() {
// A nil *Membership is a value this codebase deliberately produces when
// no peer-reachable address can be derived, so the asymmetry with Stop
// would be a trap for the next caller.
var m *cluster.Membership
Expect(func() { m.SetTunnels(cluster.NewTunnelRegistry(nil, "me")) }).ToNot(Panic())
})
})
+168 -15
View File
@@ -2,7 +2,11 @@ package testutil
import (
"context"
"fmt"
"net/url"
"runtime"
"sync"
"sync/atomic"
"time"
"github.com/testcontainers/testcontainers-go"
@@ -16,27 +20,176 @@ import (
. "github.com/onsi/gomega"
)
// SetupTestDB creates a fresh PostgreSQL 16 container and returns a gorm.DB.
// The container is cleaned up via DeferCleanup when the test completes.
// One PostgreSQL container per test PROCESS, not per spec, with a database per
// SetupTestDB call.
//
// Starting a container per spec was both slow and flaky. Slow because a
// postgres:16 start is seconds and the packages behind this helper hold several
// hundred specs; flaky because every start was a fresh chance to miss the
// readiness deadline, and a miss lands in the caller's BeforeEach as a failure
// of whichever spec happened to be running. That is the exact shape of the
// intermittent single-spec failure seen twice in this package and never
// reproduced: one spec of many, no pattern, never twice in the same place.
// Starting the container once per process leaves one chance to miss it instead
// of one per spec, and moves that chance onto a deadline that only has to be met
// while nothing else is competing for the machine.
//
// Isolation is unchanged and is what callers actually depend on: each call still
// hands back an empty database that no other spec can see. The database is
// dropped when the spec that asked for it ends. Advisory locks, sequences and
// extensions are all per-database in PostgreSQL, so nothing the packages behind
// this helper rely on leaks between specs.
//
// This mirrors the pattern already proven in tests/e2e/distributed
// (testhelpers_test.go), which is where the argument and the measurements come
// from.
//
// One container per process rather than one shared across `ginkgo -p` workers is
// deliberate: parallel Ginkgo processes are separate OS processes, each gets its
// own container, and nothing has to coordinate database names across them.
var (
sharedOnce sync.Once
sharedPG *tcpostgres.PostgresContainer
sharedDSN string
sharedErr error
// dbCounter makes each database name unique within this process. The
// container is not shared across processes, so a process-local counter is
// enough.
dbCounter atomic.Int64
)
// The container outlives every spec, so its teardown belongs to the suite. This
// registers one AfterSuite in every suite that imports this package, which is
// every suite that could have started a container; it is a no-op in the ones
// that never call SetupTestDB.
//
// Package-level rather than something callers have to remember: a helper whose
// cleanup depends on 56 test files each declaring a hook is a helper that leaks
// containers the first time someone forgets. Registration happens during package
// initialisation, which is before RunSpecs, so Ginkgo is still building its tree.
var _ = AfterSuite(func() {
if sharedPG == nil {
return
}
// Best-effort: a failed terminate must not fail a suite whose specs all
// passed. Testcontainers' reaper removes it in that case.
_ = sharedPG.Terminate(context.Background())
})
// sharedPostgres returns the DSN of this process's PostgreSQL container,
// starting it on first use.
//
// The error is remembered rather than only asserted inside the sync.Once: an
// assertion there fails the one spec that happened to be first, and every later
// spec would then find a nil container and fail for some unrelated-looking
// reason. Re-asserting the stored error makes every affected spec say the same
// true thing.
func sharedPostgres() string {
GinkgoHelper()
sharedOnce.Do(func() {
ctx := context.Background()
sharedPG, sharedErr = tcpostgres.Run(ctx, "postgres:16",
tcpostgres.WithDatabase("testdb"),
tcpostgres.WithUsername("test"),
tcpostgres.WithPassword("test"),
// The deadline is per process now, not per spec, so it is generous
// on purpose: it is paid once, and the cost of missing it is a
// whole suite rather than one spec.
testcontainers.WithWaitStrategyAndDeadline(120*time.Second,
wait.ForLog("database system is ready to accept connections").WithOccurrence(2)),
)
if sharedErr != nil {
return
}
sharedDSN, sharedErr = sharedPG.ConnectionString(ctx, "sslmode=disable")
})
Expect(sharedErr).ToNot(HaveOccurred(), "the suite's PostgreSQL container could not be started")
return sharedDSN
}
// SetupTestDB returns a gorm.DB on a PostgreSQL database created for the calling
// spec. The database is dropped, and its connection pool closed, when the spec
// ends.
func SetupTestDB() *gorm.DB {
GinkgoHelper()
if runtime.GOOS == "darwin" {
Skip("testcontainers requires Docker, not available on macOS CI")
}
ctx := context.Background()
pgC, err := tcpostgres.Run(ctx, "postgres:16",
tcpostgres.WithDatabase("testdb"),
tcpostgres.WithUsername("test"),
tcpostgres.WithPassword("test"),
testcontainers.WithWaitStrategyAndDeadline(60*time.Second,
wait.ForLog("database system is ready to accept connections").WithOccurrence(2)),
)
Expect(err).ToNot(HaveOccurred())
DeferCleanup(func() { pgC.Terminate(context.Background()) })
connStr, err := pgC.ConnectionString(ctx, "sslmode=disable")
Expect(err).ToNot(HaveOccurred())
db, err := gorm.Open(postgres.Open(connStr), &gorm.Config{
dsn := sharedPostgres()
name := fmt.Sprintf("testdb_%d", dbCounter.Add(1))
// Scoped so a failed CREATE cannot leak the pool: the assertion panics out
// of this function, and a leaked pool per failing spec exhausts the
// server's connection limit for every spec after it.
//
// CREATE and DROP DATABASE cannot run against the target database itself,
// so both go through a short-lived connection to the container's own
// maintenance database.
func() {
admin := openPool(dsn)
defer closePool(admin)
Expect(admin.Exec(fmt.Sprintf("CREATE DATABASE %q", name)).Error).To(Succeed())
}()
db, err := gorm.Open(postgres.Open(replaceDBName(dsn, name)), &gorm.Config{
Logger: logger.Default.LogMode(logger.Silent),
})
Expect(err).ToNot(HaveOccurred())
DeferCleanup(func() {
// The caller's own DeferCleanups were registered later and so run
// first, which is what lets a spec keep using this handle in its
// teardown.
closePool(db)
drop, err := gorm.Open(postgres.Open(dsn), &gorm.Config{Logger: logger.Discard})
if err != nil {
// Reported, never asserted. A cleanup that fails the spec turns one
// database hiccup into a failure that buries whatever the spec was
// actually about.
AddReportEntry("drop test database skipped", fmt.Sprintf("%s: %v", name, err))
return
}
defer closePool(drop)
// FORCE terminates whatever connections the spec left open, including
// any a background goroutine is still holding (PostgreSQL 13+).
if err := drop.Exec(fmt.Sprintf("DROP DATABASE IF EXISTS %q WITH (FORCE)", name)).Error; err != nil {
AddReportEntry("drop test database failed", fmt.Sprintf("%s: %v", name, err))
}
})
return db
}
// openPool connects to dsn with logging off. Used for the short-lived
// maintenance connections only; the database a spec is handed keeps gorm's
// silent logger so a caller can still swap it.
func openPool(dsn string) *gorm.DB {
GinkgoHelper()
db, err := gorm.Open(postgres.Open(dsn), &gorm.Config{Logger: logger.Discard})
Expect(err).ToNot(HaveOccurred())
return db
}
func closePool(db *gorm.DB) {
if db == nil {
return
}
if sqlDB, err := db.DB(); err == nil {
_ = sqlDB.Close()
}
}
// replaceDBName swaps the database component of a DSN, preserving credentials,
// host, port and query parameters.
func replaceDBName(dsn, name string) string {
GinkgoHelper()
u, err := url.Parse(dsn)
Expect(err).ToNot(HaveOccurred())
u.Path = "/" + name
return u.String()
}
+52
View File
@@ -0,0 +1,52 @@
package testutil_test
import (
"testing"
"github.com/mudler/LocalAI/core/services/testutil"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
)
func TestTestutil(t *testing.T) {
RegisterFailHandler(Fail)
RunSpecs(t, "Test Utilities Suite")
}
// The container is shared per process now, so the isolation callers depend on
// comes from a database per call rather than from a server per call. That is
// the property 69 call sites across eleven packages assume without saying so,
// and nothing else in the tree asserts it.
var _ = Describe("SetupTestDB", func() {
type row struct {
ID int
}
It("hands back a database no other caller can see into", func() {
first := testutil.SetupTestDB()
second := testutil.SetupTestDB()
Expect(first.Exec(`CREATE TABLE isolation_probe (id int)`).Error).To(Succeed())
Expect(first.Exec(`INSERT INTO isolation_probe VALUES (1)`).Error).To(Succeed())
var found []row
err := second.Raw(`SELECT id FROM isolation_probe`).Scan(&found).Error
Expect(err).To(HaveOccurred(),
"two SetupTestDB calls landed on the same database, so every spec can now see every other spec's rows")
// The second database must also be usable, not merely different: an
// isolation check that passed because the second handle was broken
// would prove nothing.
Expect(second.Exec(`CREATE TABLE isolation_probe (id int)`).Error).To(Succeed())
})
It("hands back an empty database", func() {
db := testutil.SetupTestDB()
var tables int64
Expect(db.Raw(
`SELECT count(*) FROM information_schema.tables WHERE table_schema = 'public'`,
).Scan(&tables).Error).To(Succeed())
Expect(tables).To(BeZero(), "a spec was handed a database another spec had already migrated")
})
})
+6 -4
View File
@@ -104,15 +104,17 @@ The peer link is served at `/api/cluster/peer` and authenticates with `LOCALAI_R
A worker can open one long-lived, multiplexed tunnel to the frontend instead of listening on a port of its own. It dials `GET /api/cluster/connect?id=<node id>`, the connection is upgraded to a WebSocket, and every subsequent request the frontend makes to that worker travels as a stream inside it. Nothing dials *into* the worker, so a worker behind NAT, in another Kubernetes cluster or on a laptop needs no inbound port and no reachable address.
The dial is authenticated against **that node's own stored token**, not the deployment-wide registration token: the frontend hashes the presented bearer token and compares it with the hash recorded on the node's row at registration. A worker that presents a token belonging to no node, or names a node ID the frontend has never seen, is refused with `401` before the WebSocket upgrade happens. A frontend that cannot read its node table answers `500` rather than `401`, so a worker retries instead of re-registering under a new identity.
The dial is authenticated against **the hash stored on that node's row, which today is still the registration token's hash**: the frontend hashes the presented bearer token and compares it with what registration recorded. A worker that presents a token belonging to no node, or names a node ID the frontend has never seen, is refused with `401` before the WebSocket upgrade happens. A node still awaiting admin approval is refused with `403`. A frontend that cannot read its node table answers `500` rather than `401`, so a worker retries instead of re-registering under a new identity.
One honest caveat about that today: a worker currently registers by presenting the deployment's registration token, so the hash stored on its row *is* that token's hash, and a leaked registration token plus a known node ID still gets a tunnel. What the tunnel endpoint does not do is trust the configured token directly, so when workers are issued their own secrets at registration the isolation becomes real with no change to this route.
So the isolation is not there yet, and the caveat is worth stating plainly: because a worker registers by presenting the deployment's registration token, a leaked registration token plus a known node ID still gets a tunnel. What the tunnel endpoint does not do is trust the configured token directly, so when workers are issued their own per-node secrets and those secrets are what registration stores, the isolation becomes real with no change to this route.
The tunnel lands on exactly one frontend replica, and that replica records itself as the owner of the worker's connection in the `node_connections` table. When the socket dies the claim is dropped with it. If the replica stalls long enough for its peers to reap it, it re-claims every tunnel it still holds as soon as it re-registers.
**Worker tunnels need `LOCALAI_REGISTRATION_TOKEN` set.** A worker registering against a frontend with no registration token configured sends no token, so nothing is stored on its row, and the tunnel has nothing to authenticate it against: every dial is refused with `401`. LocalAI warns about this at startup. Setting the token later is not enough on its own; the workers have to register again for the hash to be written.
The tunnel lands on exactly one frontend replica, and that replica records itself as the owner of the worker's connection in the `node_connections` table. When the socket dies the claim is dropped with it. If the replica stalls long enough for its peers to reap it, it re-claims the tunnels it still holds on a live session as soon as it re-registers, skipping any whose socket has already gone. That re-claim needs the replica to have an advertised address: without one it never had an instance row to begin with, and its tunnels are usable only by the replica holding them.
| Method | Path | Description |
|--------|------|-------------|
| `GET` | `/api/cluster/connect?id=<node id>` | Worker opens its multiplexed tunnel (`Authorization: Bearer <the worker's own token>`) |
| `GET` | `/api/cluster/connect?id=<node id>` | Worker opens its multiplexed tunnel (`Authorization: Bearer <the token the worker registered with>`) |
The route is exempt from the normal session/API-key authentication (it authenticates itself, like `/api/cluster/peer`) and is registered in every deployment. Outside distributed mode there is no node table to check a token against, so it answers `503`.