diff --git a/.golangci.yml b/.golangci.yml index d25d1ccb4789..7a68e69a75f2 100644 --- a/.golangci.yml +++ b/.golangci.yml @@ -19,7 +19,22 @@ linters: - staticcheck enable: - forbidigo + # gocritic is enabled for ONE checker: ruleguard, which runs the rules in + # hack/lint/. Every other gocritic check is off (disable-all below), so + # this adds no style noise; it is here purely as the gate that catches a + # gRPC backend wrapper written without Unwrap. See + # hack/lint/backend_wrappers.go for why that cannot be a compile-time + # assertion. + - gocritic settings: + gocritic: + disable-all: true + enabled-checks: + - ruleguard + settings: + ruleguard: + failOn: all + rules: '${base-path}/hack/lint/backend_wrappers.go' forbidigo: forbid: - pattern: '^t\.Errorf$' @@ -126,3 +141,10 @@ linters: - path: ^backend/go/whisper/sources/ text: 'http\.(DefaultClient|Get|Post|PostForm|Head)' linters: [forbidigo] + # Test doubles embed grpc.Backend to inherit the interface's method set + # over a NIL value; they decorate nothing, hold no inner client, and have + # no transport answer to forward. The rule targets production wrappers, + # which is where swallowing that answer deletes replica rows. + # gocritic here is only the backend-wrapper ruleguard rule. + - path: _test\.go$ + linters: [gocritic] diff --git a/core/application/distributed.go b/core/application/distributed.go index b7dc0bf91351..9835b5ced9df 100644 --- a/core/application/distributed.go +++ b/core/application/distributed.go @@ -5,6 +5,8 @@ import ( "encoding/json" "fmt" "io" + "net" + "strconv" "strings" "sync" "time" @@ -12,12 +14,14 @@ import ( "github.com/google/uuid" "github.com/mudler/LocalAI/core/config" "github.com/mudler/LocalAI/core/services/agents" + "github.com/mudler/LocalAI/core/services/cluster" "github.com/mudler/LocalAI/core/services/distributed" "github.com/mudler/LocalAI/core/services/jobs" "github.com/mudler/LocalAI/core/services/messaging" "github.com/mudler/LocalAI/core/services/nodes" "github.com/mudler/LocalAI/core/services/nodes/prefixcache" "github.com/mudler/LocalAI/core/services/storage" + "github.com/mudler/LocalAI/internal" "github.com/mudler/LocalAI/pkg/distributedhdr" "github.com/mudler/LocalAI/pkg/sanitize" "github.com/mudler/xlog" @@ -43,6 +47,35 @@ type DistributedServices struct { Unloader *nodes.RemoteUnloaderAdapter ModelCleanup *nodes.ModelCleanupService + // Cluster is the replica-membership registry: which frontend replicas are + // alive, at which address, and which of them holds a given worker's tunnel. + Cluster *cluster.Registry + // Membership publishes this replica's row and reaps the dead. Nil when no + // peer-reachable address could be determined, which leaves this replica + // invisible to its peers but otherwise fully functional. + Membership *cluster.Membership + // PeerSessions owns the peer links other replicas dialled into this one, + // and relays the streams that arrive on them onto the worker tunnels this + // replica holds. + PeerSessions *cluster.SessionStore + // Peers owns the peer links this replica dialled OUT, the mirror of + // PeerSessions. It is what the relaying dialer opens a stream on when a + // request arrives here for a worker another replica holds. + Peers *cluster.PeerPool + // Tunnels holds the worker tunnels this replica has accepted and keeps the + // node_connections table agreeing with them. It is handed to the membership + // loop, which re-claims what it holds after this replica has been reaped, + // and to the route that accepts a worker's dial. + Tunnels *cluster.TunnelRegistry + // WorkerDialer is how anything in this process reaches a worker: locally + // when this replica holds the tunnel, and through the owning replica when + // it does not. The HTTP layer takes its WebSocket log proxy from here. + WorkerDialer *cluster.WorkerDialer + // BackendClients builds the gRPC clients for worker backend processes, over + // WorkerDialer. Exposed so the model store built in startup.go reaches + // remote models the same way every other caller does. + BackendClients nodes.BackendClientFactory + shutdownOnce sync.Once } @@ -53,6 +86,22 @@ func (ds *DistributedServices) Shutdown() { return } ds.shutdownOnce.Do(func() { + // Peer state first: a replica that is going away should stop claiming + // to be alive before it stops answering, so peers re-home rather than + // dial a process in teardown. + if ds.Membership != nil { + ds.Membership.Stop() + } + if ds.PeerSessions != nil { + ds.PeerSessions.CloseAll() + } + // Both halves of the peer mesh go down together. A pool left open + // holds a WebSocket and two yamux loop goroutines per peer for as long + // as the process lives, and an Open after this reports ErrPoolClosed, + // which is a fact about this process and never node absence. + if ds.Peers != nil { + ds.Peers.Close() + } if ds.Health != nil { ds.Health.Stop() } @@ -162,6 +211,97 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade } xlog.Info("Node registry initialized") + // Replica membership. NewNodeRegistry has just migrated the tables this + // reads, so it has to come after it. + clusterRegistry := cluster.NewRegistry(authDB) + var membership *cluster.Membership + if advertised, err := advertisedPeerAddr(cfg); err != nil { + // Not fatal, and the cost is worth stating exactly rather than as + // "peers cannot reach it", because it is larger than that now. + // + // Without a row in the instances table this replica is not a live + // owner as far as Registry.Owner is concerned: that read joins a + // connection against a live instance, so a worker whose tunnel lands + // HERE is answered as unroutable at every OTHER replica, for as long + // as it stays here. This replica serves that worker perfectly well + // itself; nobody else can. On N replicas behind round robin that is + // (N-1)/N of the traffic for that worker. + // + // It does not refuse to START. Refusing would take out every existing + // single-host deployment, whose route to a local database is loopback + // and which has no peers to be unreachable by; the deployments this + // hurts are multi-replica ones, and telling those two apart at startup + // is a change with its own design and its own specs rather than a line + // here. + // + // What it does not get to do is stay quiet. One startup line scrolls + // away in seconds and the cost is paid for the whole life of the + // process, on a symptom (workers that 5xx from most of the fleet) whose + // obvious reading is "the worker is broken". So this is an ERROR, not a + // warning, and nagUnadvertisedReplica below repeats it for as long as + // the state lasts, naming the workers it is currently costing. + xlog.Error("This replica is not registered in the cluster: no advertised address. Peers cannot reach it, and any worker whose tunnel lands here will be unroutable from every other replica", + "error", err, "knob", "LOCALAI_DISTRIBUTED_ADVERTISE_ADDR") + } else { + membership = cluster.NewMembership(clusterRegistry, cfg.Distributed.InstanceID, advertised, internal.PrintableVersion()) + if err := membership.Start(cfg.Context); err != nil { + return nil, fmt.Errorf("registering this replica in the cluster: %w", err) + } + } + + // The worker tunnels this replica accepts. It claims as the SAME instance + // ID membership registers under, because that is the ID a peer's Owner + // lookup joins a claim against to decide the owner is alive; two IDs here + // would make every claim this replica writes look like it belongs to a + // replica that does not exist. + tunnels := cluster.NewTunnelRegistry(clusterRegistry, cfg.Distributed.InstanceID) + // Without this the re-claim in the heartbeat loop is dead code: a replica + // stalled long enough to be swept loses the connection rows it owned, and + // nothing would ever write them back, so every other replica would answer + // "not connected" for workers that are connected right here. + // + // Nil when no peer-reachable address could be determined above. There is no + // heartbeat loop to hand it to in that case, and no other replica can reach + // this one anyway; the registry is still built, because it is what the + // tunnel endpoint attaches to and what this replica opens its own streams + // through. + if membership != nil { + membership.SetTunnels(tunnels) + } else { + // The runtime symptom the startup line cannot be. See + // nagUnadvertisedReplica. + go nagUnadvertisedReplica(cfg.Context, tunnels.Held, unadvertisedNagInterval, logUnroutableWorkers) + } + + // The links peers dial IN, with the relay installed on them. This is what + // makes more than one replica work: a worker holds one tunnel, it lands on + // one replica, and every request that arrives anywhere else reaches the + // worker through this handler. Passing nil here would leave every such + // request refused, promptly and only at debug level, which presents as a + // worker that is connected and unusable from most of the deployment. + peerSessions := cluster.NewSessionStore(cluster.NewRelay(tunnels).Stream) + // The links this replica dials OUT, the other half of the same mesh. It + // authenticates with the registration token because that is the token the + // peer route checks (see RegisterClusterRoutes); two different tokens here + // would make every peer dial 401 with nothing naming the mismatch. + peers := cluster.NewPeerPool(cfg.Distributed.InstanceID, cfg.Distributed.RegistrationToken, clusterRegistry) + // The one door to every worker. Nothing in the frontend may dial a worker's + // advertised address any more: a worker holds ONE tunnel, it lands on ONE + // replica, and this resolves which replica that is and relays through it + // when it is not this one. The three transports the frontend speaks to a + // worker (gRPC to backend processes, HTTP for file staging and logs, a + // WebSocket for live log streaming) are all pointed at it below. + workerDialer := cluster.NewWorkerDialer(tunnels, peers) + backendClients, err := nodes.NewTunnelClientFactory(cfg.Distributed.RegistrationToken, workerDialer.GRPCDialerFor) + if err != nil { + return nil, fmt.Errorf("wiring the worker backend client factory: %w", err) + } + // Bound to the http tag: the worker ignores the target for it and routes to + // its own file-transfer and log server, wherever that bound. + workerHTTPDialer := nodes.WorkerNetDialerFor(func(nodeID string) func(ctx context.Context, network, addr string) (net.Conn, error) { + return workerDialer.DialerFor(nodeID, cluster.StreamTagHTTP) + }) + // Let scheduling rules be keyed by a model alias. The registry resolves a // rule's name through the config loader to find the model it governs, so an // operator can pin placement to a stable name like "production" and have it @@ -202,6 +342,7 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade cfg.Distributed.StaleNodeThresholdOrDefault(), routerAuthToken, !cfg.Distributed.DisablePerModelHealthCheck, + backendClients, ) // Initialize job store @@ -257,11 +398,13 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade if err != nil { return "", err } - if node.HTTPAddress == "" { - return "", fmt.Errorf("node %s has no HTTP address for file transfer", nodeID) - } - return node.HTTPAddress, nil - }, cfg.Distributed.RegistrationToken) + // An empty HTTPAddress is no longer a refusal. A tunnel-only worker + // reports none and does not need one: the http stream tag ignores + // the target and the worker routes to its own server. The host is + // only ever the URL's host component here, and WorkerHTTPHost + // supplies one that resolves nowhere so it cannot become a dial. + return nodes.WorkerHTTPHost(nodeID, node.HTTPAddress), nil + }, cfg.Distributed.RegistrationToken, workerHTTPDialer) xlog.Info("File stager initialized (HTTP direct transfer)") } // Create RemoteUnloaderAdapter — needed by SmartRouter and startup.go @@ -363,6 +506,7 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade FileStager: fileStager, GalleriesJSON: routerGalleriesJSON, AuthToken: routerAuthToken, + ClientFactory: backendClients, DB: authDB, ConflictResolver: conflictResolver, PrefixProvider: prefixProvider, @@ -421,6 +565,7 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade Unloader: remoteUnloader, Adapter: remoteUnloader, RegistrationToken: cfg.Distributed.RegistrationToken, + ClientFactory: backendClients, DB: authDB, Interval: 30 * time.Second, ScaleDownDelay: 5 * time.Minute, @@ -434,25 +579,127 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade success = true return &DistributedServices{ - Nats: natsClient, - Store: store, - Registry: registry, - Router: router, - Health: healthMon, - Reconciler: reconciler, - JobStore: jobStore, - Dispatcher: dispatcher, - AgentStore: agentStore, - AgentBridge: agentBridge, - DistStores: distStores, - FileMgr: fileMgr, - FileStager: fileStager, - ModelAdapter: modelAdapter, - Unloader: remoteUnloader, - ModelCleanup: modelCleanup, + Nats: natsClient, + Store: store, + Registry: registry, + Router: router, + Health: healthMon, + Reconciler: reconciler, + JobStore: jobStore, + Dispatcher: dispatcher, + AgentStore: agentStore, + AgentBridge: agentBridge, + DistStores: distStores, + FileMgr: fileMgr, + FileStager: fileStager, + ModelAdapter: modelAdapter, + Unloader: remoteUnloader, + ModelCleanup: modelCleanup, + Cluster: clusterRegistry, + Membership: membership, + PeerSessions: peerSessions, + Peers: peers, + Tunnels: tunnels, + WorkerDialer: workerDialer, + BackendClients: backendClients, }, nil } +// unadvertisedNagInterval is how often a replica that could not advertise +// itself says so again. +// +// Five minutes is chosen against the log it lands in, not against the urgency: +// the condition never clears on its own, so this line is either read once and +// acted on or it is noise for the life of the process, and a noisy line gets +// filtered rather than fixed. It is still frequent enough that the state is +// visible in any window of logs an operator pulls while investigating the +// symptom it causes. +const unadvertisedNagInterval = 5 * time.Minute + +// nagUnadvertisedReplica repeats, for as long as the process runs, that this +// replica is invisible to its peers, and names what that is currently costing. +// +// It exists because the deferral it accompanies changed cost between phases and +// nothing about the deployment says so. Before workers held tunnels, a replica +// with no advertised address was merely unreachable BY peers and could still +// dial every worker directly, so a startup warning was proportionate. Now a +// worker's tunnel lands on one replica and every other replica reaches it by +// relaying to the owner, and the owner is resolved by joining the connection +// row against a LIVE INSTANCES ROW - which this replica does not have. So every +// worker that lands here is answered as unroutable everywhere else: on N +// replicas behind round robin, (N-1)/N of that worker's traffic fails, while +// this replica serves it perfectly and reports nothing. +// +// held is passed as a function rather than the registry so this can be driven +// without one, and alarm is passed rather than logged inline so a spec can +// observe the alarms instead of scraping a log. +func nagUnadvertisedReplica(ctx context.Context, held func() []string, every time.Duration, alarm func([]string)) { + ticker := time.NewTicker(every) + defer ticker.Stop() + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + alarm(held()) + } + } +} + +// logUnroutableWorkers says what the state costs RIGHT NOW. +// +// The two cases are kept apart because they call for different urgency and an +// operator can tell them apart at a glance. With no worker held this is a +// misconfiguration that has not been paid for yet; with workers held, every one +// of them is named, because "which worker is broken" is the question the +// symptom sends an operator to ask and the answer is that none of them is. +func logUnroutableWorkers(held []string) { + if len(held) == 0 { + xlog.Warn("This replica is still not registered in the cluster: no advertised address. No worker holds a tunnel here yet; the first that does will be unroutable from every other replica", + "knob", "LOCALAI_DISTRIBUTED_ADVERTISE_ADDR") + return + } + xlog.Error("This replica is not registered in the cluster and holds worker tunnels: those workers are unroutable from every OTHER replica, and requests for their models fail there with no route. The workers are healthy; this replica is invisible", + "workers", held, "worker_count", len(held), "knob", "LOCALAI_DISTRIBUTED_ADVERTISE_ADDR") +} + +// advertisedPeerAddr is the host:port peers dial to reach this replica. +// +// The operator's value wins outright. Otherwise it is derived from the port +// this process serves on and the local address that routes to PostgreSQL, which +// is only a peer-reachable answer when the database is on another host; +// DiscoverAdvertisedAddr refuses rather than guessing when it is not. +func advertisedPeerAddr(cfg *config.ApplicationConfig) (string, error) { + if configured := cfg.Distributed.AdvertiseAddr; configured != "" { + // A configured address skips discovery, so it also skips every check + // discovery makes. Unusable is refused; merely questionable (a + // loopback address, correct on one host and wrong on three) is said + // once and honoured, because refusing it would refuse single-host + // deployments that use it correctly. + reason, err := cluster.CheckAdvertisedAddr(configured) + if err != nil { + return "", err + } + if reason != "" { + xlog.Warn("Configured peer address is not one another host can dial", + "address", configured, "reason", reason, "knob", "LOCALAI_DISTRIBUTED_ADVERTISE_ADDR") + } + return configured, nil + } + if cfg.APIAddress == "" { + return "", fmt.Errorf("no API address to derive a peer port from") + } + _, port, err := net.SplitHostPort(cfg.APIAddress) + if err != nil { + return "", fmt.Errorf("reading the peer port out of API address %q: %w", cfg.APIAddress, err) + } + portNumber, err := strconv.Atoi(port) + if err != nil { + return "", fmt.Errorf("API address %q has a non-numeric port: %w", cfg.APIAddress, err) + } + return cluster.DiscoverAdvertisedAddr(cfg.Auth.DatabaseURL, portNumber) +} + func isPostgresURL(url string) bool { return strings.HasPrefix(url, "postgres://") || strings.HasPrefix(url, "postgresql://") } diff --git a/core/application/startup.go b/core/application/startup.go index abc2f4a17571..c790a1403617 100644 --- a/core/application/startup.go +++ b/core/application/startup.go @@ -283,9 +283,15 @@ func New(opts ...config.AppOption) (*Application, error) { // Wire ModelRouter so grpcModel() delegates to SmartRouter in distributed mode application.modelLoader.SetModelRouter(distSvc.ModelAdapter.AsModelRouter()) // Wire DistributedModelStore so shutdown/list/watchdog can find remote models + // The client factory is not optional here. Without it the store builds + // remote models with no client, and pkg/model.Model.GRPC then dials the + // worker's raw address with gRPC's own dialer, which is the direct dial + // the tunnel replaces; ShutdownModel's Free and the backend monitor's + // Status both reach it. distStore := nodes.NewDistributedModelStore( model.NewInMemoryModelStore(), distSvc.Registry, + distSvc.BackendClients, ) application.modelLoader.SetModelStore(distStore) // Drop the local stub when a model's last replica leaves the registry. diff --git a/core/application/unadvertised_replica_test.go b/core/application/unadvertised_replica_test.go new file mode 100644 index 000000000000..ed9eeb4d6399 --- /dev/null +++ b/core/application/unadvertised_replica_test.go @@ -0,0 +1,73 @@ +// SPDX-License-Identifier: MIT + +package application + +import ( + "context" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +// The runtime symptom for a deferral whose cost changed between phases. +// +// Not refusing to start without an advertised address stays deferred on +// purpose: refusing would take out every single-host deployment. What is not +// deferred is telling the operator, repeatedly, that this replica is invisible +// and which workers that is costing - because the symptom it produces (a worker +// that 5xxs from most of the fleet) reads as a worker problem, and a single +// startup line has scrolled away long before anyone goes looking. +var _ = Describe("the alarm for a replica with no advertised address", func() { + It("keeps firing for as long as the state lasts, and names the workers it costs", func() { + // Repetition is the property. A one-shot alarm is the startup line + // again, which is what was already there and was not enough. + ctx, cancel := context.WithCancel(context.Background()) + DeferCleanup(cancel) + + alarms := make(chan []string, 8) + go nagUnadvertisedReplica(ctx, func() []string { return []string{"w1", "w2"} }, + time.Millisecond, func(held []string) { alarms <- held }) + + // Two, not one: the second is what a one-shot implementation fails. + var first, second []string + Eventually(alarms, "10s").Should(Receive(&first)) + Eventually(alarms, "10s").Should(Receive(&second)) + Expect(first).To(ConsistOf("w1", "w2"), + "the workers this is costing are the answer to the question the symptom provokes") + Expect(second).To(ConsistOf("w1", "w2")) + }) + + It("reads the held set on every tick rather than the one it started with", func() { + // A replica accumulates tunnels while it runs, so an alarm bound to the + // set at startup would name an empty list forever on exactly the + // deployment where the cost is real. + ctx, cancel := context.WithCancel(context.Background()) + DeferCleanup(cancel) + + workers := make(chan []string, 32) + for range 32 { + workers <- []string{"w-late"} + } + alarms := make(chan []string, 8) + go nagUnadvertisedReplica(ctx, func() []string { return <-workers }, + time.Millisecond, func(held []string) { alarms <- held }) + + var got []string + Eventually(alarms, "10s").Should(Receive(&got)) + Expect(got).To(ConsistOf("w-late")) + }) + + It("stops when the process context ends", func() { + ctx, cancel := context.WithCancel(context.Background()) + stopped := make(chan struct{}) + go func() { + defer GinkgoRecover() + nagUnadvertisedReplica(ctx, func() []string { return nil }, time.Hour, func([]string) {}) + close(stopped) + }() + + cancel() + Eventually(stopped, "10s").Should(BeClosed()) + }) +}) diff --git a/core/cli/run.go b/core/cli/run.go index 6b9b3e4dc0b9..4408edd1203c 100644 --- a/core/cli/run.go +++ b/core/cli/run.go @@ -165,6 +165,7 @@ type RunCMD struct { Distributed bool `env:"LOCALAI_DISTRIBUTED" default:"false" help:"Enable distributed mode (requires PostgreSQL + NATS)" group:"distributed"` InstanceID string `env:"LOCALAI_INSTANCE_ID" help:"Unique instance ID for distributed mode (auto-generated UUID if empty)" group:"distributed"` NatsURL string `env:"LOCALAI_NATS_URL" help:"NATS server URL (e.g., nats://localhost:4222)" group:"distributed"` + DistributedAdvertiseAddr string `env:"LOCALAI_DISTRIBUTED_ADVERTISE_ADDR" help:"host:port other frontend replicas dial to reach this one (peer link). Empty = derived from the local address that routes to PostgreSQL, which only works when the database is on another host." group:"distributed"` StorageURL string `env:"LOCALAI_STORAGE_URL" help:"S3-compatible storage endpoint URL (e.g., http://minio:9000)" group:"distributed"` StorageBucket string `env:"LOCALAI_STORAGE_BUCKET" default:"localai" help:"S3 bucket name for object storage" group:"distributed"` StorageRegion string `env:"LOCALAI_STORAGE_REGION" default:"us-east-1" help:"S3 region" group:"distributed"` @@ -351,6 +352,9 @@ func (r *RunCMD) Run(ctx *cliContext.Context) error { if r.InstanceID != "" { opts = append(opts, config.WithDistributedInstanceID(r.InstanceID)) } + if r.DistributedAdvertiseAddr != "" { + opts = append(opts, config.WithDistributedAdvertiseAddr(r.DistributedAdvertiseAddr)) + } if r.NatsURL != "" { opts = append(opts, config.WithNatsURL(r.NatsURL)) } diff --git a/core/cli/workerregistry/client.go b/core/cli/workerregistry/client.go index cf46455c95c0..fb00fb3f1d86 100644 --- a/core/cli/workerregistry/client.go +++ b/core/cli/workerregistry/client.go @@ -8,7 +8,9 @@ import ( "bytes" "context" "encoding/json" + "errors" "fmt" + "io" "net/http" "strings" "sync" @@ -58,9 +60,15 @@ func (c *RegistrationClient) setAuth(req *http.Request) { // RegisterResponse is the JSON body returned by /api/node/register. type RegisterResponse struct { - ID string `json:"id"` - Status string `json:"status,omitempty"` // "pending" until an admin approves the node - APIToken string `json:"api_token,omitempty"` + ID string `json:"id"` + Status string `json:"status,omitempty"` // "pending" until an admin approves the node + APIToken string `json:"api_token,omitempty"` + // TunnelToken is this node's own credential for GET /api/cluster/connect. + // The frontend mints a fresh one on every registration and keeps only its + // hash, so this is the ONLY time the plaintext exists anywhere but in this + // worker's memory: a worker that discards it cannot get it back without + // registering again. + TunnelToken string `json:"tunnel_token,omitempty"` NatsJWT string `json:"nats_jwt,omitempty"` NatsUserSeed string `json:"nats_user_seed,omitempty"` } @@ -87,7 +95,7 @@ func (c *RegistrationClient) RegisterFull(ctx context.Context, body map[string]a defer resp.Body.Close() if resp.StatusCode < 200 || resp.StatusCode >= 300 { - return nil, fmt.Errorf("registration failed with status %d", resp.StatusCode) + return nil, registrationStatusError(resp) } var result RegisterResponse @@ -97,6 +105,58 @@ func (c *RegistrationClient) RegisterFull(ctx context.Context, body map[string]a return &result, nil } +// ErrRegistrationRejected marks a registration the frontend REFUSED, as opposed +// to one it could not answer. +// +// Retrying a refusal cannot change it: the request is wrong, or this worker is +// not allowed to make it. The one that matters in practice is a worker of this +// release registering against a frontend that predates it, which answers +// "address is required for backend workers" with 400, because a worker no +// longer has an address to send. Without this the retry ladder spends four +// minutes on a verdict the frontend reached instantly, and the operator watches +// it before being told anything. +// +// 408 and 429 are deliberately NOT rejections. Both are the frontend asking for +// the same request again later, which is exactly what a retry does. +var ErrRegistrationRejected = errors.New("the frontend refused this registration") + +// maxRegistrationErrorBody bounds how much of a refusal's body is quoted back. +// Enough for a message, not enough for an HTML error page to bury the log line +// it is meant to explain. +const maxRegistrationErrorBody = 512 + +// registrationStatusError turns a non-2xx response into an error that says WHY. +// +// The body is the point. The frontend explains its refusals there +// ("address is required for backend workers", "invalid registration token"), +// and discarding it left an operator with a bare status code: the one line that +// would tell them which of several possible mistakes they made was read off the +// socket and thrown away. +func registrationStatusError(resp *http.Response) error { + detail, err := io.ReadAll(io.LimitReader(resp.Body, maxRegistrationErrorBody)) + if err != nil { + xlog.Debug("Could not read the frontend's registration error body", "status", resp.StatusCode, "error", err) + } + msg := strings.Join(strings.Fields(string(detail)), " ") + base := fmt.Sprintf("registration failed with status %d", resp.StatusCode) + if msg != "" { + base = fmt.Sprintf("%s: %s", base, msg) + } + if isRegistrationRejection(resp.StatusCode) { + return fmt.Errorf("%s: %w", base, ErrRegistrationRejected) + } + return errors.New(base) +} + +// isRegistrationRejection reports whether a status is a verdict rather than a +// condition that may pass. +func isRegistrationRejection(status int) bool { + if status == http.StatusRequestTimeout || status == http.StatusTooManyRequests { + return false + } + return status >= 400 && status < 500 +} + // Register sends a single registration request and returns the node ID and // optional credentials (API token for agent workers, NATS JWT when configured). func (c *RegistrationClient) Register(ctx context.Context, body map[string]any) (nodeID, apiToken, natsJWT, natsSeed string, err error) { @@ -108,27 +168,48 @@ func (c *RegistrationClient) Register(ctx context.Context, body map[string]any) } // RegisterWithRetry retries registration with exponential backoff. +// +// It drops every field of the response it does not name, the tunnel credential +// among them. Callers that need one use RegisterFullWithRetry. func (c *RegistrationClient) RegisterWithRetry(ctx context.Context, body map[string]any, maxRetries int) (nodeID, apiToken, natsJWT, natsSeed string, err error) { + res, err := c.RegisterFullWithRetry(ctx, body, maxRetries) + if err != nil { + return "", "", "", "", err + } + return res.ID, res.APIToken, res.NatsJWT, res.NatsUserSeed, nil +} + +// RegisterFullWithRetry retries registration with exponential backoff and +// returns the whole response. +func (c *RegistrationClient) RegisterFullWithRetry(ctx context.Context, body map[string]any, maxRetries int) (*RegisterResponse, error) { backoff := 2 * time.Second maxBackoff := 30 * time.Second + var err error for attempt := 1; attempt <= maxRetries; attempt++ { - nodeID, apiToken, natsJWT, natsSeed, err = c.Register(ctx, body) + var res *RegisterResponse + res, err = c.RegisterFull(ctx, body) if err == nil { - return nodeID, apiToken, natsJWT, natsSeed, nil + return res, nil + } + if errors.Is(err, ErrRegistrationRejected) { + // A verdict, not an outage. Reported on the first attempt so the + // reason the frontend gave is the first thing in the log rather + // than the last, after the ladder. + return nil, err } if attempt == maxRetries { - return "", "", "", "", fmt.Errorf("failed after %d attempts: %w", maxRetries, err) + return nil, fmt.Errorf("failed after %d attempts: %w", maxRetries, err) } xlog.Warn("Registration failed, retrying", "attempt", attempt, "next_retry", backoff, "error", err) select { case <-ctx.Done(): - return "", "", "", "", ctx.Err() + return nil, ctx.Err() case <-time.After(backoff): } backoff = min(backoff*2, maxBackoff) } - return nodeID, apiToken, natsJWT, natsSeed, err + return nil, err } // Heartbeat sends a single heartbeat POST with the given body. diff --git a/core/cli/workerregistry/client_test.go b/core/cli/workerregistry/client_test.go new file mode 100644 index 000000000000..5870d2524b5c --- /dev/null +++ b/core/cli/workerregistry/client_test.go @@ -0,0 +1,138 @@ +package workerregistry + +import ( + "context" + "errors" + "net/http" + "net/http/httptest" + "sync/atomic" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +// The case these specs exist for: a worker of this release registering against +// a frontend that predates it. A worker no longer sends an address, the old +// frontend requires one, and it answers 400 with the reason in the body. Two +// things used to go wrong there at once. The reason was discarded, so the +// operator saw only "status 400" and had to guess which of several mistakes +// they had made; and the retry ladder spent four minutes on a verdict the +// frontend reached instantly. +var _ = Describe("Registration client refusals", func() { + var ( + attempts atomic.Int32 + status atomic.Int32 + body atomic.Value // string + server *httptest.Server + client *RegistrationClient + // seen carries one token per request the handler served, so a spec can + // wait for the Nth attempt instead of sleeping for however long the + // ladder's backoff happens to be. + seen chan struct{} + ) + + BeforeEach(func() { + attempts.Store(0) + status.Store(int32(http.StatusBadRequest)) + body.Store(`{"error":{"code":400,"message":"address is required for backend workers"}}`) + seen = make(chan struct{}, 64) + server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + attempts.Add(1) + select { + case seen <- struct{}{}: + default: + } + w.WriteHeader(int(status.Load())) + _, _ = w.Write([]byte(body.Load().(string))) + })) + client = &RegistrationClient{FrontendURL: server.URL, HTTPTimeout: 2 * time.Second} + }) + + AfterEach(func() { server.Close() }) + + It("quotes what the frontend said", func() { + _, err := client.RegisterFull(context.Background(), map[string]any{"name": "w1"}) + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring("status 400")) + Expect(err.Error()).To(ContainSubstring("address is required for backend workers")) + }) + + It("marks a 4xx as a refusal", func() { + _, err := client.RegisterFull(context.Background(), map[string]any{"name": "w1"}) + Expect(err).To(MatchError(ErrRegistrationRejected)) + }) + + It("does not mark a 5xx as a refusal", func() { + // A frontend that is restarting or wedged has not judged anything, and + // retrying it is the whole reason the ladder exists. + status.Store(int32(http.StatusBadGateway)) + body.Store("bad gateway") + _, err := client.RegisterFull(context.Background(), map[string]any{"name": "w1"}) + Expect(err).To(HaveOccurred()) + Expect(errors.Is(err, ErrRegistrationRejected)).To(BeFalse()) + }) + + DescribeTable("treats a status that asks for the same request again as retryable", + func(code int) { + status.Store(int32(code)) + body.Store("later") + _, err := client.RegisterFull(context.Background(), map[string]any{"name": "w1"}) + Expect(err).To(HaveOccurred()) + Expect(errors.Is(err, ErrRegistrationRejected)).To(BeFalse()) + }, + Entry("408 Request Timeout", http.StatusRequestTimeout), + Entry("429 Too Many Requests", http.StatusTooManyRequests), + ) + + It("stops the retry ladder on the first refusal", func() { + // Ten attempts on a 400 is roughly four minutes of backoff before the + // operator is told anything, and the answer is the same one the + // frontend gave immediately. + _, err := client.RegisterFullWithRetry(context.Background(), map[string]any{"name": "w1"}, 10) + Expect(err).To(MatchError(ErrRegistrationRejected)) + Expect(err.Error()).To(ContainSubstring("address is required for backend workers")) + Expect(attempts.Load()).To(Equal(int32(1))) + }) + + It("still retries something that is not a refusal", func() { + // The control. Without it, a change that returned on EVERY error would + // pass the spec above and silently delete the retry behaviour a worker + // booting alongside its frontend depends on. + status.Store(int32(http.StatusServiceUnavailable)) + body.Store("starting up") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + done := make(chan error, 1) + go func() { + _, err := client.RegisterFullWithRetry(ctx, map[string]any{"name": "w1"}, 10) + done <- err + }() + + // Two tokens is the whole assertion: the ladder came back for a second + // attempt on a status that is not a verdict. Waiting on the handler + // rather than on a duration makes it exact instead of tolerant. + Eventually(seen).Should(Receive()) + Eventually(seen, "10s").Should(Receive()) + cancel() + + var ladderErr error + Eventually(done).Should(Receive(&ladderErr)) + Expect(ladderErr).To(HaveOccurred()) + Expect(errors.Is(ladderErr, ErrRegistrationRejected)).To(BeFalse()) + Expect(attempts.Load()).To(BeNumerically(">=", 2)) + }) + + It("stops the credential manager's acquire loop on a refusal", func() { + // The default worker path goes through Acquire, not the ladder above, + // and its bound is 100 attempts rather than 10. A refusal there is the + // same verdict and has to end the same way. + mgr := NewNATSCredentialManager(func(ctx context.Context) (*RegisterResponse, error) { + return client.RegisterFull(ctx, map[string]any{"name": "w1"}) + }, true) + _, err := mgr.Acquire(context.Background()) + Expect(err).To(MatchError(ErrRegistrationRejected)) + Expect(attempts.Load()).To(Equal(int32(1))) + }) +}) diff --git a/core/cli/workerregistry/credentials.go b/core/cli/workerregistry/credentials.go index 24dd6f3c8ed7..b023b9916918 100644 --- a/core/cli/workerregistry/credentials.go +++ b/core/cli/workerregistry/credentials.go @@ -2,6 +2,7 @@ package workerregistry import ( "context" + "errors" "fmt" "sync" "time" @@ -50,6 +51,11 @@ type NATSCredentialManager struct { jwt string seed string nodeID string + // tunnelToken is the node's own tunnel credential from the most recent + // registration. It is kept here because every re-registration this manager + // performs ROTATES it, so the tunnel client has to read the current value + // at dial time rather than be handed one at startup. + tunnelToken string } // NewNATSCredentialManager builds a manager over register. When requireCreds is @@ -87,6 +93,22 @@ func (m *NATSCredentialManager) store(res *RegisterResponse) { if res.NatsJWT != "" && res.NatsUserSeed != "" { m.jwt, m.seed = res.NatsJWT, res.NatsUserSeed } + // Guarded the same way the NATS pair is: a response that carries no tunnel + // token (a frontend that predates them, or one whose minting failed) must + // not wipe a working credential this worker already holds. Overwriting with + // "" would lock the tunnel out until the next registration that did carry + // one, which is the opposite of what an empty field means. + if res.TunnelToken != "" { + m.tunnelToken = res.TunnelToken + } +} + +// TunnelToken returns the node's current tunnel credential, empty until one has +// been issued. It is the callback the tunnel client reads on every dial. +func (m *NATSCredentialManager) TunnelToken() string { + m.mu.RLock() + defer m.mu.RUnlock() + return m.tunnelToken } // Current returns the latest NATS credentials (both empty until acquired). @@ -125,6 +147,11 @@ func (m *NATSCredentialManager) Acquire(ctx context.Context) (*RegisterResponse, for attempt := 1; m.maxAttempts <= 0 || attempt <= m.maxAttempts; attempt++ { res, err := m.register(ctx) switch { + case errors.Is(err, ErrRegistrationRejected): + // The frontend refused rather than failed. Waiting through the full + // attempt ladder would delay the operator's only explanation by the + // length of the ladder and change nothing about the answer. + return nil, err case err != nil: lastReason = err xlog.Warn("Registration failed, retrying", "attempt", attempt, "next_retry", backoff, "error", err) diff --git a/core/config/distributed_config.go b/core/config/distributed_config.go index 5a48a84e9b44..bbb141592934 100644 --- a/core/config/distributed_config.go +++ b/core/config/distributed_config.go @@ -13,8 +13,15 @@ import ( // DistributedConfig holds configuration for horizontal scaling mode. // When Enabled is true, PostgreSQL and NATS are required. type DistributedConfig struct { - Enabled bool // --distributed / LOCALAI_DISTRIBUTED - InstanceID string // --instance-id / LOCALAI_INSTANCE_ID (auto-generated UUID if empty) + Enabled bool // --distributed / LOCALAI_DISTRIBUTED + InstanceID string // --instance-id / LOCALAI_INSTANCE_ID (auto-generated UUID if empty) + // AdvertiseAddr is the host:port OTHER REPLICAS dial to reach this one, + // which is not the address this process binds: a replica behind a service + // or a NAT binds one and is reached at another. Empty means "work it out", + // by asking the kernel which local address routes to PostgreSQL; that + // answer is only usable when the database is remote, so a deployment with + // a local or sidecar database has to set this. + AdvertiseAddr string // LOCALAI_DISTRIBUTED_ADVERTISE_ADDR NatsURL string // --nats-url / LOCALAI_NATS_URL StorageURL string // --storage-url / LOCALAI_STORAGE_URL (S3 endpoint) RegistrationToken string // --registration-token / LOCALAI_REGISTRATION_TOKEN (required token for node registration) @@ -195,6 +202,14 @@ func WithDistributedInstanceID(id string) AppOption { } } +// WithDistributedAdvertiseAddr pins the host:port peers dial to reach this +// replica, overriding the route-based discovery. +func WithDistributedAdvertiseAddr(addr string) AppOption { + return func(o *ApplicationConfig) { + o.Distributed.AdvertiseAddr = addr + } +} + func WithNatsURL(url string) AppOption { return func(o *ApplicationConfig) { o.Distributed.NatsURL = url diff --git a/core/http/app.go b/core/http/app.go index 2e1453ac0e38..ec2b864ee8b7 100644 --- a/core/http/app.go +++ b/core/http/app.go @@ -1,12 +1,14 @@ package http import ( + "context" "embed" "errors" "fmt" "io/fs" "math" "mime" + "net" "net/http" "os" "path/filepath" @@ -28,6 +30,7 @@ import ( "github.com/mudler/LocalAI/core/application" "github.com/mudler/LocalAI/core/schema" + clustersvc "github.com/mudler/LocalAI/core/services/cluster" "github.com/mudler/LocalAI/core/services/distributed" "github.com/mudler/LocalAI/core/services/finetune" "github.com/mudler/LocalAI/core/services/galleryop" @@ -566,15 +569,79 @@ func API(application *application.Application) (*echo.Echo, error) { distCfg := application.ApplicationConfig().Distributed var registry *nodes.NodeRegistry var remoteUnloader nodes.NodeCommandSender + // How the admin log-proxy routes reach a worker's own HTTP server. Left nil + // outside distributed mode, where there are no workers and no tunnels; the + // routes then refuse rather than dialling an address directly. + var workerHTTPDialFor nodes.WorkerNetDialerFor if d := application.Distributed(); d != nil { registry = d.Registry if d.Router != nil { remoteUnloader = d.Router.Unloader() } + if d.WorkerDialer != nil { + workerHTTPDialFor = func(nodeID string) func(ctx context.Context, network, addr string) (net.Conn, error) { + return d.WorkerDialer.DialerFor(nodeID, clustersvc.StreamTagHTTP) + } + } } natsCfg := distCfg.NatsAuthConfig() routes.RegisterNodeSelfServiceRoutes(e, registry, distCfg.RegistrationToken, distCfg.AutoApproveNodes, application.AuthDB(), application.ApplicationConfig().Auth.APIKeyHMACSecret, natsCfg) - routes.RegisterNodeAdminRoutes(e, registry, remoteUnloader, application.GalleryService(), opcache, application.ApplicationConfig(), adminMiddleware, application.AuthDB(), application.ApplicationConfig().Auth.APIKeyHMACSecret, application.ApplicationConfig().Distributed.RegistrationToken, natsCfg) + routes.RegisterNodeAdminRoutes(e, registry, remoteUnloader, application.GalleryService(), opcache, application.ApplicationConfig(), adminMiddleware, application.AuthDB(), application.ApplicationConfig().Auth.APIKeyHMACSecret, application.ApplicationConfig().Distributed.RegistrationToken, natsCfg, workerHTTPDialFor) + + // Replica-to-replica peer link. Registered only in distributed mode: in + // single-node mode there are no peers, and the route authenticates with the + // registration token, so publishing it unconditionally would put a + // multiplexer on every single-binary install. + if d := application.Distributed(); d != nil && d.PeerSessions != nil { + if distCfg.RegistrationToken == "" { + // The handler fails closed on an empty token, which is right and + // invisible: without this line an operator sees only 401s on a + // route they never configured, and nothing connecting them to the + // token they did not set. + xlog.Warn("Replica peer link will refuse every dial: no registration token is configured", + "route", clustersvc.PeerPath, "knob", "LOCALAI_REGISTRATION_TOKEN") + } + routes.RegisterClusterRoutes(e, distCfg.RegistrationToken, d.PeerSessions.Accept) + } + + // The worker tunnel, registered unconditionally. Both arguments are nil + // outside distributed mode and the handler refuses every dial then, which + // is what makes registering it always safe; what it buys is the + // route-coverage test walking the route in a plain single-binary + // application, and that test is what holds the rule that an unauthenticated + // dial is refused BEFORE the WebSocket upgrade. + var tunnels *clustersvc.TunnelRegistry + if d := application.Distributed(); d != nil { + tunnels = d.Tunnels + if distCfg.RegistrationToken == "" { + // A different warning from the peer link's, for the same missing + // knob, because what breaks is different. Tunnels themselves work + // without a registration token: each node is minted its own tunnel + // credential at registration whether or not one is configured. What + // is missing is the gate in FRONT of that. With no registration + // token, RegisterNodeEndpoint validates nothing, so anyone who can + // reach this frontend can register a node and be issued a tunnel + // credential for it. + // + // How far that gets them depends on the OTHER knob. With + // auto-approve on, the node is healthy at once and the credential + // works immediately. With it off, the node is pending, and the + // tunnel route refuses a pending node on every dial, so the + // credential is inert until an admin approves it and approval is + // the real gate. Worth stating precisely, because the same commit + // argues exactly this distinction three files away to justify + // minting for pending nodes at all. + // + // This warning replaced one that said the opposite, that tunnels + // would refuse every dial without this token. That was true while + // the tunnel authenticated against the registration token's own + // hash, and stopped being true when nodes got credentials of their + // own. + xlog.Warn("Node registration is unauthenticated, so any caller that can reach this frontend can register a worker and be issued a tunnel credential", + "route", clustersvc.ConnectPath, "knob", "LOCALAI_REGISTRATION_TOKEN") + } + } + routes.RegisterWorkerTunnelRoute(e, registry, tunnels) // Distributed SSE routes (job progress + agent events via NATS) if d := application.Distributed(); d != nil { diff --git a/core/http/auth/public_routes.go b/core/http/auth/public_routes.go index 658205a78f8f..2c1e1dc3cdd1 100644 --- a/core/http/auth/public_routes.go +++ b/core/http/auth/public_routes.go @@ -74,8 +74,26 @@ func isPublicRoute(method, path string) bool { return false } +// 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 +// must not import core/http/auth. Nothing in the compiler holds the two +// together, so a spec does instead, driving a peer request through this +// middleware in core/http/endpoints/cluster/peer_test.go. Moving either string +// without the other turns that spec red, which is the whole reason it exists. +const ClusterPathPrefix = "/api/cluster/" + // usesAlternativeAuthentication identifies requests whose credentials are // validated by route-group middleware instead of the global auth middleware. func usesAlternativeAuthentication(path string) bool { - return strings.HasPrefix(path, "/api/node/") + return strings.HasPrefix(path, "/api/node/") || + strings.HasPrefix(path, ClusterPathPrefix) } diff --git a/core/http/endpoints/cluster/cluster_suite_test.go b/core/http/endpoints/cluster/cluster_suite_test.go new file mode 100644 index 000000000000..00f45e37a672 --- /dev/null +++ b/core/http/endpoints/cluster/cluster_suite_test.go @@ -0,0 +1,13 @@ +package cluster_test + +import ( + "testing" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +func TestClusterEndpoints(t *testing.T) { + RegisterFailHandler(Fail) + RunSpecs(t, "Cluster Endpoints Suite") +} diff --git a/core/http/endpoints/cluster/connect.go b/core/http/endpoints/cluster/connect.go new file mode 100644 index 000000000000..8d175704669c --- /dev/null +++ b/core/http/endpoints/cluster/connect.go @@ -0,0 +1,262 @@ +// SPDX-License-Identifier: MIT + +package cluster + +import ( + "crypto/sha256" + "crypto/subtle" + "encoding/hex" + "errors" + "net/http" + "strings" + + "github.com/gorilla/websocket" + "github.com/labstack/echo/v4" + "github.com/libp2p/go-yamux/v5" + clustersvc "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/nodes" + "github.com/mudler/xlog" + "gorm.io/gorm" +) + +// ConnectHandler serves the door a worker knocks on: it authenticates the dial +// against the node's OWN stored token, upgrades it to a WebSocket, wraps that as +// a yamux server session and attaches it to the tunnel registry. +// +// The worker dials out and never listens, which is the whole point of the +// tunnel: a worker behind NAT, in another cluster, or on a laptop needs no +// inbound port. It is therefore the yamux CLIENT and this side the SERVER, so +// this side owns the even stream IDs and is the side that opens streams. +// Nothing here accepts streams: in this design the frontend asks and the worker +// answers, so a worker that opened a stream into this session would park on the +// accept backlog rather than be served. +// +// The route is registered in every deployment, including single-binary ones, so +// that the route-coverage test under build tag `auth` walks it. A nil registry +// or a nil tunnel registry therefore has to be a real answer rather than a +// panic; see the 503 below. +// +// It is deliberately absent from auth.RouteFeatureRegistry. That registry gates +// 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 + // is what this link wants. The same choice PeerHandler makes. + upgrader := websocket.Upgrader{} + + return func(c echo.Context) error { + // Everything below happens BEFORE the upgrade, and the order inside it + // is load-bearing. The credential is read first because a dial with no + // Authorization header at all is the anonymous case, and the + // route-coverage test issues exactly that, with no query string: it + // must see 401 rather than a 400 about a missing node id. + token, ok := bearerToken(c.Request()) + if !ok { + return echo.NewHTTPError(http.StatusUnauthorized, "unauthorized") + } + + // Not 401. A frontend with no cluster cannot authenticate anybody, and + // 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") + } + + nodeID := c.QueryParam("id") + if nodeID == "" { + return echo.NewHTTPError(http.StatusBadRequest, "missing node id") + } + + node, err := registry.Get(c.Request().Context(), nodeID) + switch { + case errors.Is(err, gorm.ErrRecordNotFound): + // A node this frontend has never seen. Reported as 401 rather than + // 404 so a caller cannot enumerate node IDs by status code. + xlog.Debug("worker tunnel dial named an unknown node", "node", nodeID) + return echo.NewHTTPError(http.StatusUnauthorized, "unauthorized") + case err != nil: + // A query that FAILED is neither a rejection nor an absence. This + // is the phase's standing rule in its HTTP form: telling a worker + // its credentials are wrong when the database merely could not be + // read sends it re-registering instead of retrying, and a worker + // that re-registers has thrown away the identity its tunnel and its + // loaded models are keyed by. + xlog.Error("Looking up a worker for its tunnel dial failed", "node", nodeID, "error", err) + 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 this node + // last registered against a LocalAI that predates per-node tunnel + // credentials, so it holds no secret this route can check and must + // register again; a mismatch means the worker is presenting the wrong + // one, usually a stale credential from before a rotation. One log line + // for both leaves an operator reading "wrong token" while a whole fleet + // of not-yet-restarted workers fails identically. + if node.TunnelTokenHash == "" { + // Debug, not Warn. Every worker in that state fails this way on + // every reconnect, so warning per dial buries the log. + xlog.Debug("refusing a worker tunnel: this node has no tunnel credential, so it has not registered since they were introduced", + "node", nodeID) + return echo.NewHTTPError(http.StatusUnauthorized, "unauthorized") + } + if !authorizedWorker(token, node.TunnelTokenHash) { + 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 + // (provisionAgentWorkerKey, guarded at its call site in + // core/http/endpoints/localai/nodes.go) and its NATS credential + // (attachNatsJWT in the same file). Cited by NAME, not by line: the + // previous version of this comment cited line numbers into a file this + // same commit was editing, and both were stale before it landed. + // + // 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. + xlog.Debug("worker tunnel upgrade failed", "node", nodeID, "error", err) + return nil + } + + sess, err := yamux.Server(clustersvc.WebsocketConn(ws), nil, nil) + if err != nil { + xlog.Error("Worker tunnel session setup failed", "node", nodeID, "error", err) + _ = ws.Close() + 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. + epoch, err := tunnels.Attach(c.Request().Context(), nodeID, sess) + if err != nil { + xlog.Error("Attaching a worker tunnel failed", "node", nodeID, "error", err) + _ = sess.Close() + return nil + } + + xlog.Info("Worker tunnel established", "node", nodeID, "remote", ws.RemoteAddr().String()) + // The session outlives this handler, so something other than the + // request goroutine has to notice it die. yamux closes shutdownCh from + // its receive loop the moment the underlying conn fails + // (go-yamux/v5@v5.1.0/session.go:691-695 calling close at + // session.go:297-311), and its default config keepalives every 30s + // (mux.go:73-74), so a worker that vanishes without a FIN is noticed + // too rather than held forever. + go func() { + <-sess.CloseChan() + // The token Attach returned, never a fresh or zero one. Detach + // matches it by EQUALITY: it identifies THIS attachment, so a + // worker that has already re-dialled onto this replica is not + // evicted by its predecessor's teardown. + tunnels.Detach(nodeID, epoch) + xlog.Debug("worker tunnel closed", "node", nodeID) + }() + return nil + } +} + +// bearerToken returns the token from an Authorization: Bearer header, and +// whether one was present at all. +// +// The presence of a credential and its correctness are separate answers on +// purpose: "no credential" is what decides the pre-upgrade 401, and it has to be +// decidable before anything about the node is known. +func bearerToken(r *http.Request) (string, bool) { + // RFC 7235 makes the scheme case-insensitive; the token after it is not. + const prefix = "Bearer " + header := r.Header.Get("Authorization") + if len(header) < len(prefix) || !strings.EqualFold(header[:len(prefix)], prefix) { + return "", false + } + token := header[len(prefix):] + if token == "" { + return "", false + } + return token, true +} + +// authorizedWorker compares a presented token against the hash stored on the +// node's own row, in constant time. +// +// Against the NODE's OWN tunnel credential, not the deployment's registration +// token. A tunnel is a durable, multiplexed pipe into a worker, and a +// credential that authorizes every worker at once would mean one leak lets an +// attacker impersonate any worker whose ID it can read and take over that +// worker's traffic by claiming its tunnel. The secret compared here is minted +// per node at registration (attachTunnelToken in +// core/http/endpoints/localai/nodes.go), returned to that worker once, and +// stored only as this hash, so knowing the registration token no longer gets +// anyone a tunnel. +// +// Note which column: BackendNode.TunnelTokenHash, not TokenHash. TokenHash is +// still the hash of whatever token the worker registered WITH, which on most +// deployments is the shared registration token, and comparing against it is +// exactly the weakness this replaced. +// +// 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 "a node with no credential authorizes nobody" from a length rule, and +// because the caller logs that case separately. +func authorizedWorker(token, storedHash string) bool { + if storedHash == "" { + return false + } + sum := sha256.Sum256([]byte(token)) + return subtle.ConstantTimeCompare([]byte(hex.EncodeToString(sum[:])), []byte(storedHash)) == 1 +} diff --git a/core/http/endpoints/cluster/connect_test.go b/core/http/endpoints/cluster/connect_test.go new file mode 100644 index 000000000000..15206d50444c --- /dev/null +++ b/core/http/endpoints/cluster/connect_test.go @@ -0,0 +1,391 @@ +package cluster_test + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "net/http" + "net/http/httptest" + "net/url" + "strings" + + "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/core/http/auth" + "github.com/mudler/LocalAI/core/http/routes" + clustersvc "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/nodes" + "github.com/mudler/LocalAI/core/services/testutil" + + "github.com/gorilla/websocket" + "github.com/labstack/echo/v4" + "github.com/libp2p/go-yamux/v5" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "gorm.io/gorm" +) + +// workerToken is the tunnel credential one worker holds. Registration mints it +// per node and keeps only its hash, which is the whole point of the check the +// specs below pin: a second worker's credential, and the deployment-wide +// registration token, are both wrong for this node. +const workerToken = "worker-1-secret" + +// registrationToken stands in for the shared secret every worker in a +// deployment registers with. It is stored on the row too, in a DIFFERENT +// column, and a spec below pins that presenting it does not open a tunnel. +const registrationToken = "deployment-registration-token" + +func tokenHash(token string) string { + sum := sha256.Sum256([]byte(token)) + return hex.EncodeToString(sum[:]) +} + +// bearer builds the header a worker dials with. +func bearer(token string) http.Header { + h := http.Header{} + h.Set("Authorization", "Bearer "+token) + return h +} + +// wsConnectURL is the worker tunnel route on a test server, named as nodeID. +func wsConnectURL(s *httptest.Server, nodeID string) string { + return "ws" + strings.TrimPrefix(s.URL, "http") + clustersvc.ConnectPath + + "?id=" + url.QueryEscape(nodeID) +} + +var _ = Describe("Worker tunnel handler", func() { + var ( + srv *httptest.Server + db *gorm.DB + reg *clustersvc.Registry + tun *clustersvc.TunnelRegistry + nodeID string + ctx context.Context + ) + + 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", + // Both hashes are set, and they differ. That is what a real + // registration produces: TokenHash is the shared token the worker + // registered WITH, TunnelTokenHash is the secret minted FOR it. + TokenHash: tokenHash(registrationToken), + TunnelTokenHash: tokenHash(workerToken), + } + Expect(nodeReg.Register(ctx, node, true)).To(Succeed()) + nodeID = node.ID + Expect(nodeID).ToNot(BeEmpty()) + + reg = clustersvc.NewRegistry(db) + Expect(reg.Register(ctx, "me", "10.0.0.1:8080", "v1")).To(Succeed()) + tun = clustersvc.NewTunnelRegistry(reg, "me") + + e := echo.New() + routes.RegisterWorkerTunnelRoute(e, nodeReg, tun) + srv = httptest.NewServer(e) + DeferCleanup(srv.Close) + }) + + It("refuses an anonymous dial before upgrading", func() { + // A plain GET, not a WebSocket dial: this is exactly what the + // route-coverage test issues, and a handler that upgrades first answers + // it with gorilla's own 400 handshake failure instead of a 401. + req, err := http.NewRequestWithContext(ctx, http.MethodGet, srv.URL+clustersvc.ConnectPath+"?id="+nodeID, nil) + Expect(err).ToNot(HaveOccurred()) + resp, err := http.DefaultClient.Do(req) + Expect(err).ToNot(HaveOccurred()) + defer func() { _ = resp.Body.Close() }() + + Expect(resp.StatusCode).To(Equal(http.StatusUnauthorized)) + Expect(resp.Header.Get("Upgrade")).To(BeEmpty(), + "the handler upgraded an unauthenticated dial") + Expect(tun.Held()).To(BeEmpty()) + }) + + It("refuses a dial that carries no credentials at all", func() { + _, resp, err := websocket.DefaultDialer.Dial(wsConnectURL(srv, nodeID), nil) + Expect(err).To(HaveOccurred()) + Expect(resp).ToNot(BeNil()) + Expect(resp.StatusCode).To(Equal(http.StatusUnauthorized)) + Expect(tun.Held()).To(BeEmpty()) + }) + + It("refuses the deployment's registration token, which this node also stores", func() { + // The registration token is the credential every worker in the + // deployment holds, and it IS on this node's row, in TokenHash. + // Accepting it would mean one leaked shared secret impersonates any + // worker whose ID an attacker can read; the node's own tunnel + // credential is the only thing this route accepts. + _, resp, err := websocket.DefaultDialer.Dial(wsConnectURL(srv, nodeID), bearer(registrationToken)) + Expect(err).To(HaveOccurred()) + Expect(resp).ToNot(BeNil()) + Expect(resp.StatusCode).To(Equal(http.StatusUnauthorized)) + Expect(tun.Held()).To(BeEmpty()) + }) + + It("refuses a dial that names a node it has never seen", func() { + _, resp, err := websocket.DefaultDialer.Dial(wsConnectURL(srv, "no-such-node"), bearer(workerToken)) + Expect(err).To(HaveOccurred()) + Expect(resp).ToNot(BeNil()) + Expect(resp.StatusCode).To(Equal(http.StatusUnauthorized)) + Expect(tun.Held()).To(BeEmpty()) + }) + + It("refuses an authenticated dial that names no node", func() { + _, resp, err := websocket.DefaultDialer.Dial( + "ws"+strings.TrimPrefix(srv.URL, "http")+clustersvc.ConnectPath, bearer(workerToken)) + Expect(err).To(HaveOccurred()) + Expect(resp).ToNot(BeNil()) + Expect(resp.StatusCode).To(Equal(http.StatusBadRequest)) + }) + + It("attaches an authenticated worker and carries bytes to it", func() { + conn, _, err := websocket.DefaultDialer.Dial(wsConnectURL(srv, nodeID), bearer(workerToken)) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = conn.Close() }) + + // The worker is the side that dials, so its half of the mux is the + // yamux CLIENT and the frontend's is the server. + workerSess, err := yamux.Client(clustersvc.WebsocketConn(conn), nil, nil) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = workerSess.Close() }) + + Eventually(tun.Held, "10s").Should(ConsistOf(nodeID)) + + owner, _, err := reg.OwnerRow(ctx, nodeID) + Expect(err).ToNot(HaveOccurred()) + Expect(owner).To(Equal("me"), + "the tunnel was stored without the claim that tells other replicas where it is") + + go func() { + defer GinkgoRecover() + stream, aerr := workerSess.AcceptStream() + if aerr != nil { + return + } + defer func() { _ = stream.Close() }() + buf := make([]byte, 4) + if _, rerr := stream.Read(buf); rerr != nil { + return + } + _, _ = stream.Write(buf) + }() + + stream, err := tun.Open(ctx, nodeID) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = stream.Close() }) + _, err = stream.Write([]byte("ping")) + Expect(err).ToNot(HaveOccurred()) + echoed := make([]byte, 4) + _, err = stream.Read(echoed) + Expect(err).ToNot(HaveOccurred()) + Expect(string(echoed)).To(Equal("ping")) + }) + + It("detaches the tunnel and drops its claim when the worker goes away", func() { + conn, _, err := websocket.DefaultDialer.Dial(wsConnectURL(srv, nodeID), bearer(workerToken)) + Expect(err).ToNot(HaveOccurred()) + Eventually(tun.Held, "10s").Should(ConsistOf(nodeID)) + + Expect(conn.Close()).To(Succeed()) + + Eventually(tun.Held, "10s").Should(BeEmpty(), + "a dead tunnel is still held here, so every dialer routed to this replica gets a socket that carries nothing") + Eventually(func() error { + _, _, err := reg.OwnerRow(ctx, nodeID) + return err + }, "10s").Should(MatchError(clustersvc.ErrNoConnection), + "the claim outlived the socket, so this replica keeps being named the owner of a worker it no longer holds") + }) + + It("refuses a node with no tunnel credential, without falling back to its registration token", func() { + // A node registered by a LocalAI predating per-node tunnel credentials + // produces exactly this row: a registration-token hash in token_hash + // and nothing in tunnel_token_hash. It cannot be back-filled, because + // the plaintext only ever existed in the response that minted it, so + // such a node must register again. + // + // The dial presents the REGISTRATION token, which is still on the row. + // A handler that fell back to token_hash when the tunnel hash is empty + // would let it in, which is the weakness this whole change removed. + Expect(db.Exec(`UPDATE backend_nodes SET tunnel_token_hash = '' WHERE id = ?`, nodeID).Error).To(Succeed()) + + _, resp, err := websocket.DefaultDialer.Dial(wsConnectURL(srv, nodeID), bearer(registrationToken)) + 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 + // the frontend simply could not look them up sends it re-registering + // instead of retrying. + Expect(db.Exec(`DROP TABLE backend_nodes CASCADE`).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.StatusInternalServerError)) + }) +}) + +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", TunnelTokenHash: 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 + // node registry there is nothing to authenticate against, so it must refuse + // every dial rather than publish an unauthenticated multiplexer. + var srv *httptest.Server + + BeforeEach(func() { + e := echo.New() + routes.RegisterWorkerTunnelRoute(e, nil, nil) + srv = httptest.NewServer(e) + DeferCleanup(srv.Close) + }) + + It("refuses an anonymous dial with 401", func() { + req, err := http.NewRequestWithContext(GinkgoT().Context(), http.MethodGet, srv.URL+clustersvc.ConnectPath, nil) + Expect(err).ToNot(HaveOccurred()) + resp, err := http.DefaultClient.Do(req) + Expect(err).ToNot(HaveOccurred()) + defer func() { _ = resp.Body.Close() }() + Expect(resp.StatusCode).To(Equal(http.StatusUnauthorized)) + }) + + It("tells a credentialed worker the frontend has no cluster, rather than rejecting it", func() { + // A single-binary frontend cannot authenticate anybody, and saying + // "unauthorized" would send an operator hunting a token problem that + // does not exist. + _, resp, err := websocket.DefaultDialer.Dial(wsConnectURL(srv, "w1"), bearer(workerToken)) + Expect(err).To(HaveOccurred()) + Expect(resp).ToNot(BeNil()) + Expect(resp.StatusCode).To(Equal(http.StatusServiceUnavailable)) + }) +}) + +var _ = Describe("Worker tunnel auth coverage", func() { + // The same argument the peer link's coverage specs make: the tunnel route + // authenticates a worker against its own stored token, not a session, so it + // only works while its path sits under the prefix the global auth + // middleware exempts. + var srv *httptest.Server + + BeforeEach(func() { + e := echo.New() + // A nil DB with one legacy API key is the cheapest configuration that + // turns the middleware ON without a database. + e.Use(auth.Middleware(nil, &config.ApplicationConfig{ApiKeys: []string{"an-api-key"}})) + routes.RegisterWorkerTunnelRoute(e, nil, nil) + e.GET("/api/nodes", func(c echo.Context) error { return c.NoContent(http.StatusOK) }) + srv = httptest.NewServer(e) + DeferCleanup(srv.Close) + }) + + It("refuses an uncredentialed request to a route outside the cluster prefix", func() { + resp, err := http.Get(srv.URL + "/api/nodes") + Expect(err).ToNot(HaveOccurred()) + defer func() { _ = resp.Body.Close() }() + Expect(resp.StatusCode).To(Equal(http.StatusUnauthorized), + "the global auth middleware is not actually guarding this server, so the assertion below would prove nothing") + }) + + It("lets a worker dial reach the handler, which is the only thing that can authenticate it", func() { + // The worker's token is not one of the API keys the middleware knows, + // so a 503 from the handler's own no-cluster check can only mean the + // request was let through by the middleware. + req, err := http.NewRequestWithContext(GinkgoT().Context(), http.MethodGet, srv.URL+clustersvc.ConnectPath+"?id=w1", nil) + Expect(err).ToNot(HaveOccurred()) + req.Header.Set("Authorization", "Bearer "+workerToken) + + resp, err := http.DefaultClient.Do(req) + Expect(err).ToNot(HaveOccurred()) + defer func() { _ = resp.Body.Close() }() + Expect(resp.StatusCode).To(Equal(http.StatusServiceUnavailable), + "a worker dial must reach the handler; 401 here means the tunnel route left the auth-exempt prefix %q", auth.ClusterPathPrefix) + }) +}) diff --git a/core/http/endpoints/cluster/peer.go b/core/http/endpoints/cluster/peer.go new file mode 100644 index 000000000000..9dcaea8b38de --- /dev/null +++ b/core/http/endpoints/cluster/peer.go @@ -0,0 +1,146 @@ +// SPDX-License-Identifier: MIT + +// Package cluster serves the replica-to-replica link that a LocalAI frontend +// uses to reach a worker tunnel it does not own. A peer dials +// GET /api/cluster/peer, the connection becomes one multiplexed yamux session, +// and the relay opens a stream on it per request. +package cluster + +import ( + "crypto/subtle" + "net/http" + "strings" + + "github.com/gorilla/websocket" + "github.com/labstack/echo/v4" + "github.com/libp2p/go-yamux/v5" + clustersvc "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/xlog" +) + +// PeerHandler upgrades an authenticated peer dial to a WebSocket, wraps it as +// a yamux server session and hands it to onSession. +// +// onSession runs on the request goroutine, so it must return promptly; the +// session outlives the handler because the upgrade hijacks the connection, and +// closing it is the caller's job. +func PeerHandler(token string, onSession func(peerID string, sess *yamux.Session)) echo.HandlerFunc { + // gorilla's default CheckOrigin already restricts a browser to same-origin + // and lets a header-less client (which every peer is) through, so the + // zero value is what this link wants. + upgrader := websocket.Upgrader{} + + return func(c echo.Context) error { + // Reject before upgrading. Upgrading and then closing would give the + // dialer a WebSocket error in place of an HTTP status, and both the + // route-coverage test and a peer's own retry logic read the status. + if !authorizedPeer(c.Request(), token) { + return echo.NewHTTPError(http.StatusUnauthorized, "unauthorized") + } + + // SELF-DECLARED, and knowingly so. Unlike the worker route next door, + // which resolves ?id= to a node row and checks that node's OWN minted + // credential, this route has only the shared cluster token to check, + // so the id is a label and not a claim anything verifies. + // + // What that costs, exactly, for anything already holding the shared + // token (every worker holds it, and it is the same token that + // authenticates registration): it can relay to every worker tunnel this + // replica owns, reaching every backend gRPC process and every worker's + // file-transfer server; by declaring a legitimate replica's id it can + // make SessionStore.Accept evict that replica's inbound link, at will; + // and it can aim the per-session receive window, which PeerLinkConfig + // sizes at roughly 31 GiB of unread data per session, at one replica's + // memory. The first two are not new capabilities in KIND - before + // workers stopped listening, a holder of that token could already dial + // any worker's advertised ports directly - but the token is now the + // only thing between an attacker and the whole fleet's tunnels, and the + // third is a figure written down as sizing guidance that is also a + // budget on a route this open. + // + // It is deferred rather than patched, because the cheap patch does not + // work: checking ?id= against the instances table stops an invented id + // and stops nothing else, since the attack declares a REAL replica's + // id, and it would buy a false sense of a closed hole. Closing it takes + // a credential per replica, minted where a replica joins the instances + // table and presented here, which is a design with its own migration + // and its own specs. Tracked as the phase-3 item named at + // nodes.BackendNode.TunnelTokenHash. + peerID := c.QueryParam("id") + if peerID == "" { + return echo.NewHTTPError(http.StatusBadRequest, "missing peer id") + } + + ws, err := upgrader.Upgrade(c.Response(), c.Request(), nil) + if err != nil { + // Upgrade has already written its own failure to the client. + xlog.Debug("cluster peer link upgrade failed", "peer", peerID, "error", err) + return nil + } + + // Server side of the mux: the dialing peer is the client, so it owns + // the odd stream IDs and this side the even ones. + // + // The SAME configuration the dialler uses, and that is load bearing + // rather than symmetry for its own sake. A yamux receive window is + // advertised by the receiving side, so a nil here left this end on the + // 256 KiB default while the dialler ran at 4 MiB, and the direction + // governed by this end is the one that carries a relayed model artifact + // INTO the replica that owns the worker's tunnel. That direction was + // measured at roughly half the throughput of the same transfer without + // a relay in it. + // + // It also puts the same ceiling on unread data at this end that + // PeerLinkConfig already documents for the dialling end, so a replica + // is now sized against that figure per link in BOTH directions. That + // is the cost of the window being useful at all: a window is a bound + // on data received and not yet read, so a receiver that will not + // buffer cannot advertise one. + sess, err := yamux.Server(clustersvc.WebsocketConn(ws), clustersvc.PeerLinkConfig(), nil) + if err != nil { + xlog.Error("cluster peer link session setup failed", "peer", peerID, "error", err) + _ = ws.Close() + return nil + } + + if onSession == nil { + // Nothing will ever read from this session, so do not leave the + // peer believing it has a live link. + _ = sess.Close() + return nil + } + + xlog.Debug("cluster peer link established", "peer", peerID, "remote", ws.RemoteAddr().String()) + // net/http recovers a panic from this goroutine but does not close a + // hijacked connection afterwards, so a panicking callback would leave + // the peer holding a link nobody accepts streams on: its opens would + // fill the 256-deep backlog and then hang without an error. + defer func() { + if r := recover(); r != nil { + _ = sess.Close() + panic(r) + } + }() + onSession(peerID, sess) + return nil + } +} + +// authorizedPeer compares the request's bearer token with the cluster token in +// constant time, matching the check the worker file-transfer server makes. +// +// Unlike that one, an empty configured token authorizes nobody: this route is +// registered in every deployment, so failing open would publish an +// unauthenticated mux to any caller that can reach the port. +func authorizedPeer(r *http.Request, expected string) bool { + if expected == "" { + return false + } + // RFC 7235 makes the scheme case-insensitive; the token after it is not. + const prefix = "Bearer " + header := r.Header.Get("Authorization") + if len(header) < len(prefix) || !strings.EqualFold(header[:len(prefix)], prefix) { + return false + } + return subtle.ConstantTimeCompare([]byte(header[len(prefix):]), []byte(expected)) == 1 +} diff --git a/core/http/endpoints/cluster/peer_test.go b/core/http/endpoints/cluster/peer_test.go new file mode 100644 index 000000000000..dc8cba3da9e6 --- /dev/null +++ b/core/http/endpoints/cluster/peer_test.go @@ -0,0 +1,256 @@ +package cluster_test + +import ( + "net/http" + "net/http/httptest" + "strings" + + "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/core/http/auth" + "github.com/mudler/LocalAI/core/http/routes" + clustersvc "github.com/mudler/LocalAI/core/services/cluster" + + "github.com/gorilla/websocket" + "github.com/labstack/echo/v4" + "github.com/libp2p/go-yamux/v5" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +// wsPeerURL is the peer route on a test server, named as peer-1. +func wsPeerURL(s *httptest.Server) string { + return "ws" + strings.TrimPrefix(s.URL, "http") + clustersvc.PeerPath + "?id=peer-1" +} + +var _ = Describe("Peer link handler", func() { + var ( + srv *httptest.Server + sessions chan *yamux.Session + ) + + BeforeEach(func() { + sessions = make(chan *yamux.Session, 1) + e := echo.New() + routes.RegisterClusterRoutes(e, "peer-token", func(_ string, s *yamux.Session) { + sessions <- s + }) + srv = httptest.NewServer(e) + DeferCleanup(srv.Close) + }) + + It("rejects a connection with no token", func() { + _, resp, err := websocket.DefaultDialer.Dial(wsPeerURL(srv), nil) + Expect(err).To(HaveOccurred()) + Expect(resp).ToNot(BeNil()) + Expect(resp.StatusCode).To(Equal(http.StatusUnauthorized)) + }) + + It("rejects a connection with the wrong token", func() { + h := http.Header{} + h.Set("Authorization", "Bearer wrong") + _, resp, err := websocket.DefaultDialer.Dial(wsPeerURL(srv), h) + Expect(err).To(HaveOccurred()) + Expect(resp.StatusCode).To(Equal(http.StatusUnauthorized)) + }) + + It("accepts an authenticated peer and yields a usable yamux session", func() { + h := http.Header{} + h.Set("Authorization", "Bearer peer-token") + conn, _, err := websocket.DefaultDialer.Dial(wsPeerURL(srv), h) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = conn.Close() }) + + var serverSess *yamux.Session + Eventually(sessions, "5s").Should(Receive(&serverSess)) + Expect(serverSess).ToNot(BeNil()) + + // The client wraps its side as a yamux CLIENT and opens a stream; the + // server must accept it. This proves the WebSocket was adapted into a + // stream-oriented conn correctly, which is the part most likely to be + // subtly wrong. + clientSess, err := yamux.Client(clustersvc.WebsocketConn(conn), nil, nil) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = clientSess.Close() }) + + go func() { + defer GinkgoRecover() + st, e := clientSess.OpenStream(GinkgoT().Context()) + if e == nil { + _, _ = st.Write([]byte("hello")) + } + }() + + accepted := make(chan []byte, 1) + go func() { + defer GinkgoRecover() + st, e := serverSess.AcceptStream() + if e != nil { + return + } + buf := make([]byte, 5) + if _, e := st.Read(buf); e == nil { + accepted <- buf + } + }() + Eventually(accepted, "10s").Should(Receive(Equal([]byte("hello")))) + }) + + It("reports the peer id it was given", func() { + h := http.Header{} + h.Set("Authorization", "Bearer peer-token") + ids := make(chan string, 1) + e := echo.New() + routes.RegisterClusterRoutes(e, "peer-token", func(id string, _ *yamux.Session) { ids <- id }) + s2 := httptest.NewServer(e) + DeferCleanup(s2.Close) + + conn, _, err := websocket.DefaultDialer.Dial(wsPeerURL(s2), h) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = conn.Close() }) + + Eventually(ids, "5s").Should(Receive(Equal("peer-1"))) + }) + It("rejects every dial when no cluster token is configured", func() { + // The route is registered in every deployment, so an empty configured + // token must authorize nobody. Failing open the way the worker + // file-transfer server's checkBearerToken does would publish an + // unauthenticated yamux multiplexer to anyone who can reach the port. + e := echo.New() + accepted := make(chan *yamux.Session, 1) + routes.RegisterClusterRoutes(e, "", func(_ string, sess *yamux.Session) { accepted <- sess }) + s2 := httptest.NewServer(e) + DeferCleanup(s2.Close) + + for _, header := range []http.Header{nil, {"Authorization": []string{"Bearer "}}, {"Authorization": []string{"Bearer anything"}}} { + _, resp, err := websocket.DefaultDialer.Dial(wsPeerURL(s2), header) + Expect(err).To(HaveOccurred()) + Expect(resp).ToNot(BeNil()) + Expect(resp.StatusCode).To(Equal(http.StatusUnauthorized)) + } + Expect(accepted).ToNot(Receive()) + }) + + It("accepts the bearer scheme in any case", func() { + // RFC 7235 makes the scheme case-insensitive. The token after it is not. + h := http.Header{} + h.Set("Authorization", "bearer peer-token") + conn, _, err := websocket.DefaultDialer.Dial(wsPeerURL(srv), h) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = conn.Close() }) + Eventually(sessions, "5s").Should(Receive()) + }) + + It("closes the session when the callback panics", func() { + // net/http recovers the panic but leaves the hijacked socket open, so + // without the handler's own recover the peer would keep a link nobody + // ever accepts streams on. + e := echo.New() + routes.RegisterClusterRoutes(e, "peer-token", func(_ string, _ *yamux.Session) { + panic("callback exploded") + }) + s2 := httptest.NewServer(e) + DeferCleanup(func() { + // A hijacked connection the handler never closed would park + // httptest's Close forever, turning the assertion below into a + // suite hang. Forcing the conns shut keeps the failure legible. + s2.CloseClientConnections() + s2.Close() + }) + + h := http.Header{} + h.Set("Authorization", "Bearer peer-token") + conn, _, err := websocket.DefaultDialer.Dial(wsPeerURL(s2), h) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = conn.Close() }) + + clientSess, err := yamux.Client(clustersvc.WebsocketConn(conn), nil, nil) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = clientSess.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. The session's own liveness + // is the observable that answers in both directions. + Eventually(clientSess.IsClosed, "10s").Should(BeTrue()) + }) + + It("rejects an authenticated dial that names no peer", func() { + // The session is keyed by peer id, so a nameless link could never be + // looked up again; refusing it is cheaper than leaking it. + h := http.Header{} + h.Set("Authorization", "Bearer peer-token") + _, resp, err := websocket.DefaultDialer.Dial( + "ws"+strings.TrimPrefix(srv.URL, "http")+"/api/cluster/peer", h) + Expect(err).To(HaveOccurred()) + Expect(resp).ToNot(BeNil()) + Expect(resp.StatusCode).To(Equal(http.StatusBadRequest)) + Expect(sessions).ToNot(Receive()) + }) +}) + +var _ = Describe("Peer link auth coverage", func() { + // These specs put the REAL global auth middleware in front of the REAL + // registrar and prove a peer dial reaches the handler anyway. The peer link + // authenticates with the cluster token, not a session, so it only works + // while its path sits under the prefix auth exempts; moving either one + // alone 401s every peer dial, and the two live in packages that must not + // import each other. + // + // The predicate that grants the exemption is unexported, so this asserts on + // its effect rather than on it: what a caller can observe is whether the + // request reaches the handler. + var ( + srv *httptest.Server + sessions chan *yamux.Session + ) + + BeforeEach(func() { + sessions = make(chan *yamux.Session, 1) + e := echo.New() + // A nil DB with one legacy API key is the cheapest configuration that + // turns the middleware ON without a database. With neither, Middleware + // short-circuits to next() and every assertion below would pass against + // a server that has no auth at all. + e.Use(auth.Middleware(nil, &config.ApplicationConfig{ApiKeys: []string{"an-api-key"}})) + routes.RegisterClusterRoutes(e, "peer-token", func(_ string, s *yamux.Session) { sessions <- s }) + // A route outside the cluster prefix, registered on the same server, is + // the control: it proves the middleware in front of both is live. + e.GET("/api/nodes", func(c echo.Context) error { return c.NoContent(http.StatusOK) }) + srv = httptest.NewServer(e) + DeferCleanup(srv.Close) + }) + + It("refuses an uncredentialed request to a route outside the cluster prefix", func() { + resp, err := http.Get(srv.URL + "/api/nodes") + Expect(err).ToNot(HaveOccurred()) + defer func() { _ = resp.Body.Close() }() + Expect(resp.StatusCode).To(Equal(http.StatusUnauthorized), + "the global auth middleware is not actually guarding this server, so the peer-route assertions below would prove nothing") + }) + + It("lets a peer dial reach the handler, which is the only thing that can authenticate it", func() { + // The cluster token is not one of the API keys the middleware knows, so + // a 400 from the handler's own missing-id check can only mean the + // request was let through unauthenticated by the middleware. + req, err := http.NewRequestWithContext(GinkgoT().Context(), http.MethodGet, srv.URL+clustersvc.PeerPath, nil) + Expect(err).ToNot(HaveOccurred()) + req.Header.Set("Authorization", "Bearer peer-token") + + resp, err := http.DefaultClient.Do(req) + Expect(err).ToNot(HaveOccurred()) + defer func() { _ = resp.Body.Close() }() + Expect(resp.StatusCode).To(Equal(http.StatusBadRequest), + "a peer dial must reach the handler; 401 here means the peer route left the auth-exempt prefix %q", auth.ClusterPathPrefix) + }) + + It("completes a full peer handshake through the guarded server", func() { + // The status-code assertion above cannot see the upgrade, and the + // upgrade is what a peer actually does. + h := http.Header{} + h.Set("Authorization", "Bearer peer-token") + conn, _, err := websocket.DefaultDialer.Dial(wsPeerURL(srv), h) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = conn.Close() }) + Eventually(sessions, "5s").Should(Receive()) + }) +}) diff --git a/core/http/endpoints/localai/nodes.go b/core/http/endpoints/localai/nodes.go index bbae523b1025..0571e7d7f8df 100644 --- a/core/http/endpoints/localai/nodes.go +++ b/core/http/endpoints/localai/nodes.go @@ -2,6 +2,7 @@ package localai import ( "context" + "crypto/rand" "crypto/sha256" "crypto/subtle" "encoding/hex" @@ -9,6 +10,7 @@ import ( "errors" "fmt" "io" + "net" "net/http" "net/url" "sync" @@ -75,10 +77,14 @@ func GetNodeEndpoint(registry *nodes.NodeRegistry) echo.HandlerFunc { // RegisterNodeRequest is the request body for registering a new worker node. type RegisterNodeRequest struct { - Name string `json:"name"` - NodeType string `json:"node_type,omitempty"` // "backend" (default) or "agent" - Address string `json:"address"` - HTTPAddress string `json:"http_address,omitempty"` + Name string `json:"name"` + NodeType string `json:"node_type,omitempty"` // "backend" (default) or "agent" + // No address and no http_address. A worker has no inbound endpoint to + // register: it holds one outbound tunnel to a frontend replica and every + // protocol the frontend speaks to it travels on that. An older worker still + // sends both keys and they are ignored, which is the intended outcome: + // storing them would put a dialable-looking endpoint back in the API for + // something nothing dials. Token string `json:"token,omitempty"` TotalVRAM uint64 `json:"total_vram,omitempty"` AvailableVRAM uint64 `json:"available_vram,omitempty"` @@ -140,22 +146,15 @@ func RegisterNodeEndpoint(registry *nodes.NodeRegistry, expectedToken string, au fmt.Sprintf("invalid node_type %q; must be %q or %q", nodeType, nodes.NodeTypeBackend, nodes.NodeTypeAgent))) } - // Backend workers require address; agent workers don't serve gRPC + // A backend worker no longer has to state an address; the tunnel it + // dials is what makes it reachable, and requiring one here would refuse + // exactly the workers this design is for. if req.Name == "" { return c.JSON(http.StatusBadRequest, nodeError(http.StatusBadRequest, "name is required")) } - if nodeType == nodes.NodeTypeBackend && req.Address == "" { - return c.JSON(http.StatusBadRequest, nodeError(http.StatusBadRequest, "address is required for backend workers")) - } if len(req.Name) > 255 { return c.JSON(http.StatusBadRequest, nodeError(http.StatusBadRequest, "name exceeds 255 characters")) } - if len(req.Address) > 512 { - return c.JSON(http.StatusBadRequest, nodeError(http.StatusBadRequest, "address exceeds 512 characters")) - } - if len(req.HTTPAddress) > 512 { - return c.JSON(http.StatusBadRequest, nodeError(http.StatusBadRequest, "http_address exceeds 512 characters")) - } // Hash the token for storage (if provided) var tokenHash string @@ -175,8 +174,6 @@ func RegisterNodeEndpoint(registry *nodes.NodeRegistry, expectedToken string, au node := &nodes.BackendNode{ Name: req.Name, NodeType: nodeType, - Address: req.Address, - HTTPAddress: req.HTTPAddress, TokenHash: tokenHash, TotalVRAM: req.TotalVRAM, AvailableVRAM: req.AvailableVRAM, @@ -244,6 +241,7 @@ func RegisterNodeEndpoint(registry *nodes.NodeRegistry, expectedToken string, au } } + attachTunnelToken(ctx, response, registry, node) attachNatsJWT(response, node, natsCfg) return c.JSON(http.StatusCreated, response) @@ -288,6 +286,78 @@ func ApproveNodeEndpoint(registry *nodes.NodeRegistry, authDB *gorm.DB, hmacSecr } } +// attachTunnelToken mints this node a fresh tunnel credential, stores only its +// hash, and puts the plaintext in the registration response. +// +// It is minted for EVERY node that registers, pending ones included, which is a +// deliberate divergence from the two other per-node credentials in this file: +// the agent worker's API key (provisionAgentWorkerKey) and its NATS JWT +// (attachNatsJWT) are both withheld from a node awaiting approval. Those two +// are bearer grants that WORK the moment they are issued, so issuing one to an +// unapproved node would route around the admin. A tunnel credential is not: +// core/http/endpoints/cluster/connect.go re-reads the node's status on every +// dial and refuses a pending node with 403, so the credential is inert until an +// admin approves and stays inert if approval is revoked. Withholding it would +// instead strand every worker that registers exactly once (the static-NATS +// path in core/services/worker/worker.go does), because approval alone does not +// prompt a re-registration and nothing else can hand it the secret. +// +// Only BACKEND nodes get one, and that is a decision rather than an oversight. +// An agent worker serves no gRPC backends and no file staging; nothing dials +// into it at all, so a tunnel replaces nothing for it and there is no client on +// the agent side that would ever open one. Minting anyway would hand out a +// working credential for a pipe nobody drives, which is surface without a +// feature, and it would contradict every comment in this change that says +// "backend workers, the ones that tunnel". +// +// The gate lives HERE and not in ConnectHandler, which never looks at NodeType. +// It does not need to, PROVIDED an ineligible node ends up with no credential +// rather than merely being handed no new one, because the handler's empty-hash +// branch is what does the refusing. So this CLEARS the column instead of +// returning early, and the difference is not theoretical: Register upserts by +// NAME, so a backend node re-registering as an agent keeps its ID, and +// Register's struct Updates zero-skips TunnelTokenHash while writing the new +// node_type. An early return left a live credential on a row that had become an +// agent. Clearing is what makes "enforcement is structural" true. +// +// It clears unconditionally rather than only when something is there, so the +// invariant holds without depending on what the row happened to contain. The +// cost is one UPDATE per agent registration. +// +// The day agent workers want a tunnel, relaxing the eligibility condition is +// the whole change, and it has to be a deliberate one. +// +// A failure to mint or to store is logged and the response goes out without the +// token. Registration is what gets a worker into the cluster at all, and +// failing it over a credential the worker does not need until it tunnels would +// turn a tunnel problem into a node that cannot join. The worker sees no +// tunnel_token, reports that it has no credential, and retries at its next +// registration. +func attachTunnelToken(ctx context.Context, response map[string]any, registry *nodes.NodeRegistry, node *nodes.BackendNode) { + if node == nil { + return + } + if node.NodeType != nodes.NodeTypeBackend { + // Cleared, not skipped. SetTunnelTokenHash writes the single column + // directly rather than through a struct update, so unlike Register it + // can write an empty value; see its doc. + if err := registry.SetTunnelTokenHash(ctx, node.ID, ""); err != nil { + xlog.Error("Failed to clear the tunnel credential of a node that is not a backend worker", + "node", node.Name, "type", node.NodeType, "error", err) + } + return + } + // crypto/rand.Text: at least 128 bits of randomness, no error to handle and + // no length constant to get wrong. + plaintext := rand.Text() + sum := sha256.Sum256([]byte(plaintext)) + if err := registry.SetTunnelTokenHash(ctx, node.ID, hex.EncodeToString(sum[:])); err != nil { + xlog.Error("Failed to store a tunnel credential for node", "node", node.Name, "error", err) + return + } + response["tunnel_token"] = plaintext +} + // attachNatsJWT adds a per-node NATS user JWT to a register/approve response when minting is enabled. func attachNatsJWT(response map[string]any, node *nodes.BackendNode, natsCfg natsauth.Config) { if !natsCfg.CanMintWorkers() || node == nil || node.Status == nodes.StatusPending { @@ -717,7 +787,7 @@ func DeleteModelOnNodeEndpoint(unloader nodes.NodeCommandSender, registry *nodes // NodeBackendLogsListEndpoint proxies a request to a worker node's /v1/backend-logs // endpoint to list model IDs that have backend logs. -func NodeBackendLogsListEndpoint(registry *nodes.NodeRegistry, registrationToken string) echo.HandlerFunc { +func NodeBackendLogsListEndpoint(registry *nodes.NodeRegistry, registrationToken string, dialFor nodes.WorkerNetDialerFor) echo.HandlerFunc { return func(c echo.Context) error { ctx := c.Request().Context() nodeID := c.Param("id") @@ -726,11 +796,11 @@ func NodeBackendLogsListEndpoint(registry *nodes.NodeRegistry, registrationToken return c.JSON(http.StatusNotFound, nodeError(http.StatusNotFound, "node not found")) } - if node.HTTPAddress == "" { - return c.JSON(http.StatusBadGateway, nodeError(http.StatusBadGateway, "node has no HTTP address")) - } - - resp, err := proxyHTTPToWorker(node.HTTPAddress, "/v1/backend-logs", registrationToken) + // No HTTPAddress guard: a tunnel-only worker reports none, and the + // http stream tag ignores the target anyway. WorkerHTTPHost fills the + // URL's host with something that identifies the node and resolves + // nowhere; the tunnel decides where the bytes go. + resp, err := proxyHTTPToWorker(ctx, dialFor, nodeID, nodes.WorkerHTTPHost(nodeID, node.HTTPAddress), "/v1/backend-logs", registrationToken) if err != nil { return c.JSON(http.StatusBadGateway, nodeError(http.StatusBadGateway, fmt.Sprintf("failed to reach worker: %v", err))) } @@ -745,7 +815,7 @@ func NodeBackendLogsListEndpoint(registry *nodes.NodeRegistry, registrationToken // NodeBackendLogsLinesEndpoint proxies a request to a worker node's // /v1/backend-logs/{modelId} endpoint to get buffered log lines. -func NodeBackendLogsLinesEndpoint(registry *nodes.NodeRegistry, registrationToken string) echo.HandlerFunc { +func NodeBackendLogsLinesEndpoint(registry *nodes.NodeRegistry, registrationToken string, dialFor nodes.WorkerNetDialerFor) echo.HandlerFunc { return func(c echo.Context) error { ctx := c.Request().Context() nodeID := c.Param("id") @@ -756,12 +826,8 @@ func NodeBackendLogsLinesEndpoint(registry *nodes.NodeRegistry, registrationToke return c.JSON(http.StatusNotFound, nodeError(http.StatusNotFound, "node not found")) } - if node.HTTPAddress == "" { - return c.JSON(http.StatusBadGateway, nodeError(http.StatusBadGateway, "node has no HTTP address")) - } - path := "/v1/backend-logs/" + url.PathEscape(modelID) - resp, err := proxyHTTPToWorker(node.HTTPAddress, path, registrationToken) + resp, err := proxyHTTPToWorker(ctx, dialFor, nodeID, nodes.WorkerHTTPHost(nodeID, node.HTTPAddress), path, registrationToken) if err != nil { return c.JSON(http.StatusBadGateway, nodeError(http.StatusBadGateway, fmt.Sprintf("failed to reach worker: %v", err))) } @@ -776,7 +842,7 @@ func NodeBackendLogsLinesEndpoint(registry *nodes.NodeRegistry, registrationToke // NodeBackendLogsWSEndpoint proxies a WebSocket connection to a worker node's // /v1/backend-logs/{modelId}/ws endpoint for real-time log streaming. -func NodeBackendLogsWSEndpoint(registry *nodes.NodeRegistry, registrationToken string) echo.HandlerFunc { +func NodeBackendLogsWSEndpoint(registry *nodes.NodeRegistry, registrationToken string, dialFor nodes.WorkerNetDialerFor) echo.HandlerFunc { browserUpgrader := websocket.Upgrader{ CheckOrigin: func(r *http.Request) bool { origin := r.Header.Get("Origin") @@ -808,15 +874,41 @@ func NodeBackendLogsWSEndpoint(registry *nodes.NodeRegistry, registrationToken s return err } - // Dial the worker WebSocket - workerURL := fmt.Sprintf("ws://%s/v1/backend-logs/%s/ws", node.HTTPAddress, url.PathEscape(modelID)) + // Dial the worker WebSocket over that worker's tunnel. The URL still + // names the worker's registered address, for the Host header; the + // NetDialContext below is what decides where the connection goes. A + // missing dialer is a failure, not a direct dial: see + // nodes.ErrNoWorkerDialer. + workerURL := fmt.Sprintf("ws://%s/v1/backend-logs/%s/ws", nodes.WorkerHTTPHost(nodeID, node.HTTPAddress), url.PathEscape(modelID)) workerHeaders := http.Header{} if registrationToken != "" { workerHeaders.Set("Authorization", "Bearer "+registrationToken) } - workerDialer := websocket.Dialer{HandshakeTimeout: 10 * time.Second} - workerWS, _, err := workerDialer.Dial(workerURL, workerHeaders) + var workerDial func(ctx context.Context, network, addr string) (net.Conn, error) + if dialFor != nil { + workerDial = dialFor(nodeID) + } + if workerDial == nil { + // A JSON body cannot be written here: the response writer was + // hijacked by the upgrade above, so the status line is long gone + // and the write lands nowhere. The browser has to be told the same + // way every other failure past the upgrade tells it, with a close + // frame, and the socket has to be closed or it leaks for the life + // of the process. + // Best-effort: the browser may already have gone, and there is + // nothing left to report the failure to either way. The CLOSE is + // what matters and it is unconditional. + _ = browserWS.WriteMessage(websocket.CloseMessage, + websocket.FormatCloseMessage(websocket.CloseInternalServerErr, "no route to worker")) + _ = browserWS.Close() + xlog.Error("Cannot stream backend logs: no way to reach the worker", + "node", nodeID, "error", nodes.ErrNoWorkerDialer) + return nil + } + + workerDialer := websocket.Dialer{HandshakeTimeout: 10 * time.Second, NetDialContext: workerDial} + workerWS, _, err := workerDialer.DialContext(ctx, workerURL, workerHeaders) if err != nil { browserWS.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseInternalServerErr, "failed to connect to worker")) @@ -1273,10 +1365,25 @@ func DeleteSchedulingEndpoint(registry *nodes.NodeRegistry) echo.HandlerFunc { } } -// proxyHTTPToWorker makes a GET request to a worker's HTTP server with bearer token auth. -func proxyHTTPToWorker(httpAddress, path, token string) (*http.Response, error) { +// proxyHTTPToWorker makes a GET request to a worker's HTTP server with bearer +// token auth, over that worker's tunnel. +// +// The URL still names the worker's registered HTTP address, because that is +// what the Host header and every error message should say; what it no longer +// decides is where the bytes go. dialFor supplies the transport, and a nil one +// is an error rather than a fall back to connecting to httpAddress: a worker +// behind NAT has no address to connect to, and a direct dial is the bypass this +// whole change removes. +func proxyHTTPToWorker(ctx context.Context, dialFor nodes.WorkerNetDialerFor, nodeID, httpAddress, path, token string) (*http.Response, error) { + if dialFor == nil { + return nil, fmt.Errorf("reaching node %s: %w", nodeID, nodes.ErrNoWorkerDialer) + } + dial := dialFor(nodeID) + if dial == nil { + return nil, fmt.Errorf("reaching node %s: %w", nodeID, nodes.ErrNoWorkerDialer) + } reqURL := fmt.Sprintf("http://%s%s", httpAddress, path) - req, err := http.NewRequest("GET", reqURL, nil) + req, err := http.NewRequestWithContext(ctx, "GET", reqURL, nil) if err != nil { return nil, err } @@ -1284,6 +1391,6 @@ func proxyHTTPToWorker(httpAddress, path, token string) (*http.Response, error) req.Header.Set("Authorization", "Bearer "+token) } - client := httpclient.NewWithTimeout(15 * time.Second) + client := httpclient.NewWithTimeout(15*time.Second, httpclient.WithTransport(&http.Transport{DialContext: dial})) return client.Do(req) } diff --git a/core/http/endpoints/localai/nodes_test.go b/core/http/endpoints/localai/nodes_test.go index 19e6a6b07eea..2255116aba8c 100644 --- a/core/http/endpoints/localai/nodes_test.go +++ b/core/http/endpoints/localai/nodes_test.go @@ -4,6 +4,7 @@ import ( "context" "crypto/sha256" "crypto/subtle" + "encoding/hex" "encoding/json" "net/http" "net/http/httptest" @@ -20,6 +21,12 @@ import ( . "github.com/onsi/gomega" ) +// hashOf is how the node registry stores a secret: hex-encoded SHA-256. +func hashOf(secret string) string { + sum := sha256.Sum256([]byte(secret)) + return hex.EncodeToString(sum[:]) +} + var _ = DescribeTable("token validation", func(expectedToken, providedToken string, wantMatch bool) { if expectedToken == "" { @@ -77,6 +84,149 @@ var _ = Describe("Node HTTP handlers", func() { Expect(resp["status"]).To(Equal(nodes.StatusHealthy)) }) + // register posts one registration and returns the decoded response. + register := func(body string, expectedToken string, autoApprove bool) map[string]any { + e := echo.New() + req := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(body)) + req.Header.Set(echo.HeaderContentType, echo.MIMEApplicationJSON) + rec := httptest.NewRecorder() + c := e.NewContext(req, rec) + + handler := RegisterNodeEndpoint(registry, expectedToken, autoApprove, nil, "", natsauth.Config{}) + ExpectWithOffset(1, handler(c)).To(Succeed()) + ExpectWithOffset(1, rec.Code).To(Equal(http.StatusCreated)) + + var resp map[string]any + ExpectWithOffset(1, json.Unmarshal(rec.Body.Bytes(), &resp)).To(Succeed()) + return resp + } + + It("mints a per-node tunnel credential and stores only its hash", func() { + resp := register(`{"name":"worker-tunnel","address":"10.0.0.3:50051","token":"shared-registration-token"}`, + "shared-registration-token", true) + + plaintext, _ := resp["tunnel_token"].(string) + Expect(plaintext).ToNot(BeEmpty()) + // Not the registration token. That is the whole point: a leaked + // registration token plus a known node ID used to open a tunnel, + // because the tunnel authenticated against the hash of exactly the + // value every worker in the deployment holds. + Expect(plaintext).ToNot(Equal("shared-registration-token")) + + node, err := registry.Get(context.Background(), resp["id"].(string)) + Expect(err).ToNot(HaveOccurred()) + // Stored as a hash, never as the secret. + Expect(node.TunnelTokenHash).To(Equal(hashOf(plaintext))) + Expect(node.TunnelTokenHash).ToNot(Equal(plaintext)) + // And it is a DIFFERENT column from the registration token's hash, + // which is what the tunnel used to compare against. + Expect(node.TunnelTokenHash).ToNot(Equal(node.TokenHash)) + Expect(node.TokenHash).To(Equal(hashOf("shared-registration-token"))) + + // The security property, which none of the above actually pins: a + // second node registering with the SAME shared token gets a + // DIFFERENT credential. Everything above is satisfied by a secret + // derived deterministically from the registration token, which + // would isolate nothing; a mutation that did exactly that passed + // every assertion before this one. + other := register(`{"name":"worker-tunnel-2","address":"10.0.0.3:50052","token":"shared-registration-token"}`, + "shared-registration-token", true) + Expect(other["tunnel_token"]).ToNot(Equal(plaintext)) + }) + + It("rotates the tunnel credential on every re-registration", func() { + body := `{"name":"worker-rotate","address":"10.0.0.4:50051"}` + first := register(body, "", true) + second := register(body, "", true) + + Expect(second["id"]).To(Equal(first["id"]), "re-registration must keep the node identity") + firstToken := first["tunnel_token"].(string) + secondToken := second["tunnel_token"].(string) + // Only the hash is stored, so a re-registering worker cannot be told + // the secret it already holds; the alternative to rotating would be + // storing the plaintext. + Expect(secondToken).ToNot(Equal(firstToken)) + + node, err := registry.Get(context.Background(), first["id"].(string)) + Expect(err).ToNot(HaveOccurred()) + Expect(node.TunnelTokenHash).To(Equal(hashOf(secondToken))) + Expect(node.TunnelTokenHash).ToNot(Equal(hashOf(firstToken))) + }) + + It("issues a tunnel credential to a node still awaiting approval", func() { + // Deliberately unlike the agent API key and the NATS JWT, which are + // both withheld from a pending node. Those work the moment they are + // issued; this one does not, because the tunnel endpoint re-reads + // the node's status on every dial and refuses a pending node. A + // worker that registers exactly once would otherwise never receive + // one, since approval alone prompts no re-registration. + first := register(`{"name":"worker-pending","address":"10.0.0.5:50051","token":"shared"}`, "shared", false) + Expect(first["status"]).To(Equal(nodes.StatusPending)) + plaintext, _ := first["tunnel_token"].(string) + Expect(plaintext).ToNot(BeEmpty()) + + // Non-empty alone does not pin per-node-ness, and a review's + // variant of the "derived from the shared token" mutation stayed + // green on exactly that gap. A pending node's credential has to be + // as unpredictable and as per-node as an approved one's, since it + // becomes live the moment an admin approves. + Expect(plaintext).ToNot(Equal("shared")) + node, err := registry.Get(context.Background(), first["id"].(string)) + Expect(err).ToNot(HaveOccurred()) + Expect(node.TunnelTokenHash).To(Equal(hashOf(plaintext))) + Expect(node.TunnelTokenHash).ToNot(Equal(node.TokenHash)) + + second := register(`{"name":"worker-pending-2","address":"10.0.0.5:50052","token":"shared"}`, "shared", false) + Expect(second["status"]).To(Equal(nodes.StatusPending)) + Expect(second["tunnel_token"]).ToNot(Equal(plaintext)) + }) + + It("does not issue a tunnel credential to an agent node", func() { + // An agent worker serves no gRPC backends and no file staging; + // nothing dials into it, so a tunnel replaces nothing for it and no + // client on its side would open one. Minting anyway would be + // credential surface with no feature behind it. + // + // Enforcement is structural rather than a second check: with no + // credential minted, the node's hash stays empty and the tunnel + // route refuses it like any other node without one. + resp := register(`{"name":"agent-1","node_type":"agent"}`, "", true) + Expect(resp["node_type"]).To(Equal(nodes.NodeTypeAgent)) + Expect(resp).ToNot(HaveKey("tunnel_token")) + + node, err := registry.Get(context.Background(), resp["id"].(string)) + Expect(err).ToNot(HaveOccurred()) + Expect(node.TunnelTokenHash).To(BeEmpty()) + }) + + It("clears a tunnel credential when a node stops being a backend node", func() { + // Register upserts BY NAME, so a node can change node_type in place. + // Skipping the mint on the way through leaves the credential the + // node earned as a backend sitting on a row that is now an agent: + // Register's struct Updates zero-skips the column while writing the + // new node_type, so nothing else clears it. ConnectHandler never + // looks at node_type, so that stale hash is a usable tunnel + // credential for a node type that is not supposed to hold one. + // + // This is the same shape as the Register-upserts-by-name hazard + // already carried forward: a name is not an identity. + backend := register(`{"name":"shifty","address":"10.0.0.7:50051"}`, "", true) + Expect(backend["tunnel_token"]).ToNot(BeEmpty()) + + agent := register(`{"name":"shifty","node_type":"agent"}`, "", true) + Expect(agent["id"]).To(Equal(backend["id"]), "re-registration must keep the node identity") + Expect(agent["node_type"]).To(Equal(nodes.NodeTypeAgent)) + Expect(agent).ToNot(HaveKey("tunnel_token")) + + node, err := registry.Get(context.Background(), backend["id"].(string)) + Expect(err).ToNot(HaveOccurred()) + // The claim the gate makes is that an ineligible node HAS no + // credential, not merely that it was not handed a new one. Only + // then is the empty-hash refusal in ConnectHandler the enforcement. + Expect(node.TunnelTokenHash).To(BeEmpty(), + "the node kept the credential it earned as a backend, so the mint-site gate is not structural") + }) + It("returns nats_jwt when account seed is configured", func() { akp, err := nkeys.CreateAccount() Expect(err).ToNot(HaveOccurred()) @@ -139,7 +289,11 @@ var _ = Describe("Node HTTP handlers", func() { Expect(errObj["message"]).To(ContainSubstring("exceeds 255 characters")) }) - It("returns 400 when address is missing for backend node type", func() { + It("registers a backend worker that states no address", func() { + // This used to be a 400. It is the shape every worker now + // registers with: it has no inbound endpoint, it holds one outbound + // tunnel, and refusing it here would refuse exactly the workers the + // tunnel exists for. e := echo.New() body := `{"name":"worker-no-addr"}` req := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(body)) @@ -149,13 +303,35 @@ var _ = Describe("Node HTTP handlers", func() { handler := RegisterNodeEndpoint(registry, "", true, nil, "", natsauth.Config{}) Expect(handler(c)).To(Succeed()) - Expect(rec.Code).To(Equal(http.StatusBadRequest)) + Expect(rec.Code).To(Equal(http.StatusCreated)) - var resp map[string]any - Expect(json.Unmarshal(rec.Body.Bytes(), &resp)).To(Succeed()) - errObj, ok := resp["error"].(map[string]any) - Expect(ok).To(BeTrue()) - Expect(errObj["message"]).To(ContainSubstring("address is required")) + stored, err := registry.GetByName(context.Background(), "worker-no-addr") + Expect(err).ToNot(HaveOccurred()) + Expect(stored.NodeType).To(Equal(nodes.NodeTypeBackend)) + Expect(stored.Address).To(BeEmpty()) + Expect(stored.HTTPAddress).To(BeEmpty()) + }) + + It("stores no address even when a worker still sends one", func() { + // An older worker keeps sending both keys. Storing them would put a + // dialable-looking endpoint back into the API and the Nodes page for + // something nothing dials, and would leave a reader of either one + // unsure which workers are reached how. + e := echo.New() + body := `{"name":"worker-legacy-addr","address":"10.0.0.9:50051","http_address":"10.0.0.9:50050"}` + req := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(body)) + req.Header.Set(echo.HeaderContentType, echo.MIMEApplicationJSON) + rec := httptest.NewRecorder() + c := e.NewContext(req, rec) + + handler := RegisterNodeEndpoint(registry, "", true, nil, "", natsauth.Config{}) + Expect(handler(c)).To(Succeed()) + Expect(rec.Code).To(Equal(http.StatusCreated)) + + stored, err := registry.GetByName(context.Background(), "worker-legacy-addr") + Expect(err).ToNot(HaveOccurred()) + Expect(stored.Address).To(BeEmpty()) + Expect(stored.HTTPAddress).To(BeEmpty()) }) It("returns 400 when node_type is invalid", func() { diff --git a/core/http/react-ui/src/components/nodes/NodePanel.jsx b/core/http/react-ui/src/components/nodes/NodePanel.jsx index 623db00936bf..b73714ab9972 100644 --- a/core/http/react-ui/src/components/nodes/NodePanel.jsx +++ b/core/http/react-ui/src/components/nodes/NodePanel.jsx @@ -19,7 +19,11 @@ export default function NodePanel({ node, models = [], onApprove, onDrain, onRes
{node.name} - {node.address} + {/* A worker has no address to show: it holds an outbound tunnel and + binds nothing routable. Its id is what identifies it in routing + logs, so that is what an operator needs here. Pre-tunnel nodes + may still carry an address until they re-register. */} + {node.address || node.id}
e.stopPropagation()}> {node.status === 'pending' && ( diff --git a/core/http/react-ui/src/pages/NodeDetail.jsx b/core/http/react-ui/src/pages/NodeDetail.jsx index bff7db526ad2..52d3fea1d5f2 100644 --- a/core/http/react-ui/src/pages/NodeDetail.jsx +++ b/core/http/react-ui/src/pages/NodeDetail.jsx @@ -78,7 +78,7 @@ export default function NodeDetail() { navigate('/app/nodes')} className="link-plain">