From 7d1aacb19fb9e337c08bac7f7a41c3a0bbb51ee6 Mon Sep 17 00:00:00 2001 From: Emil Date: Mon, 3 Aug 2026 04:01:45 +0300 Subject: [PATCH] Reject blank worker names and probe the real credential and venv in --check --- coordinator/cmd/worker-agent/main.go | 21 ++++- coordinator/internal/agent/check.go | 89 +++++++++++++++++++- coordinator/internal/agent/check_test.go | 68 +++++++++++++++ coordinator/internal/agent/configfile.go | 17 ++++ coordinator/internal/agent/setupui/server.go | 2 +- coordinator/internal/domain/worker.go | 5 ++ coordinator/internal/domain/worker_test.go | 8 ++ 7 files changed, 207 insertions(+), 3 deletions(-) create mode 100644 coordinator/internal/agent/check_test.go diff --git a/coordinator/cmd/worker-agent/main.go b/coordinator/cmd/worker-agent/main.go index f2c15b3..06e8d07 100644 --- a/coordinator/cmd/worker-agent/main.go +++ b/coordinator/cmd/worker-agent/main.go @@ -58,7 +58,17 @@ func main() { fmt.Println("check: no coordinator URL (pass --coordinator-url or set COORDINATOR_URL)") os.Exit(1) } - report := agent.RunCheck(ctx, url, "") + // Probe the managed venv when the wizard has installed it: workloads + // run with that interpreter, so checking the bare system python3 + // would report a false negative. + checkPython, checkToken, checkKey, checkUsers := "", "", "", "" + if configPath := checkConfigPath(); configPath != "" { + checkPython = agent.VenvPython(configPath) + if config, err := agent.LoadConfigFile(configPath); err == nil { + checkToken, checkKey, checkUsers = config.Token, config.WorkerKey, config.UserserviceURL + } + } + report := agent.RunCheck(ctx, url, checkPython, checkToken, checkKey, checkUsers) printCheck(report) if !report.Coordinator.OK || !report.Python.OK || !report.Scimesh.OK { os.Exit(1) @@ -87,6 +97,15 @@ func main() { } } +// checkConfigPath resolves where the wizard's config would be, honouring +// SCIMESH_WORKER_CONFIG like the rest of the agent. +func checkConfigPath() string { + if env := os.Getenv("SCIMESH_WORKER_CONFIG"); env != "" { + return env + } + return agent.DefaultConfigPath() +} + // loadConfig prefers a --config file; environment variables override the file // (see agent.ConfigFile.Config). Without a file, the plain environment path is // used exactly as before. diff --git a/coordinator/internal/agent/check.go b/coordinator/internal/agent/check.go index 7ca74e4..3aa23d7 100644 --- a/coordinator/internal/agent/check.go +++ b/coordinator/internal/agent/check.go @@ -9,6 +9,8 @@ import ( "runtime" "strings" "time" + + guuid "github.com/google/uuid" ) // CheckItem is one line of the preflight report the setup wizard shows. @@ -105,8 +107,13 @@ func CheckEnvironmentWithPython(ctx context.Context, python string) CheckReport // is the body behind `worker-agent --check` and the wizard's test step. A // non-empty python overrides the interpreter probed for the scimesh package // (the managed venv after a runtime install). -func RunCheck(ctx context.Context, coordinatorURL, python string) CheckReport { +// RunCheck combines the coordinator probe, a credential probe and the local +// environment probe; it is the body behind `worker-agent --check` and the +// wizard's test step. A non-empty python overrides the interpreter probed for +// the scimesh package (the managed venv after a runtime install). +func RunCheck(ctx context.Context, coordinatorURL, python, token, workerKey, userserviceURL string) CheckReport { report := CheckCoordinator(ctx, coordinatorURL, 15*time.Second) + report.Auth = CheckAuth(ctx, coordinatorURL, token, workerKey, userserviceURL) var env CheckReport if python != "" { env = CheckEnvironmentWithPython(ctx, python) @@ -124,3 +131,83 @@ var Version = "dev" // Platform is the host platform string shown on the wizard. func Platform() string { return runtime.GOOS + "/" + runtime.GOARCH } + +// CheckAuth verifies the configured credential against the coordinator +// without mutating anything: with a worker key it first exchanges it at the +// userservice for a short-lived JWT, then it probes /tasks/claim with a +// throwaway worker id and no capabilities. A 401 anywhere means the +// credential was rejected; any other status proves it was accepted. +func CheckAuth(ctx context.Context, url, token, workerKey, userserviceURL string) CheckItem { + item := CheckItem{Name: "auth"} + if token == "" && workerKey == "" { + item.OK = true + item.Detail = "no credential configured — will be checked at registration" + return item + } + client := &http.Client{Timeout: 30 * time.Second} + if workerKey != "" && userserviceURL != "" { + payload, _ := json.Marshal(map[string]string{"key": workerKey}) + req, err := http.NewRequestWithContext(ctx, http.MethodPost, strings.TrimRight(userserviceURL, "/")+"/worker-tokens/exchange", strings.NewReader(string(payload))) + if err != nil { + item.OK = false + item.Detail = "invalid userservice URL" + return item + } + req.Header.Set("Content-Type", "application/json") + resp, err := client.Do(req) + if err != nil { + item.OK = false + item.Detail = "userservice unreachable: " + err.Error() + return item + } + defer func() { _ = resp.Body.Close() }() + if resp.StatusCode == http.StatusUnauthorized { + item.OK = false + item.Detail = "worker key rejected by the userservice" + return item + } + if resp.StatusCode != http.StatusOK { + item.OK = false + item.Detail = fmt.Sprintf("userservice exchange: HTTP %d", resp.StatusCode) + return item + } + var exchanged struct { + Token string `json:"token"` + } + if err := json.NewDecoder(resp.Body).Decode(&exchanged); err != nil || exchanged.Token == "" { + item.OK = false + item.Detail = "userservice exchange returned no token" + return item + } + token = exchanged.Token + } + if token == "" { + item.OK = false + item.Detail = "no usable credential after the key exchange" + return item + } + payload, _ := json.Marshal(map[string]any{"worker_id": guuid.NewString(), "capabilities": []string{}}) + req, err := http.NewRequestWithContext(ctx, http.MethodPost, strings.TrimRight(url, "/")+"/tasks/claim", strings.NewReader(string(payload))) + if err != nil { + item.OK = false + item.Detail = "invalid coordinator URL" + return item + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+token) + resp, err := client.Do(req) + if err != nil { + item.OK = false + item.Detail = "coordinator unreachable: " + err.Error() + return item + } + defer func() { _ = resp.Body.Close() }() + if resp.StatusCode == http.StatusUnauthorized { + item.OK = false + item.Detail = "token rejected by the coordinator" + return item + } + item.OK = true + item.Detail = "credential accepted" + return item +} diff --git a/coordinator/internal/agent/check_test.go b/coordinator/internal/agent/check_test.go new file mode 100644 index 0000000..8e427d1 --- /dev/null +++ b/coordinator/internal/agent/check_test.go @@ -0,0 +1,68 @@ +package agent + +import ( + "context" + "net/http" + "net/http/httptest" + "testing" +) + +func TestCheckAuthAcceptsToken(t *testing.T) { + stub := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/tasks/claim" && r.Header.Get("Authorization") == "Bearer good-token" { + w.WriteHeader(http.StatusNoContent) + return + } + w.WriteHeader(http.StatusUnauthorized) + })) + defer stub.Close() + + item := CheckAuth(context.Background(), stub.URL, "good-token", "", "") + if !item.OK { + t.Errorf("good token: %+v", item) + } + item = CheckAuth(context.Background(), stub.URL, "bad-token", "", "") + if item.OK { + t.Error("bad token must fail") + } +} + +func TestCheckAuthNoCredentialIsNotAFailure(t *testing.T) { + item := CheckAuth(context.Background(), "http://coord:8080", "", "", "") + if !item.OK { + t.Errorf("no credential must not fail the preflight: %+v", item) + } +} + +func TestCheckAuthWorkerKeyExchange(t *testing.T) { + var exchanged bool + users := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/worker-tokens/exchange" { + t.Errorf("unexpected userservice path %q", r.URL.Path) + } + exchanged = true + _, _ = w.Write([]byte(`{"token":"jwt-after-exchange"}`)) + })) + defer users.Close() + coord := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Header.Get("Authorization") == "Bearer jwt-after-exchange" { + w.WriteHeader(http.StatusNoContent) + return + } + w.WriteHeader(http.StatusUnauthorized) + })) + defer coord.Close() + + item := CheckAuth(context.Background(), coord.URL, "", "smk_key", users.URL) + if !item.OK || !exchanged { + t.Errorf("key exchange flow: %+v exchanged=%v", item, exchanged) + } + + rejected := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusUnauthorized) + })) + defer rejected.Close() + if item := CheckAuth(context.Background(), coord.URL, "", "smk_bad", rejected.URL); item.OK { + t.Error("rejected key must fail the preflight") + } +} diff --git a/coordinator/internal/agent/configfile.go b/coordinator/internal/agent/configfile.go index fc9e54d..1637013 100644 --- a/coordinator/internal/agent/configfile.go +++ b/coordinator/internal/agent/configfile.go @@ -132,3 +132,20 @@ func SaveConfigFile(path string, file ConfigFile) error { } return nil } + +// VenvPython returns the managed venv python next to the given config file, +// or "" when the runtime installer has not created one yet. Mirrors the +// wizard's probe so `worker-agent --check` and the preflight agree on what +// interpreter will actually execute workloads. +func VenvPython(configPath string) string { + dir := filepath.Dir(configPath) + for _, candidate := range []string{ + filepath.Join(dir, "venv", "bin", "python"), + filepath.Join(dir, "venv", "Scripts", "python.exe"), + } { + if info, err := os.Stat(candidate); err == nil && !info.IsDir() { + return candidate + } + } + return "" +} diff --git a/coordinator/internal/agent/setupui/server.go b/coordinator/internal/agent/setupui/server.go index 6b3a15d..cafacdd 100644 --- a/coordinator/internal/agent/setupui/server.go +++ b/coordinator/internal/agent/setupui/server.go @@ -414,7 +414,7 @@ func (s *Server) handleTest(w http.ResponseWriter, r *http.Request) { // After the runtime installer created the venv, probe that interpreter: // checking the bare system python3 would keep reporting scimesh as // missing even though the worker would run with the venv. - report := agent.RunCheck(r.Context(), url, s.venvPython()) + report := agent.RunCheck(r.Context(), url, s.venvPython(), req.Token, req.WorkerKey, req.UserserviceURL) writeJSON(w, http.StatusOK, report) } diff --git a/coordinator/internal/domain/worker.go b/coordinator/internal/domain/worker.go index 75ee27b..4e1edac 100644 --- a/coordinator/internal/domain/worker.go +++ b/coordinator/internal/domain/worker.go @@ -1,6 +1,7 @@ package domain import ( + "strings" "time" "github.com/google/uuid" @@ -52,6 +53,10 @@ func NewWorker(name string, capabilities []string, now time.Time) (*Worker, erro if len(capabilities) == 0 { return nil, ErrInvalidInput } + name = strings.TrimSpace(name) + if name == "" { + return nil, ErrInvalidInput + } return &Worker{ ID: uuid.New(), Name: name, diff --git a/coordinator/internal/domain/worker_test.go b/coordinator/internal/domain/worker_test.go index d39716e..3aa52e2 100644 --- a/coordinator/internal/domain/worker_test.go +++ b/coordinator/internal/domain/worker_test.go @@ -29,3 +29,11 @@ func TestNewWorkerRejectsNoCapabilities(t *testing.T) { t.Errorf("empty slice: err = %v, want ErrInvalidInput", err) } } + +func TestNewWorkerRejectsBlankName(t *testing.T) { + for _, name := range []string{"", " ", "\t\n"} { + if _, err := NewWorker(name, []string{"similarity-search"}, testNow); !errors.Is(err, ErrInvalidInput) { + t.Errorf("name %q: got %v, want ErrInvalidInput", name, err) + } + } +}