Reject blank worker names and probe the real credential and venv in --check

2026-08-03 04:01:45 +03:00
parent 0be114c8da
commit 7d1aacb19f
7 changed files with 207 additions and 3 deletions
@@ -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.
@@ -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
+}
@@ -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")
+	}
+}
@@ -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 ""
+}
@@ -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)
 }

@@ -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,
@@ -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)
+		}
+	}
+}