diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 0e2fd5242..980a62ced 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -52,7 +52,6 @@ jobs: run: make tests lint_and_test_go_pkg: - needs: lint_and_test_python_sdk runs-on: ubuntu-latest steps: - uses: actions/checkout@v3 @@ -75,7 +74,6 @@ jobs: make test-pkg verify_protobufs: - needs: lint_and_test_go_pkg runs-on: ubuntu-latest steps: - uses: actions/checkout@v3 diff --git a/manifests/k3d/beta9.yaml b/manifests/k3d/beta9.yaml index f3f81edec..edda75873 100644 --- a/manifests/k3d/beta9.yaml +++ b/manifests/k3d/beta9.yaml @@ -277,7 +277,11 @@ spec: postgresql: auth: username: root - password: password + # Was the literal string "password" (security baseline P0 #4, + # plaintext infra creds). Substituted at apply time by + # beta9_deploy.sh's envsubst pass, same pattern as SENTRY_DSN below + # — export POSTGRES_PASSWORD before running the deploy script. + password: ${POSTGRES_PASSWORD} database: main image: registry: registry.agentosaurus.com @@ -381,7 +385,10 @@ data: host: postgresql port: 5432 username: root - password: password + # Must match the postgresql HelmChart's auth.password above — + # same POSTGRES_PASSWORD substituted at apply time (security + # baseline P0 #4). + password: ${POSTGRES_PASSWORD} name: main timezone: UTC redis: diff --git a/pkg/agent/agent.go b/pkg/agent/agent.go index 53d339bcf..61eb16c88 100644 --- a/pkg/agent/agent.go +++ b/pkg/agent/agent.go @@ -357,12 +357,14 @@ func (a *Agent) StartInference() error { } if a.ollama.IsRunning() { + status := a.ollama.GetStatus() log.Info(). Int("port", DefaultOllamaPort). Str("tailscale_ip", a.ollama.TailscaleIP()). + Str("gpu_type", status.GPUType). Msg("Inference server ready") a.state.AddLog("Inference: ready on :" + fmt.Sprintf("%d", DefaultOllamaPort)) - a.state.UpdateInference("running", a.ollama.TailscaleIP(), DefaultOllamaPort, nil) + a.state.UpdateInferenceWithGPU("running", a.ollama.TailscaleIP(), DefaultOllamaPort, nil, status.GPUType) } return nil diff --git a/pkg/agent/control.go b/pkg/agent/control.go index b7e84893e..55008f7aa 100644 --- a/pkg/agent/control.go +++ b/pkg/agent/control.go @@ -3,9 +3,11 @@ package agent import ( "bytes" "context" + "crypto/subtle" "encoding/json" "fmt" "net/http" + "os" "strings" "time" @@ -14,6 +16,49 @@ import ( const DefaultControlPort = 9999 +// DefaultControlBindAddr binds the control server to loopback only by +// default (P0-2: it used to hardcode 0.0.0.0, reachable by any Tailnet +// peer with zero authentication). Set BETA9_AGENT_CONTROL_BIND_ADDR to +// override for setups that genuinely need the gateway to reach this over +// the Tailnet — the shared-secret check below still gates every request +// either way. +const DefaultControlBindAddr = "127.0.0.1" + +// isValidOllamaModelName allows only simple local model references +// (name, name:tag, namespace/name:tag) and rejects anything shaped like a +// fully-qualified third-party registry reference. P0-2b: `req.Model` used +// to be forwarded verbatim to Ollama's /api/pull, so a caller could smuggle +// "registry.attacker.com/model:tag" — a supply-chain pull plus a disk-fill +// DoS via the 30-minute pull timeout. Charset is an explicit allowlist +// (not a regex, per this repo's ban on hand-rolled regex for validation); +// the "does the first path segment contain a dot" check is the same +// heuristic Docker/Ollama use to tell a registry hostname apart from a +// plain namespace. +func isValidOllamaModelName(name string) bool { + if name == "" || len(name) > 255 { + return false + } + for _, r := range name { + switch { + case r >= 'a' && r <= 'z', r >= 'A' && r <= 'Z', r >= '0' && r <= '9': + case r == '_' || r == '.' || r == '-' || r == ':' || r == '/': + default: + return false + } + } + if strings.Contains(name, "..") || strings.Contains(name, "//") { + return false + } + parts := strings.Split(name, "/") + if len(parts) > 2 { + return false // deeper than namespace/name is a fully-qualified external registry path + } + if len(parts) == 2 && strings.Contains(parts[0], ".") { + return false // first segment looks like a registry hostname (e.g. registry.attacker.com) + } + return true +} + // ControlServer handles external commands to the agent type ControlServer struct { agent *Agent @@ -44,27 +89,63 @@ func NewControlServer(agent *Agent, port int) *ControlServer { } } +// requireControlToken wraps a handler so it 401s unless the caller presents +// the shared secret from BETA9_AGENT_CONTROL_TOKEN as a Bearer token or +// X-Control-Token header. P0-2: this server previously had zero +// authentication at all — any Tailnet peer reaching :9999 could start/stop +// inference or trigger a model pull. Fails closed (503) if the operator +// hasn't set a token, matching the same policy used for the gateway's +// cluster-admin auth rather than silently allowing everything. +func (c *ControlServer) requireControlToken(next http.HandlerFunc) http.HandlerFunc { + return func(w http.ResponseWriter, r *http.Request) { + token := os.Getenv("BETA9_AGENT_CONTROL_TOKEN") + if token == "" { + writeJSON(w, http.StatusServiceUnavailable, map[string]any{ + "status": "error", + "error": "BETA9_AGENT_CONTROL_TOKEN is not configured; refusing control requests", + }) + return + } + supplied := r.Header.Get("X-Control-Token") + if supplied == "" { + supplied = strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ") + } + if subtle.ConstantTimeCompare([]byte(supplied), []byte(token)) != 1 { + writeJSON(w, http.StatusUnauthorized, map[string]any{ + "status": "error", + "error": "missing or invalid control token", + }) + return + } + next(w, r) + } +} + // Start starts the control server func (c *ControlServer) Start(ctx context.Context) error { mux := http.NewServeMux() // Inference control - mux.HandleFunc("/inference/start", c.handleInferenceStart) - mux.HandleFunc("/inference/stop", c.handleInferenceStop) - mux.HandleFunc("/inference/status", c.handleInferenceStatus) - mux.HandleFunc("/inference/pull", c.handleInferencePull) + mux.HandleFunc("/inference/start", c.requireControlToken(c.handleInferenceStart)) + mux.HandleFunc("/inference/stop", c.requireControlToken(c.handleInferenceStop)) + mux.HandleFunc("/inference/status", c.requireControlToken(c.handleInferenceStatus)) + mux.HandleFunc("/inference/pull", c.requireControlToken(c.handleInferencePull)) // Agent status - mux.HandleFunc("/status", c.handleStatus) - mux.HandleFunc("/health", c.handleHealth) + mux.HandleFunc("/status", c.requireControlToken(c.handleStatus)) + mux.HandleFunc("/health", c.handleHealth) // unauthenticated liveness probe only + bindAddr := os.Getenv("BETA9_AGENT_CONTROL_BIND_ADDR") + if bindAddr == "" { + bindAddr = DefaultControlBindAddr + } c.server = &http.Server{ - Addr: fmt.Sprintf("0.0.0.0:%d", c.port), + Addr: fmt.Sprintf("%s:%d", bindAddr, c.port), Handler: mux, } go func() { - log.Info().Int("port", c.port).Msg("Control server starting") + log.Info().Str("bind_addr", bindAddr).Int("port", c.port).Msg("Control server starting") if err := c.server.ListenAndServe(); err != nil && err != http.ErrServerClosed { log.Error().Err(err).Msg("Control server error") } @@ -130,32 +211,6 @@ func (c *ControlServer) handleInferenceStop(w http.ResponseWriter, r *http.Reque }) } -// isAllowedModelName enforces a narrow allowlist pattern for model names -// accepted by /inference/pull. Ollama model names are of the form -// "/:" or ":" — only A-Z a-z 0-9 . _ - / : -// are permitted, length <= 128. Explicitly rejects ".." to block path -// traversal payloads from reaching the Ollama daemon. -func isAllowedModelName(name string) bool { - if name == "" || len(name) > 128 { - return false - } - // Block path-traversal sequences outright. - if strings.Contains(name, "..") { - return false - } - for _, ch := range name { - switch { - case ch >= 'a' && ch <= 'z': - case ch >= 'A' && ch <= 'Z': - case ch >= '0' && ch <= '9': - case ch == '.' || ch == '_' || ch == '-' || ch == '/' || ch == ':': - default: - return false - } - } - return true -} - // handleInferencePull pulls a model and streams progress to TUI logs func (c *ControlServer) handleInferencePull(w http.ResponseWriter, r *http.Request) { if r.Method != http.MethodPost { @@ -186,12 +241,10 @@ func (c *ControlServer) handleInferencePull(w http.ResponseWriter, r *http.Reque return } - // P0-A model-name allowlist: reject anything with shell metacharacters - // or path-traversal sequences before forwarding to the Ollama daemon. - if !isAllowedModelName(req.Model) { + if !isValidOllamaModelName(req.Model) { writeJSON(w, http.StatusBadRequest, map[string]any{ "status": "error", - "error": "Invalid model name", + "error": "Model name must be a simple local reference (name[:tag] or namespace/name[:tag]), not a fully-qualified registry path", }) return } diff --git a/pkg/agent/keepalive.go b/pkg/agent/keepalive.go index 4bad7b18e..1c50a7f4c 100644 --- a/pkg/agent/keepalive.go +++ b/pkg/agent/keepalive.go @@ -24,10 +24,11 @@ type KeepalivePayload struct { } type InferenceStatus struct { - Status string `json:"status"` // stopped, starting, running, error - IP string `json:"ip,omitempty"` - Port int `json:"port,omitempty"` - Models []string `json:"models,omitempty"` + Status string `json:"status"` // stopped, starting, running, error + IP string `json:"ip,omitempty"` + Port int `json:"port,omitempty"` + Models []string `json:"models,omitempty"` + GPUType string `json:"gpu_type,omitempty"` // e.g. "MPS", "CUDA", "" } // KeepaliveLoop manages periodic keepalive updates to the gateway @@ -167,9 +168,10 @@ func (k *KeepaliveLoop) sendKeepalive(ctx context.Context) bool { // Add inference status if state is available if k.state != nil { inferenceObj := &InferenceStatus{ - Status: k.state.InferenceStatus, - Port: k.state.InferencePort, - Models: k.state.InferenceModels, + Status: k.state.InferenceStatus, + Port: k.state.InferencePort, + Models: k.state.InferenceModels, + GPUType: k.state.InferenceGPUType, } // If inference is running, use Tailscale IP if k.state.InferenceStatus == "running" && k.state.InferenceIP != "" { diff --git a/pkg/agent/state.go b/pkg/agent/state.go index e94675fdc..8c13d4448 100644 --- a/pkg/agent/state.go +++ b/pkg/agent/state.go @@ -55,10 +55,11 @@ type AgentState struct { TotalJobs int // Inference - InferenceStatus string // "stopped", "starting", "running", "error" - InferenceIP string - InferencePort int - InferenceModels []string // List of loaded models + InferenceStatus string // "stopped", "starting", "running", "error" + InferenceIP string + InferencePort int + InferenceModels []string // List of loaded models + InferenceGPUType string // e.g. "MPS", "CUDA", "" // Logs (ring buffer for TUI display) Logs []string @@ -74,18 +75,19 @@ type AgentStateSnapshot struct { CPUPercent float64 MemoryPercent float64 GPUCount int - StartTime time.Time - LastHeartbeat time.Time - HeartbeatStatus string - Jobs []JobInfo - RunningJobs int - TotalJobs int - InferenceStatus string - InferenceIP string - InferencePort int - InferenceModels []string - Logs []string - MaxLogs int + StartTime time.Time + LastHeartbeat time.Time + HeartbeatStatus string + Jobs []JobInfo + RunningJobs int + TotalJobs int + InferenceStatus string + InferenceIP string + InferencePort int + InferenceModels []string + InferenceGPUType string + Logs []string + MaxLogs int } // Uptime returns the agent uptime @@ -215,22 +217,23 @@ func (s *AgentState) GetSnapshot() AgentStateSnapshot { // Build snapshot field-by-field to avoid copying the mutex snapshot := AgentStateSnapshot{ - MachineID: s.MachineID, - PoolName: s.PoolName, - Gateway: s.Gateway, - Status: s.Status, - CPUPercent: s.CPUPercent, - MemoryPercent: s.MemoryPercent, - GPUCount: s.GPUCount, - StartTime: s.StartTime, - LastHeartbeat: s.LastHeartbeat, - HeartbeatStatus: s.HeartbeatStatus, - RunningJobs: s.RunningJobs, - TotalJobs: s.TotalJobs, - InferenceStatus: s.InferenceStatus, - InferenceIP: s.InferenceIP, - InferencePort: s.InferencePort, - MaxLogs: s.MaxLogs, + MachineID: s.MachineID, + PoolName: s.PoolName, + Gateway: s.Gateway, + Status: s.Status, + CPUPercent: s.CPUPercent, + MemoryPercent: s.MemoryPercent, + GPUCount: s.GPUCount, + StartTime: s.StartTime, + LastHeartbeat: s.LastHeartbeat, + HeartbeatStatus: s.HeartbeatStatus, + RunningJobs: s.RunningJobs, + TotalJobs: s.TotalJobs, + InferenceStatus: s.InferenceStatus, + InferenceIP: s.InferenceIP, + InferencePort: s.InferencePort, + InferenceGPUType: s.InferenceGPUType, + MaxLogs: s.MaxLogs, } // Deep copy slices @@ -246,11 +249,17 @@ func (s *AgentState) GetSnapshot() AgentStateSnapshot { // UpdateInference updates inference server status func (s *AgentState) UpdateInference(status, ip string, port int, models []string) { + s.UpdateInferenceWithGPU(status, ip, port, models, "") +} + +// UpdateInferenceWithGPU updates inference server status including GPU type +func (s *AgentState) UpdateInferenceWithGPU(status, ip string, port int, models []string, gpuType string) { s.mu.Lock() defer s.mu.Unlock() s.InferenceStatus = status s.InferenceIP = ip s.InferencePort = port + s.InferenceGPUType = gpuType // Copy slice to prevent data races from caller mutations s.InferenceModels = make([]string, len(models)) copy(s.InferenceModels, models) diff --git a/pkg/api/v1/machine.go b/pkg/api/v1/machine.go index 792e32130..0e2e8c52e 100644 --- a/pkg/api/v1/machine.go +++ b/pkg/api/v1/machine.go @@ -2,10 +2,12 @@ package apiv1 import ( "fmt" + "net" "net/http" "strconv" "github.com/labstack/echo/v4" + "github.com/rs/zerolog/log" "github.com/beam-cloud/beta9/pkg/auth" "github.com/beam-cloud/beta9/pkg/network" @@ -37,10 +39,11 @@ type MachineKeepaliveRequest struct { } type InferenceStatus struct { - Status string `json:"status"` // stopped, starting, running, error - IP string `json:"ip,omitempty"` - Port int `json:"port,omitempty"` - Models []string `json:"models,omitempty"` + Status string `json:"status"` // stopped, starting, running, error + IP string `json:"ip,omitempty"` + Port int `json:"port,omitempty"` + Models []string `json:"models,omitempty"` + GPUType string `json:"gpu_type,omitempty"` // e.g. "MPS", "CUDA", "" } type MachineGroup struct { @@ -49,37 +52,56 @@ type MachineGroup struct { routerGroup *echo.Group config types.AppConfig workerRepo repository.WorkerRepository - inferenceRegistry interface { - RegisterNode(info interface{}) - UpdateHeartbeat(nodeID string) - UpdateNodeModels(nodeID string, models interface{}) - } + inferenceRegistry InferenceRegistry } -// Interface for model registry to avoid circular imports if possible, -// or we need to import it. Since gateway imports apiv1, apiv1 cannot import gateway. -// We'll define the interface here or pass it as 'any' and cast it, or better, move model registry to a shared package. -// For now, let's use a narrow interface matching methods we need. +// InferenceRegistry mirrors the subset of gateway.ModelRegistry's methods +// this package needs. apiv1 can't import gateway (gateway already imports +// apiv1), so this narrow interface breaks the cycle — but it must use the +// SAME concrete parameter types gateway.ModelRegistry actually implements +// (pkg/types.NodeInferenceInfo / pkg/types.ModelInfo, both already shared +// and importable here). The previous version declared RegisterNode(any) / +// UpdateNodeModels(nodeID string, any): Go requires exact signature matches +// for interface satisfaction, so *gateway.ModelRegistry never satisfied it +// and `inferenceRegistry.(InferenceRegistry)` below panicked on every +// gateway startup (P0, beta9 security baseline 2026-04-20 onward). type InferenceRegistry interface { - RegisterNode(info any) + RegisterNode(info *types.NodeInferenceInfo) UpdateHeartbeat(nodeID string) - UpdateNodeModels(nodeID string, models any) + UpdateNodeModels(nodeID string, models map[string]*types.ModelInfo) } func NewMachineGroup(g *echo.Group, providerRepo repository.ProviderRepository, tailscale *network.Tailscale, config types.AppConfig, workerRepo repository.WorkerRepository, inferenceRegistry any) *MachineGroup { + // Non-panicking assertion as defense in depth: if a caller ever passes + // something that doesn't satisfy InferenceRegistry, disable inference + // keepalive/registration for this group instead of crashing the gateway. + registry, ok := inferenceRegistry.(InferenceRegistry) + if !ok && inferenceRegistry != nil { + log.Warn(). + Str("type", fmt.Sprintf("%T", inferenceRegistry)). + Msg("inferenceRegistry does not implement apiv1.InferenceRegistry; inference keepalive/registration disabled for this machine group") + } group := &MachineGroup{routerGroup: g, providerRepo: providerRepo, tailscale: tailscale, config: config, workerRepo: workerRepo, - inferenceRegistry: inferenceRegistry.(InferenceRegistry), + inferenceRegistry: registry, } g.GET("/:workspaceId/gpus", auth.WithWorkspaceAuth(group.GPUCounts)) - g.POST("/register", group.RegisterMachine) - g.POST("/keepalive", group.MachineKeepalive) - g.GET("/config", group.GetConfig) - g.GET("/list", group.ListPoolMachines) + // RegisterMachine/MachineKeepalive/GetConfig/ListPoolMachines all do + // `cc, _ := ctx.(*auth.HttpAuthContext)` and immediately dereference + // `cc.AuthInfo...` with no ok-check. Unwrapped, a caller with no token + // hits fail-open AuthMiddleware (ctx stays plain echo.Context), so cc + // is nil and every one of these routes nil-derefs on the first line + // (security baseline P1: "nil-deref panic on unauth /v1/machine/*"). + // auth.WithAuth guarantees ctx is *HttpAuthContext before the handler + // runs, which both fixes the panic and requires real authentication. + g.POST("/register", auth.WithAuth(group.RegisterMachine)) + g.POST("/keepalive", auth.WithAuth(group.MachineKeepalive)) + g.GET("/config", auth.WithAuth(group.GetConfig)) + g.GET("/list", auth.WithAuth(group.ListPoolMachines)) return group } @@ -202,22 +224,22 @@ func (g *MachineGroup) RegisterMachine(ctx echo.Context) error { }) } -// Helper structs for reflection/dynamic typing without importing gateway -type NodeInferenceInfo struct { - NodeID string `json:"node_id"` - TailscaleIP string `json:"tailscale_ip"` - Port int `json:"port"` - GPUType string `json:"gpu_type"` - TotalVRAM int64 `json:"total_vram_mb"` - AvailableVRAM int64 `json:"available_vram_mb"` - Models map[string]*ModelInfo `json:"models"` -} +// tailscaleCGNAT / isTailscaleAddr duplicate pkg/gateway's inference_handlers.go +// helper of the same purpose — apiv1 can't import gateway (gateway already +// imports apiv1), so this is a small, intentional duplication rather than a +// new shared package for one four-line check. If a third caller ever needs +// it, move both copies to pkg/network. +var tailscaleCGNAT = func() *net.IPNet { + _, n, err := net.ParseCIDR("100.64.0.0/10") + if err != nil { + panic(err) + } + return n +}() -type ModelInfo struct { - Name string `json:"name"` - LoadState string `json:"load_state"` - LastUsed any `json:"last_used"` - LoadedAt any `json:"loaded_at"` +func isTailscaleAddr(ip string) bool { + parsed := net.ParseIP(ip) + return parsed != nil && tailscaleCGNAT.Contains(parsed) } func (g *MachineGroup) MachineKeepalive(ctx echo.Context) error { @@ -254,17 +276,24 @@ func (g *MachineGroup) MachineKeepalive(ctx echo.Context) error { // Update inference status if available if request.Inference != nil && g.inferenceRegistry != nil { + if request.Inference.Status == "running" && request.Inference.IP != "" && !isTailscaleAddr(request.Inference.IP) { + // P1 (beta9 security baseline): a compromised-but-authenticated + // worker token could otherwise redirect all chat/embed traffic + // for its machine_id to an arbitrary IP via keepalive. Same + // CGNAT check as /inference/nodes/register in the gateway package. + return HTTPBadRequest("inference.ip must be within the Tailscale CGNAT range (100.64.0.0/10)") + } if request.Inference.Status == "running" { // Register if needed (idempotent usually) or just update heartbeat // Since we don't have full VRAM info here (it's in metrics but flat), // we construct a best-effort update. // Transform models list to map - modelsMap := make(map[string]*ModelInfo) + modelsMap := make(map[string]*types.ModelInfo) for _, m := range request.Inference.Models { - modelsMap[m] = &ModelInfo{ + modelsMap[m] = &types.ModelInfo{ Name: m, - LoadState: "ready", // Assume ready if reported + LoadState: types.LoadStateReady, // Assume ready if reported } } @@ -282,18 +311,17 @@ func (g *MachineGroup) MachineKeepalive(ctx echo.Context) error { // It probably should rely on keepalive. // So let's do a RegisterNode call here with available info - info := &NodeInferenceInfo{ + gpuType := request.Inference.GPUType + if gpuType == "" { + gpuType = "MPS" // default for backwards compatibility + } + info := &types.NodeInferenceInfo{ NodeID: request.MachineID, TailscaleIP: request.Inference.IP, Port: request.Inference.Port, - GPUType: "MPS", // TODO: Infer from metrics? + GPUType: gpuType, Models: modelsMap, - // VRAM from metrics if available } - // metrics has GPU info? - // request.Metrics is *types.ProviderMachineMetrics - // It has GpuCount but maybe not VRAM details easily? - // Let's rely on update for now. g.inferenceRegistry.RegisterNode(info) } diff --git a/pkg/api/v1/workspace.go b/pkg/api/v1/workspace.go index dc76a6b6a..9bfdacadf 100644 --- a/pkg/api/v1/workspace.go +++ b/pkg/api/v1/workspace.go @@ -31,7 +31,13 @@ func NewWorkspaceGroup(g *echo.Group, backendRepo repository.BackendRepository, defaultStorageClient: defaultStorageClient, } - g.POST("", group.CreateWorkspace) + // CreateWorkspace does `cc, _ := ctx.(*auth.HttpAuthContext)` and + // immediately dereferences cc.AuthInfo — unwrapped, an unauthenticated + // caller hits fail-open AuthMiddleware, cc is nil, and the handler + // nil-derefs (same P1 class as the /v1/machine/* routes). + // WithClusterAdminAuth both fixes that and enforces the check the + // handler body was already trying (and failing) to do itself. + g.POST("", auth.WithClusterAdminAuth(group.CreateWorkspace)) g.GET("/current", auth.WithAuth(group.CurrentWorkspace)) g.GET("/:workspaceId/export", auth.WithStrictWorkspaceAuth(group.ExportWorkspaceConfig)) g.POST("/:workspaceId/set-external-storage", auth.WithStrictWorkspaceAuth(group.SetExternalWorkspaceStorage)) diff --git a/pkg/gateway/inference_handlers.go b/pkg/gateway/inference_handlers.go index 769b7ea26..e6696c1e7 100644 --- a/pkg/gateway/inference_handlers.go +++ b/pkg/gateway/inference_handlers.go @@ -2,15 +2,37 @@ package gateway import ( "context" + "net" "net/http" "time" "github.com/labstack/echo/v4" "github.com/rs/zerolog/log" + "github.com/beam-cloud/beta9/pkg/auth" "github.com/beam-cloud/beta9/pkg/types" ) +// tailscaleCGNAT is the CGNAT range Tailscale assigns node addresses from +// (100.64.0.0/10). A registered node's IP must fall inside it — otherwise a +// caller can point routed inference/chat traffic at an arbitrary host +// (loopback, RFC-1918, link-local/cloud-metadata 169.254.169.254, etc.). +// This is SSRF finding P0 (2026-05-18): independent of the WithClusterAdminAuth +// fix above, since even a legitimate cluster-admin caller should never be +// able to register a non-Tailscale address as an inference node. +var tailscaleCGNAT = func() *net.IPNet { + _, n, err := net.ParseCIDR("100.64.0.0/10") + if err != nil { + panic(err) + } + return n +}() + +func isTailscaleIP(ip string) bool { + parsed := net.ParseIP(ip) + return parsed != nil && tailscaleCGNAT.Contains(parsed) +} + // ============================================================================ // Inference HTTP Handlers - API endpoints for inference routing // ============================================================================ @@ -32,20 +54,32 @@ func NewInferenceService(ctx context.Context, registry *ModelRegistry) *Inferenc } // RegisterRoutes registers inference routes on the Echo router +// +// `g` already has auth.AuthMiddleware attached at the group level (see +// gateway.go), but that middleware fails open when no token is presented — +// it exists to populate *auth.HttpAuthContext when a token IS given, not to +// reject anonymous callers itself. Every handler below therefore needs its +// own auth.WithAuth (or auth.WithClusterAdminAuth) wrapper, or it runs for +// unauthenticated callers exactly as it did before the middleware ran. This +// was P0-1 in the beta9 security baseline: node registration, heartbeat, +// and model load/unload are cluster-admin operations (a caller can redirect +// live inference traffic by re-registering a node), so they require a +// TokenTypeClusterAdmin token; chat/embeddings/model+node listing just need +// any authenticated caller. func (s *InferenceService) RegisterRoutes(g *echo.Group) { // OpenAI-compatible endpoints - g.POST("/chat/completions", s.handleChat) - g.POST("/embeddings", s.handleEmbed) + g.POST("/chat/completions", auth.WithAuth(s.handleChat)) + g.POST("/embeddings", auth.WithAuth(s.handleEmbed)) // Model management - g.GET("/models", s.handleListModels) - g.POST("/models/:model/load", s.handleLoadModel) - g.POST("/models/:model/unload", s.handleUnloadModel) + g.GET("/models", auth.WithAuth(s.handleListModels)) + g.POST("/models/:model/load", auth.WithClusterAdminAuth(s.handleLoadModel)) + g.POST("/models/:model/unload", auth.WithClusterAdminAuth(s.handleUnloadModel)) // Node management - g.GET("/nodes", s.handleListNodes) - g.POST("/nodes/register", s.handleRegisterNode) - g.POST("/nodes/:nodeId/heartbeat", s.handleHeartbeat) + g.GET("/nodes", auth.WithAuth(s.handleListNodes)) + g.POST("/nodes/register", auth.WithClusterAdminAuth(s.handleRegisterNode)) + g.POST("/nodes/:nodeId/heartbeat", auth.WithClusterAdminAuth(s.handleHeartbeat)) // Health g.GET("/health", s.handleHealth) @@ -250,6 +284,12 @@ func (s *InferenceService) handleRegisterNode(c echo.Context) error { }) } + if !isTailscaleIP(req.TailscaleIP) { + return c.JSON(http.StatusBadRequest, map[string]string{ + "error": "tailscale_ip must be within the Tailscale CGNAT range (100.64.0.0/10)", + }) + } + if req.Port == 0 { req.Port = 11434 // Default Ollama port } diff --git a/sdk/pyproject.toml b/sdk/pyproject.toml index 57ada81c7..4e6b6a71a 100644 --- a/sdk/pyproject.toml +++ b/sdk/pyproject.toml @@ -10,7 +10,7 @@ dependencies = [ "grpcio==1.69.0", "asgiref<4.0.0,>=3.8.1", "cloudpickle<4.0.0,>=3.0.0", - "rich<14.0.0,>=13.9.4", + "rich<15.0.0,>=13.9.4", "click<9.0.0,>=8.1.7", "protobuf<5.0.0,>=4.25.1", "fastapi<1.0.0,>=0.115.11", @@ -22,7 +22,7 @@ dependencies = [ "prompt-toolkit<4.0.0,>=3.0.48", "requests<3.0.0,>=2.32.3", "paramiko<4.0.0,>=3.5.0", - "websockets>=13,<15", + "websockets>=13,<17", "pyyaml<7.0.0,>=6.0.0", "httpx<1.0.0,>=0.27.0", "sentry-sdk>=1.0.0", diff --git a/sdk/src/beta9/clients/bot/__init__.py b/sdk/src/beta9/clients/bot/__init__.py index 2ec54553c..9da5979c5 100644 --- a/sdk/src/beta9/clients/bot/__init__.py +++ b/sdk/src/beta9/clients/bot/__init__.py @@ -16,7 +16,6 @@ from betterproto.grpcstub.grpcio_client import SyncServiceStub from betterproto.grpcstub.grpclib_server import ServiceBase - if TYPE_CHECKING: import grpclib.server from betterproto.grpcstub.grpclib_client import MetadataLike diff --git a/sdk/src/beta9/clients/endpoint/__init__.py b/sdk/src/beta9/clients/endpoint/__init__.py index 6b0ff8adc..79d14b86a 100644 --- a/sdk/src/beta9/clients/endpoint/__init__.py +++ b/sdk/src/beta9/clients/endpoint/__init__.py @@ -15,7 +15,6 @@ from betterproto.grpcstub.grpcio_client import SyncServiceStub from betterproto.grpcstub.grpclib_server import ServiceBase - if TYPE_CHECKING: import grpclib.server from betterproto.grpcstub.grpclib_client import MetadataLike diff --git a/sdk/src/beta9/clients/function/__init__.py b/sdk/src/beta9/clients/function/__init__.py index f2136dfeb..9881f9e3d 100644 --- a/sdk/src/beta9/clients/function/__init__.py +++ b/sdk/src/beta9/clients/function/__init__.py @@ -17,7 +17,6 @@ from betterproto.grpcstub.grpcio_client import SyncServiceStub from betterproto.grpcstub.grpclib_server import ServiceBase - if TYPE_CHECKING: import grpclib.server from betterproto.grpcstub.grpclib_client import MetadataLike diff --git a/sdk/src/beta9/clients/gateway/__init__.py b/sdk/src/beta9/clients/gateway/__init__.py index f3332f32d..38baa0ce7 100644 --- a/sdk/src/beta9/clients/gateway/__init__.py +++ b/sdk/src/beta9/clients/gateway/__init__.py @@ -24,7 +24,6 @@ from .. import types as _types__ - if TYPE_CHECKING: import grpclib.server from betterproto.grpcstub.grpclib_client import MetadataLike diff --git a/sdk/src/beta9/clients/image/__init__.py b/sdk/src/beta9/clients/image/__init__.py index 93ca79bb2..86f182a1e 100644 --- a/sdk/src/beta9/clients/image/__init__.py +++ b/sdk/src/beta9/clients/image/__init__.py @@ -18,7 +18,6 @@ from betterproto.grpcstub.grpcio_client import SyncServiceStub from betterproto.grpcstub.grpclib_server import ServiceBase - if TYPE_CHECKING: import grpclib.server from betterproto.grpcstub.grpclib_client import MetadataLike diff --git a/sdk/src/beta9/clients/map/__init__.py b/sdk/src/beta9/clients/map/__init__.py index 533529307..772d16ad2 100644 --- a/sdk/src/beta9/clients/map/__init__.py +++ b/sdk/src/beta9/clients/map/__init__.py @@ -16,7 +16,6 @@ from betterproto.grpcstub.grpcio_client import SyncServiceStub from betterproto.grpcstub.grpclib_server import ServiceBase - if TYPE_CHECKING: import grpclib.server from betterproto.grpcstub.grpclib_client import MetadataLike diff --git a/sdk/src/beta9/clients/output/__init__.py b/sdk/src/beta9/clients/output/__init__.py index df39183f4..2069714d4 100644 --- a/sdk/src/beta9/clients/output/__init__.py +++ b/sdk/src/beta9/clients/output/__init__.py @@ -21,7 +21,6 @@ from betterproto.grpcstub.grpcio_client import SyncServiceStub from betterproto.grpcstub.grpclib_server import ServiceBase - if TYPE_CHECKING: import grpclib.server from betterproto.grpcstub.grpclib_client import MetadataLike diff --git a/sdk/src/beta9/clients/pod/__init__.py b/sdk/src/beta9/clients/pod/__init__.py index 8740aa660..a39e0da14 100644 --- a/sdk/src/beta9/clients/pod/__init__.py +++ b/sdk/src/beta9/clients/pod/__init__.py @@ -18,7 +18,6 @@ from .. import types as _types__ - if TYPE_CHECKING: import grpclib.server from betterproto.grpcstub.grpclib_client import MetadataLike diff --git a/sdk/src/beta9/clients/secret/__init__.py b/sdk/src/beta9/clients/secret/__init__.py index 5820d888b..0a433ba7e 100644 --- a/sdk/src/beta9/clients/secret/__init__.py +++ b/sdk/src/beta9/clients/secret/__init__.py @@ -17,7 +17,6 @@ from betterproto.grpcstub.grpcio_client import SyncServiceStub from betterproto.grpcstub.grpclib_server import ServiceBase - if TYPE_CHECKING: import grpclib.server from betterproto.grpcstub.grpclib_client import MetadataLike diff --git a/sdk/src/beta9/clients/shell/__init__.py b/sdk/src/beta9/clients/shell/__init__.py index 06bdd17f4..fca947c57 100644 --- a/sdk/src/beta9/clients/shell/__init__.py +++ b/sdk/src/beta9/clients/shell/__init__.py @@ -15,7 +15,6 @@ from betterproto.grpcstub.grpcio_client import SyncServiceStub from betterproto.grpcstub.grpclib_server import ServiceBase - if TYPE_CHECKING: import grpclib.server from betterproto.grpcstub.grpclib_client import MetadataLike diff --git a/sdk/src/beta9/clients/signal/__init__.py b/sdk/src/beta9/clients/signal/__init__.py index fbde48d57..3fa122a79 100644 --- a/sdk/src/beta9/clients/signal/__init__.py +++ b/sdk/src/beta9/clients/signal/__init__.py @@ -17,7 +17,6 @@ from betterproto.grpcstub.grpcio_client import SyncServiceStub from betterproto.grpcstub.grpclib_server import ServiceBase - if TYPE_CHECKING: import grpclib.server from betterproto.grpcstub.grpclib_client import MetadataLike diff --git a/sdk/src/beta9/clients/simplequeue/__init__.py b/sdk/src/beta9/clients/simplequeue/__init__.py index c009fcb1b..6a755082d 100644 --- a/sdk/src/beta9/clients/simplequeue/__init__.py +++ b/sdk/src/beta9/clients/simplequeue/__init__.py @@ -15,7 +15,6 @@ from betterproto.grpcstub.grpcio_client import SyncServiceStub from betterproto.grpcstub.grpclib_server import ServiceBase - if TYPE_CHECKING: import grpclib.server from betterproto.grpcstub.grpclib_client import MetadataLike diff --git a/sdk/src/beta9/clients/taskqueue/__init__.py b/sdk/src/beta9/clients/taskqueue/__init__.py index a3f739988..6e01393a8 100644 --- a/sdk/src/beta9/clients/taskqueue/__init__.py +++ b/sdk/src/beta9/clients/taskqueue/__init__.py @@ -17,7 +17,6 @@ from betterproto.grpcstub.grpcio_client import SyncServiceStub from betterproto.grpcstub.grpclib_server import ServiceBase - if TYPE_CHECKING: import grpclib.server from betterproto.grpcstub.grpclib_client import MetadataLike diff --git a/sdk/src/beta9/clients/volume/__init__.py b/sdk/src/beta9/clients/volume/__init__.py index 0f6f738a5..a68863403 100644 --- a/sdk/src/beta9/clients/volume/__init__.py +++ b/sdk/src/beta9/clients/volume/__init__.py @@ -22,7 +22,6 @@ from betterproto.grpcstub.grpcio_client import SyncServiceStub from betterproto.grpcstub.grpclib_server import ServiceBase - if TYPE_CHECKING: import grpclib.server from betterproto.grpcstub.grpclib_client import MetadataLike diff --git a/sdk/src/beta9/inference.py b/sdk/src/beta9/inference.py index af8efda36..b54abea19 100644 --- a/sdk/src/beta9/inference.py +++ b/sdk/src/beta9/inference.py @@ -280,8 +280,6 @@ def embed( Returns: EmbeddingResult with the embedding vector(s) """ - # Track whether caller passed a single string or a list (batch) - is_batch = not isinstance(input, str) if isinstance(input, str): input = [input] @@ -302,15 +300,15 @@ def embed( single = data.get("embedding", []) embeddings = [single] if single else [] - # Populate fields per contract: - # - `embedding` (singular): the single vector when caller passed a string - # (or first vector for back-compat when batch also returned one result) - # - `embeddings` (plural, list-of-lists): always the full batch result - single_vec: List[float] = embeddings[0] if embeddings else [] + # Return all embeddings for batch input. + # `embedding` keeps the legacy single-vector accessor (first vector); + # `embeddings` exposes the full batch. Previously the batch list was + # overloaded onto `embedding` and `embeddings` was left empty, so + # callers of result.embeddings got nothing back for batch input. return EmbeddingResult( model=model, - embedding=single_vec if not is_batch else (embeddings[0] if len(embeddings) == 1 else []), - embeddings=embeddings if is_batch else [], + embedding=embeddings[0] if embeddings else [], + embeddings=embeddings, usage={ "prompt_tokens": data.get("prompt_eval_count", 0), }, diff --git a/sdk/tests/test_cubic_fixes.py b/sdk/tests/test_cubic_fixes.py index 0666e1db5..d0ec57434 100644 --- a/sdk/tests/test_cubic_fixes.py +++ b/sdk/tests/test_cubic_fixes.py @@ -59,6 +59,7 @@ def test_batch_embeddings(self, mock_client_cls): inference.configure() result = inference.embed("model", "single input") + self.assertIsInstance(result, EmbeddingResult) self.assertEqual(len(result.embedding), 3) # Vector dimension is 3 self.assertIsInstance(result.embedding[0], float) @@ -80,6 +81,25 @@ def test_batch_embeddings(self, mock_client_cls): self.assertEqual(result.embeddings[0], [0.1, 0.1]) self.assertEqual(result.embeddings[2], [0.3, 0.3]) + @patch('beta9.inference.httpx.Client') + def test_chat_forwards_chatmessage(self, mock_client_cls): + """Test that chat() converts ChatMessage objects to dicts in the payload.""" + mock_client = MagicMock() + mock_client_cls.return_value = mock_client + mock_response = MagicMock() + mock_response.json.return_value = {"message": {"content": "hi"}, "done_reason": "stop"} + mock_response.status_code = 200 + mock_client.post.return_value = mock_response + + inference.configure() + inference.chat("model", [ChatMessage(role="user", content="hello")]) + + # ChatMessage objects must be serialized to {"role", "content"} dicts + call_args = mock_client.post.call_args + self.assertIsNotNone(call_args) + payload = call_args[1]['json'] + self.assertEqual(payload['messages'], [{"role": "user", "content": "hello"}]) + @patch('beta9.inference.httpx.Client') def test_configure_closes_client(self, mock_client_cls): """Test that configure() closes the previous client."""