Reject blank worker names and probe the real credential and venv in --check
This commit is contained in:
@@ -58,7 +58,17 @@ func main() {
|
|||||||
fmt.Println("check: no coordinator URL (pass --coordinator-url or set COORDINATOR_URL)")
|
fmt.Println("check: no coordinator URL (pass --coordinator-url or set COORDINATOR_URL)")
|
||||||
os.Exit(1)
|
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)
|
printCheck(report)
|
||||||
if !report.Coordinator.OK || !report.Python.OK || !report.Scimesh.OK {
|
if !report.Coordinator.OK || !report.Python.OK || !report.Scimesh.OK {
|
||||||
os.Exit(1)
|
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
|
// loadConfig prefers a --config file; environment variables override the file
|
||||||
// (see agent.ConfigFile.Config). Without a file, the plain environment path is
|
// (see agent.ConfigFile.Config). Without a file, the plain environment path is
|
||||||
// used exactly as before.
|
// used exactly as before.
|
||||||
|
|||||||
@@ -9,6 +9,8 @@ import (
|
|||||||
"runtime"
|
"runtime"
|
||||||
"strings"
|
"strings"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
|
guuid "github.com/google/uuid"
|
||||||
)
|
)
|
||||||
|
|
||||||
// CheckItem is one line of the preflight report the setup wizard shows.
|
// 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
|
// 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
|
// non-empty python overrides the interpreter probed for the scimesh package
|
||||||
// (the managed venv after a runtime install).
|
// (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 := CheckCoordinator(ctx, coordinatorURL, 15*time.Second)
|
||||||
|
report.Auth = CheckAuth(ctx, coordinatorURL, token, workerKey, userserviceURL)
|
||||||
var env CheckReport
|
var env CheckReport
|
||||||
if python != "" {
|
if python != "" {
|
||||||
env = CheckEnvironmentWithPython(ctx, python)
|
env = CheckEnvironmentWithPython(ctx, python)
|
||||||
@@ -124,3 +131,83 @@ var Version = "dev"
|
|||||||
|
|
||||||
// Platform is the host platform string shown on the wizard.
|
// Platform is the host platform string shown on the wizard.
|
||||||
func Platform() string { return runtime.GOOS + "/" + runtime.GOARCH }
|
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
|
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:
|
// After the runtime installer created the venv, probe that interpreter:
|
||||||
// checking the bare system python3 would keep reporting scimesh as
|
// checking the bare system python3 would keep reporting scimesh as
|
||||||
// missing even though the worker would run with the venv.
|
// 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)
|
writeJSON(w, http.StatusOK, report)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -1,6 +1,7 @@
|
|||||||
package domain
|
package domain
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
"strings"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
"github.com/google/uuid"
|
"github.com/google/uuid"
|
||||||
@@ -52,6 +53,10 @@ func NewWorker(name string, capabilities []string, now time.Time) (*Worker, erro
|
|||||||
if len(capabilities) == 0 {
|
if len(capabilities) == 0 {
|
||||||
return nil, ErrInvalidInput
|
return nil, ErrInvalidInput
|
||||||
}
|
}
|
||||||
|
name = strings.TrimSpace(name)
|
||||||
|
if name == "" {
|
||||||
|
return nil, ErrInvalidInput
|
||||||
|
}
|
||||||
return &Worker{
|
return &Worker{
|
||||||
ID: uuid.New(),
|
ID: uuid.New(),
|
||||||
Name: name,
|
Name: name,
|
||||||
|
|||||||
@@ -29,3 +29,11 @@ func TestNewWorkerRejectsNoCapabilities(t *testing.T) {
|
|||||||
t.Errorf("empty slice: err = %v, want ErrInvalidInput", err)
|
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)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user