From 743ed13417205848c6ca7c5237d4cb4bc8049626 Mon Sep 17 00:00:00 2001 From: whiskeylotus <0xpentest@proton.me> Date: Thu, 29 Oct 2026 17:04:12 -0400 Subject: [PATCH 1/6] feat(auth): [AGI-6072] log in and refresh at the tenant's auth host Router tenants run CLI login on the router, not on their own pool. Login now asks the tenant (GET /auth/cli/config) where to sign in and runs session, redeem and refresh there. A 404 is a stack that predates the endpoint and keeps its own login; any other failure, or an answer that is neither the tenant nor a known router, fails closed. Auth calls no longer follow redirects, so a 307 cannot replay a redemption secret elsewhere. Credentials are stored per auth host, with a hosts.json binding each API host to its auth host; old credential files keyed by the tenant keep working. Refresh no longer saves the server's default org over the stored scope. Plan-Step: 1/5 Co-Authored-By: Claude Opus 5.5 --- cmd/cli/cmd/auth.go | 33 ++- cmd/cli/cmd/whoami.go | 7 +- internal/auth/auth_host_test.go | 358 ++++++++++++++++++++++++++++++++ internal/auth/hosts.go | 147 +++++++++++++ internal/auth/login.go | 168 +++++++++++---- internal/auth/login_test.go | 37 +++- 6 files changed, 691 insertions(+), 59 deletions(-) create mode 100644 internal/auth/auth_host_test.go create mode 100644 internal/auth/hosts.go diff --git a/cmd/cli/cmd/auth.go b/cmd/cli/cmd/auth.go index fd46259..ff6133e 100644 --- a/cmd/cli/cmd/auth.go +++ b/cmd/cli/cmd/auth.go @@ -69,11 +69,18 @@ var logoutCmd = &cobra.Command{ return err } + authHost, err := auth.AuthHostFor(logoutHost) + if err != nil { + return fmt.Errorf("logout failed: %w", err) + } if err := auth.Logout(logoutHost); err != nil { return fmt.Errorf("logout failed: %w", err) } fmt.Printf("Logged out from %s\n", logoutHost) + if auth.CredentialKey(authHost) != auth.CredentialKey(logoutHost) { + fmt.Printf("This removed the %s login, which other tenants signed in through it share.\n", authHost) + } return nil }, } @@ -96,14 +103,11 @@ var statusCmd = &cobra.Command{ fmt.Printf("Host: %s\n", cfg.Host) - creds, err := auth.LoadCredentials() - if err != nil { - fmt.Println("Status: not authenticated") - return nil + hostCreds, authHost, ok, err := auth.LookupCredentials(cfg.Host) + if authHost != "" { + fmt.Printf("Auth host: %s\n", authHost) } - - hostCreds, ok := creds[auth.CredentialKey(cfg.Host)] - if !ok { + if err != nil || !ok { fmt.Println("Status: not authenticated") return nil } @@ -156,15 +160,15 @@ var switchCmd = &cobra.Command{ if switchHost == "" { // List available hosts - creds, err := auth.LoadCredentials() - if err != nil || len(creds) == 0 { + known, err := auth.KnownHosts() + if err != nil || len(known) == 0 { fmt.Println("No configured hosts. Run 'nullify auth login --host .nullify.ai'") return nil } cfg, _ := auth.LoadConfig() fmt.Println("Configured hosts:") - for h := range creds { + for _, h := range known { marker := " " if cfg != nil && auth.CredentialKey(cfg.Host) == h { marker = "* " @@ -206,7 +210,14 @@ var configShowCmd = &cobra.Command{ return nil } - data, _ := json.MarshalIndent(cfg, "", " ") + shown := map[string]any{"host": cfg.Host} + if cfg.Host != "" { + if authHost, err := auth.AuthHostFor(cfg.Host); err == nil { + shown["auth_host"] = authHost + } + } + + data, _ := json.MarshalIndent(shown, "", " ") fmt.Println(string(data)) return nil }, diff --git a/cmd/cli/cmd/whoami.go b/cmd/cli/cmd/whoami.go index fc49cd3..6ca9bd3 100644 --- a/cmd/cli/cmd/whoami.go +++ b/cmd/cli/cmd/whoami.go @@ -24,11 +24,14 @@ var whoamiCmd = &cobra.Command{ "host": whoamiHost, } - creds, err := auth.LoadCredentials() + hostCreds, authHost, ok, err := auth.LookupCredentials(whoamiHost) + if authHost != "" { + info["auth_host"] = authHost + } if err != nil { info["authenticated"] = false info["error"] = "no credentials found" - } else if hostCreds, ok := creds[auth.CredentialKey(whoamiHost)]; !ok { + } else if !ok { info["authenticated"] = false info["error"] = fmt.Sprintf("no credentials for host %s", whoamiHost) } else { diff --git a/internal/auth/auth_host_test.go b/internal/auth/auth_host_test.go new file mode 100644 index 0000000..97f0866 --- /dev/null +++ b/internal/auth/auth_host_test.go @@ -0,0 +1,358 @@ +package auth + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strconv" + "strings" + "sync" + "testing" + "time" +) + +const ( + tenantHost = "acme.nullify.ai" + tenantAPI = "api.acme.nullify.ai" + routerAPI = "api.nullify.ai" +) + +// hostCall is one request the routed auth server saw: the API host the CLI +// addressed and the path it asked for. +type hostCall struct{ host, path string } + +// newRoutedAuthServer serves every auth host from one httptest server, with +// the addressed host as the first path segment, so a test can tell which +// host each call went to. It isolates the credential store like +// newAuthServer. +func newRoutedAuthServer(t *testing.T, h func(w http.ResponseWriter, r *http.Request, host, path string)) *[]hostCall { + t.Helper() + var ( + mu sync.Mutex + calls []hostCall + ) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + host, path, _ := strings.Cut(strings.TrimPrefix(r.URL.Path, "/"), "/") + path = "/" + path + mu.Lock() + calls = append(calls, hostCall{host, path}) + mu.Unlock() + h(w, r, host, path) + })) + oldBase, oldClient := authBaseURL, httpClient + authBaseURL = func(host string) string { return srv.URL + "/" + apiHost(host) } + httpClient = srv.Client() + t.Cleanup(func() { + authBaseURL, httpClient = oldBase, oldClient + srv.Close() + }) + t.Setenv("HOME", t.TempDir()) + return &calls +} + +func authConfig(w http.ResponseWriter, authHost string) { + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]string{"auth_host": authHost}) +} + +func TestDiscoverAuthHost(t *testing.T) { + tests := []struct { + name string + handler func(w http.ResponseWriter) + want string + wantErr string + }{ + {"router", func(w http.ResponseWriter) { authConfig(w, "api.nullify.ai") }, routerAPI, ""}, + {"dev router", func(w http.ResponseWriter) { authConfig(w, "api.router.dev.nullify.ai") }, "api.router.dev.nullify.ai", ""}, + {"router with port 443 and case", func(w http.ResponseWriter) { authConfig(w, "API.Nullify.AI:443") }, routerAPI, ""}, + {"self", func(w http.ResponseWriter) { authConfig(w, tenantAPI) }, tenantAPI, ""}, + {"legacy stack without the endpoint", func(w http.ResponseWriter) { http.NotFound(w, nil) }, tenantAPI, ""}, + {"server error fails closed", func(w http.ResponseWriter) { w.WriteHeader(http.StatusInternalServerError) }, "", "status 500"}, + {"forbidden fails closed", func(w http.ResponseWriter) { w.WriteHeader(http.StatusForbidden) }, "", "status 403"}, + {"foreign host fails closed", func(w http.ResponseWriter) { authConfig(w, "api.evil.example") }, "", "neither this tenant nor a Nullify router"}, + {"another tenant fails closed", func(w http.ResponseWriter) { authConfig(w, "api.other.nullify.ai") }, "", "neither this tenant nor a Nullify router"}, + {"router with another port fails closed", func(w http.ResponseWriter) { authConfig(w, "api.nullify.ai:8443") }, "", "neither this tenant nor a Nullify router"}, + {"empty answer fails closed", func(w http.ResponseWriter) { authConfig(w, "") }, "", "neither this tenant nor a Nullify router"}, + {"garbage body fails closed", func(w http.ResponseWriter) { _, _ = w.Write([]byte("")) }, "", "decode auth config"}, + {"redirect fails closed", func(w http.ResponseWriter) { + w.Header().Set("Location", "https://api.evil.example/auth/cli/config") + w.WriteHeader(http.StatusTemporaryRedirect) + }, "", "status 307"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + calls := newRoutedAuthServer(t, func(w http.ResponseWriter, _ *http.Request, _, _ string) { + tt.handler(w) + }) + + got, err := DiscoverAuthHost(context.Background(), tenantHost) + if tt.wantErr != "" { + if err == nil || !strings.Contains(err.Error(), tt.wantErr) { + t.Fatalf("err = %v, want one containing %q", err, tt.wantErr) + } + if got != "" { + t.Errorf("host = %q alongside an error, want none", got) + } + } else { + if err != nil { + t.Fatalf("DiscoverAuthHost: %v", err) + } + if got != tt.want { + t.Errorf("auth host = %q, want %q", got, tt.want) + } + } + if len(*calls) != 1 || (*calls)[0] != (hostCall{tenantAPI, "/auth/cli/config"}) { + t.Errorf("calls = %v, want exactly one config lookup on the tenant", *calls) + } + }) + } +} + +// fakeRouterLogin answers the tenant's config with the router and runs the +// session and redeem on the router. Anything else is a test failure. +func fakeRouterLogin(t *testing.T) func(w http.ResponseWriter, r *http.Request, host, path string) { + return func(w http.ResponseWriter, r *http.Request, host, path string) { + switch { + case host == tenantAPI && path == "/auth/cli/config": + authConfig(w, routerAPI) + case host == routerAPI && path == "/auth/cli/session": + var body struct{ Port int } + _ = json.NewDecoder(r.Body).Decode(&body) + _ = json.NewEncoder(w).Encode(cliSessionResponse{ + SessionID: "sid", + AuthURL: "http://127.0.0.1:" + strconv.Itoa(body.Port) + "/callback?session_id=sid&redemption_secret=" + strings.Repeat("ab", 32), + }) + case host == routerAPI && path == "/auth/cli/token": + var body cliCallback + _ = json.NewDecoder(r.Body).Decode(&body) + if body.SessionID != "sid" || body.RedemptionSecret != strings.Repeat("ab", 32) { + t.Errorf("redeem body = %+v, want the callback's session and secret", body) + } + _ = json.NewEncoder(w).Encode(cliTokenResponse{ + AccessToken: "router-access", + RefreshToken: "router-refresh", + ExpiresIn: 3600, + QueryParameters: map[string]string{"githubOwnerId": "1"}, + }) + default: + t.Errorf("unexpected call to %s%s", host, path) + w.WriteHeader(http.StatusTeapot) + } + } +} + +// The browser's job in a login is to hit the local callback; the test plays +// that part by fetching the auth URL the session returned. +func stubBrowser(t *testing.T) { + t.Helper() + old := openBrowser + openBrowser = func(url string) error { + go func() { + resp, err := http.Get(url) + if err != nil { + t.Errorf("callback: %v", err) + return + } + resp.Body.Close() + }() + return nil + } + t.Cleanup(func() { openBrowser = old }) +} + +func TestLoginRunsSessionAndRedeemOnTheAuthHost(t *testing.T) { + calls := newRoutedAuthServer(t, fakeRouterLogin(t)) + stubBrowser(t) + + if err := Login(context.Background(), tenantHost); err != nil { + t.Fatalf("Login: %v", err) + } + + want := []hostCall{ + {tenantAPI, "/auth/cli/config"}, + {routerAPI, "/auth/cli/session"}, + {routerAPI, "/auth/cli/token"}, + } + if len(*calls) != len(want) { + t.Fatalf("calls = %v, want %v", *calls, want) + } + for i := range want { + if (*calls)[i] != want[i] { + t.Errorf("call %d = %v, want %v", i, (*calls)[i], want[i]) + } + } + + creds, err := LoadCredentials() + if err != nil { + t.Fatalf("LoadCredentials: %v", err) + } + if _, ok := creds[CredentialKey(tenantHost)]; ok { + t.Error("router login was stored under the tenant's own key") + } + if got := creds[CredentialKey(routerAPI)]; got.AccessToken != "router-access" || got.RefreshToken != "router-refresh" { + t.Errorf("router credentials = %+v, want the redeemed tokens", got) + } + + authHost, err := AuthHostFor(tenantHost) + if err != nil || authHost != routerAPI { + t.Errorf("AuthHostFor = %q, %v; want %q", authHost, err, routerAPI) + } + cfg, err := LoadConfig() + if err != nil || cfg.Host != tenantHost { + t.Errorf("config = %+v, %v; want host %q", cfg, err, tenantHost) + } +} + +func TestLoginFailsClosedWhenDiscoveryFails(t *testing.T) { + calls := newRoutedAuthServer(t, func(w http.ResponseWriter, _ *http.Request, _, _ string) { + w.WriteHeader(http.StatusBadGateway) + }) + stubBrowser(t) + + err := Login(context.Background(), tenantHost) + if err == nil || !strings.Contains(err.Error(), "status 502") { + t.Fatalf("err = %v, want the discovery failure", err) + } + if len(*calls) != 1 { + t.Errorf("calls = %v, want only the config lookup: no session on any host", *calls) + } +} + +// bindRouter records a router login for the tenant, as Login would. +func bindRouter(t *testing.T, access string, expiresAt int64) { + t.Helper() + seed(t, routerAPI, access, "router-refresh", expiresAt) + err := UpdateHostConfig(tenantHost, func(hc HostConfig) HostConfig { + hc.AuthHost = routerAPI + return hc + }) + if err != nil { + t.Fatalf("bind: %v", err) + } +} + +func routedRefresh(t *testing.T) func(w http.ResponseWriter, r *http.Request, host, path string) { + return func(w http.ResponseWriter, r *http.Request, host, path string) { + if path != "/auth/refresh_token" { + t.Errorf("unexpected call to %s%s", host, path) + } + refreshHandler(w, r) + } +} + +func TestGetValidTokenRefreshesAtTheAuthHost(t *testing.T) { + calls := newRoutedAuthServer(t, routedRefresh(t)) + bindRouter(t, "expired", time.Now().Add(-time.Hour).Unix()) + + got, err := GetValidToken(context.Background(), tenantHost) + if err != nil { + t.Fatalf("GetValidToken: %v", err) + } + if got != "new-id-token" { + t.Errorf("token = %q, want new-id-token", got) + } + if len(*calls) != 1 || (*calls)[0].host != routerAPI { + t.Errorf("calls = %v, want one refresh on %s", *calls, routerAPI) + } + + creds, _ := LoadCredentials() + if creds[CredentialKey(routerAPI)].AccessToken != "new-id-token" { + t.Error("refreshed token was not saved under the router's key") + } + if _, ok := creds[CredentialKey(tenantHost)]; ok { + t.Error("refresh wrote a credential under the tenant's own key") + } +} + +func TestForceRefreshTokenRefreshesAtTheAuthHost(t *testing.T) { + calls := newRoutedAuthServer(t, routedRefresh(t)) + bindRouter(t, "rejected", time.Now().Add(time.Hour).Unix()) + + got, err := ForceRefreshToken(context.Background(), tenantHost) + if err != nil { + t.Fatalf("ForceRefreshToken: %v", err) + } + if got != "new-id-token" { + t.Errorf("token = %q, want new-id-token", got) + } + if len(*calls) != 1 || (*calls)[0].host != routerAPI { + t.Errorf("calls = %v, want one refresh on %s", *calls, routerAPI) + } +} + +// A credentials.json written before auth hosts existed has the tenant's +// login under the tenant's key and no hosts.json. It must keep working, and +// keep refreshing on the tenant. +func TestLegacyCredentialFileStillWorks(t *testing.T) { + calls := newRoutedAuthServer(t, routedRefresh(t)) + dir := filepath.Join(os.Getenv("HOME"), ".nullify") + if err := os.MkdirAll(dir, 0700); err != nil { + t.Fatal(err) + } + legacy := `{"api.acme.nullify.ai": {"access_token": "legacy", "refresh_token": "legacy-rt", "expires_at": 1,` + + ` "query_parameters": {"githubOwnerId": "7"}}}` + if err := os.WriteFile(filepath.Join(dir, "credentials.json"), []byte(legacy), 0600); err != nil { + t.Fatal(err) + } + + got, err := GetValidToken(context.Background(), tenantHost) + if err != nil { + t.Fatalf("GetValidToken: %v", err) + } + if got != "new-id-token" { + t.Errorf("token = %q, want the refreshed new-id-token", got) + } + if len(*calls) != 1 || (*calls)[0].host != tenantAPI { + t.Errorf("calls = %v, want one refresh on the tenant itself", *calls) + } + + hostCreds, authHost, ok, err := LookupCredentials(tenantHost) + if err != nil || !ok { + t.Fatalf("LookupCredentials: ok=%v err=%v", ok, err) + } + if CredentialKey(authHost) != CredentialKey(tenantHost) { + t.Errorf("auth host = %q, want the tenant itself", authHost) + } + if hostCreds.QueryParameters["githubOwnerId"] != "7" { + t.Errorf("legacy scope = %v, want it kept as stored", hostCreds.QueryParameters) + } +} + +func TestLogoutRemovesTheAuthHostLogin(t *testing.T) { + newRoutedAuthServer(t, func(http.ResponseWriter, *http.Request, string, string) {}) + bindRouter(t, "live", time.Now().Add(time.Hour).Unix()) + seed(t, tenantHost, "legacy", "", 0) + + if err := Logout(tenantHost); err != nil { + t.Fatalf("Logout: %v", err) + } + + creds, err := LoadCredentials() + if err != nil { + t.Fatalf("LoadCredentials: %v", err) + } + if len(creds) != 0 { + t.Errorf("credentials left after logout: %v", creds) + } + if _, err := GetValidToken(context.Background(), tenantHost); err == nil { + t.Error("GetValidToken succeeded after logout") + } +} + +func TestKnownHostsListsTenantsNotTheRouter(t *testing.T) { + newRoutedAuthServer(t, func(http.ResponseWriter, *http.Request, string, string) {}) + bindRouter(t, "live", time.Now().Add(time.Hour).Unix()) + seed(t, "legacy.nullify.ai", "tok", "", 0) + + got, err := KnownHosts() + if err != nil { + t.Fatalf("KnownHosts: %v", err) + } + want := []string{"acme.nullify.ai", "legacy.nullify.ai"} + if strings.Join(got, ",") != strings.Join(want, ",") { + t.Errorf("KnownHosts = %v, want %v", got, want) + } +} diff --git a/internal/auth/hosts.go b/internal/auth/hosts.go new file mode 100644 index 0000000..4008d4b --- /dev/null +++ b/internal/auth/hosts.go @@ -0,0 +1,147 @@ +package auth + +import ( + "encoding/json" + "errors" + "fmt" + "maps" + "os" + "path/filepath" + "slices" +) + +// HostConfig is what the CLI remembers about one API host. AuthHost is the +// host that ran its login and serves its refresh, which on a router tenant is +// the router rather than the tenant itself. +type HostConfig struct { + AuthHost string `json:"auth_host,omitempty"` +} + +// Hosts maps an API host, in CredentialKey form, to its HostConfig. It lives +// apart from credentials.json because it holds no secrets, and apart from +// config.json because callers rewrite that file whole. +type Hosts map[string]HostConfig + +func hostsPath() (string, error) { + dir, err := configDir() + if err != nil { + return "", err + } + return filepath.Join(dir, "hosts.json"), nil +} + +// LoadHosts returns the stored host bindings, empty when none were saved. +func LoadHosts() (Hosts, error) { + path, err := hostsPath() + if err != nil { + return nil, err + } + + data, err := os.ReadFile(path) + if errors.Is(err, os.ErrNotExist) { + return Hosts{}, nil + } + if err != nil { + return nil, err + } + + var hosts Hosts + if err := json.Unmarshal(data, &hosts); err != nil { + return nil, fmt.Errorf("read %s: %w", path, err) + } + if hosts == nil { + hosts = Hosts{} + } + return hosts, nil +} + +func saveHosts(hosts Hosts) error { + if err := ensureConfigDir(); err != nil { + return err + } + + path, err := hostsPath() + if err != nil { + return err + } + + data, err := json.MarshalIndent(hosts, "", " ") + if err != nil { + return err + } + + return atomicWriteFile(path, data, 0600) +} + +// UpdateHostConfig applies update to the stored config for host and saves it. +func UpdateHostConfig(host string, update func(HostConfig) HostConfig) error { + hosts, err := LoadHosts() + if err != nil { + return err + } + + next := maps.Clone(hosts) + next[CredentialKey(host)] = update(hosts[CredentialKey(host)]) + + return saveHosts(next) +} + +// AuthHostFor returns the host whose credentials serve host: the auth host +// recorded at login, or host itself for a login made before the CLI asked +// (those credentials are stored under the tenant's own key). +func AuthHostFor(host string) (string, error) { + hosts, err := LoadHosts() + if err != nil { + return "", err + } + if hc, ok := hosts[CredentialKey(host)]; ok && hc.AuthHost != "" { + return hc.AuthHost, nil + } + return host, nil +} + +// LookupCredentials returns the stored credentials that serve host and the +// auth host they belong to. ok is false when none are stored. +func LookupCredentials(host string) (hostCreds HostCredentials, authHost string, ok bool, err error) { + authHost, err = AuthHostFor(host) + if err != nil { + return HostCredentials{}, "", false, err + } + + creds, err := LoadCredentials() + if err != nil { + return HostCredentials{}, authHost, false, err + } + + hostCreds, ok = creds[CredentialKey(authHost)] + return hostCreds, authHost, ok, nil +} + +// KnownHosts lists the API hosts the user has signed in to, sorted. Router +// credentials are shared by the tenants bound to them, so the router's own key +// is listed only when no tenant is. +func KnownHosts() ([]string, error) { + hosts, err := LoadHosts() + if err != nil { + return nil, err + } + creds, err := LoadCredentials() + if err != nil && !errors.Is(err, os.ErrNotExist) { + return nil, err + } + + known := map[string]bool{} + for h := range creds { + known[h] = true + } + for h, hc := range hosts { + if authKey := CredentialKey(hc.AuthHost); authKey != h { + delete(known, authKey) + } + } + for h := range hosts { + known[h] = true + } + + return slices.Sorted(maps.Keys(known)), nil +} diff --git a/internal/auth/login.go b/internal/auth/login.go index 1d6cfb0..6d16f7e 100644 --- a/internal/auth/login.go +++ b/internal/auth/login.go @@ -22,6 +22,22 @@ var httpClient = &http.Client{Timeout: 30 * time.Second} // tests can point the auth calls at an httptest server. var authBaseURL = func(host string) string { return "https://" + apiHost(host) } +// routerAuthHosts are the router API hosts a tenant may hand its CLI login +// to. The CLI sends the user's login and refresh token wherever the tenant +// says, so any other answer is refused rather than followed. +var routerAuthHosts = map[string]bool{ + "api.nullify.ai": true, + "api.router.dev.nullify.ai": true, +} + +// authClient is httpClient without redirects. Session, redeem and refresh +// carry secrets, and a 307 would replay them to wherever it points. +func authClient() *http.Client { + c := *httpClient + c.CheckRedirect = func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse } + return &c +} + type cliSessionResponse struct { SessionID string `json:"session_id"` AuthURL string `json:"auth_url"` @@ -36,13 +52,21 @@ type cliTokenResponse struct { } func Login(ctx context.Context, host string) error { + authHost, err := DiscoverAuthHost(ctx, host) + if err != nil { + return fmt.Errorf("failed to find the login host: %w", err) + } + if authHost != apiHost(host) { + fmt.Printf("%s signs in through %s.\n", host, authHost) + } + listener, err := net.Listen("tcp", "127.0.0.1:0") if err != nil { return fmt.Errorf("failed to start local server: %w", err) } port := listener.Addr().(*net.TCPAddr).Port - sessionResp, err := createCLISession(ctx, host, port) + sessionResp, err := createCLISession(ctx, authHost, port) if err != nil { listener.Close() return fmt.Errorf("failed to create auth session: %w", err) @@ -64,7 +88,7 @@ func Login(ctx context.Context, host string) error { fmt.Printf("\nOpening browser to authenticate...\n") fmt.Printf("If the browser doesn't open, visit:\n %s\n\n", sessionResp.AuthURL) - if err := OpenBrowser(sessionResp.AuthURL); err != nil { + if err := openBrowser(sessionResp.AuthURL); err != nil { logger.L(ctx).Debug("could not open browser automatically", logger.Err(err)) fmt.Println("(Could not open browser automatically. Please open the URL above manually.)") } @@ -93,7 +117,7 @@ waitLoop: } } - tokenResp, err := fetchCLIToken(ctx, host, callback.SessionID, callback.RedemptionSecret) + tokenResp, err := fetchCLIToken(ctx, authHost, callback.SessionID, callback.RedemptionSecret) if err != nil { return fmt.Errorf("failed to fetch tokens: %w", err) } @@ -104,7 +128,7 @@ waitLoop: expiresAt := time.Now().Add(time.Duration(tokenResp.ExpiresIn) * time.Second).Unix() - err = SaveHostCredentials(host, HostCredentials{ + err = SaveHostCredentials(authHost, HostCredentials{ AccessToken: tokenResp.AccessToken, RefreshToken: tokenResp.RefreshToken, ExpiresAt: expiresAt, @@ -114,6 +138,14 @@ waitLoop: return fmt.Errorf("failed to save credentials: %w", err) } + err = UpdateHostConfig(host, func(hc HostConfig) HostConfig { + hc.AuthHost = authHost + return hc + }) + if err != nil { + return fmt.Errorf("failed to save host config: %w", err) + } + err = SaveConfig(&Config{Host: host}) if err != nil { return fmt.Errorf("failed to save config: %w", err) @@ -123,26 +155,43 @@ waitLoop: return nil } +// Logout removes the credentials that serve host. On a router tenant those +// are the router login, shared with every other tenant signed in through it. func Logout(host string) error { + authHost, err := AuthHostFor(host) + if err != nil { + return err + } + if err := DeleteHostCredentials(authHost); err != nil { + return err + } + // A login from before the auth host was recorded sits under host itself. return DeleteHostCredentials(host) } -func GetValidToken(ctx context.Context, host string) (string, error) { - creds, err := LoadCredentials() +// storedCredentials returns the credentials that serve the API host and the +// auth host that issued them, which is where they must be refreshed. +func storedCredentials(host string) (HostCredentials, string, error) { + hostCreds, authHost, ok, err := LookupCredentials(host) if err != nil { - return "", fmt.Errorf("not authenticated - run 'nullify auth login'") + return HostCredentials{}, "", fmt.Errorf("not authenticated - run 'nullify auth login'") } - - key := CredentialKey(host) - hostCreds, ok := creds[key] if !ok { - return "", fmt.Errorf("not authenticated for %s - run 'nullify auth login --host %s'", host, host) + return HostCredentials{}, "", fmt.Errorf("not authenticated for %s - run 'nullify auth login --host %s'", host, host) + } + return hostCreds, authHost, nil +} + +func GetValidToken(ctx context.Context, host string) (string, error) { + hostCreds, authHost, err := storedCredentials(host) + if err != nil { + return "", err } if hostCreds.ExpiresAt > 0 && time.Now().Unix() > hostCreds.ExpiresAt { if hostCreds.RefreshToken != "" { - logger.L(ctx).Debug("access token expired, attempting refresh") - refreshed, err := refreshToken(ctx, host, hostCreds.RefreshToken) + logger.L(ctx).Debug("access token expired, attempting refresh", logger.String("authHost", authHost)) + refreshed, err := refreshToken(ctx, authHost, hostCreds) if err != nil { return "", fmt.Errorf("token expired and refresh failed - run 'nullify auth login': %w", err) } @@ -164,21 +213,62 @@ var ErrNotRefreshable = errors.New("no refresh token stored") // the server rejects while the CLI still believes it is valid -- revocation, a // killed session, or a local clock running behind the server's. func ForceRefreshToken(ctx context.Context, host string) (string, error) { - creds, err := LoadCredentials() + hostCreds, authHost, err := storedCredentials(host) if err != nil { - return "", fmt.Errorf("not authenticated - run 'nullify auth login'") - } - - hostCreds, ok := creds[CredentialKey(host)] - if !ok { - return "", fmt.Errorf("not authenticated for %s - run 'nullify auth login --host %s'", host, host) + return "", err } if hostCreds.RefreshToken == "" { return "", ErrNotRefreshable } - logger.L(ctx).Debug("forcing access token refresh") - return refreshToken(ctx, host, hostCreds.RefreshToken) + logger.L(ctx).Debug("forcing access token refresh", logger.String("authHost", authHost)) + return refreshToken(ctx, authHost, hostCreds) +} + +// DiscoverAuthHost asks host where its CLI login runs (GET /auth/cli/config). +// A stack that predates the endpoint answers 404 and logs in on its own host. +// Any other failure, or an answer that is neither this tenant nor a known +// router, is an error: credentials never go to a host the CLI did not expect, +// and a failed discovery is never retried on another host. +func DiscoverAuthHost(ctx context.Context, host string) (string, error) { + own := apiHost(host) + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, authBaseURL(host)+"/auth/cli/config", nil) + if err != nil { + return "", err + } + + resp, err := authClient().Do(req) + if err != nil { + return "", err + } + defer resp.Body.Close() + + switch resp.StatusCode { + case http.StatusOK: + case http.StatusNotFound: + return own, nil + default: + body, _ := io.ReadAll(io.LimitReader(resp.Body, 1024)) + return "", fmt.Errorf("auth config request returned status %d: %s", resp.StatusCode, string(body)) + } + + var config struct { + AuthHost string `json:"auth_host"` + } + if err := json.NewDecoder(io.LimitReader(resp.Body, 4<<10)).Decode(&config); err != nil { + return "", fmt.Errorf("decode auth config: %w", err) + } + + authHost := strings.TrimSuffix(strings.ToLower(strings.TrimSpace(config.AuthHost)), ":443") + switch { + case strings.EqualFold(authHost, own): + return own, nil + case routerAuthHosts[authHost]: + return authHost, nil + default: + return "", fmt.Errorf("%s names %q as its login host, which is neither this tenant nor a Nullify router", own, config.AuthHost) + } } func createCLISession(ctx context.Context, host string, port int) (*cliSessionResponse, error) { @@ -195,7 +285,7 @@ func createCLISession(ctx context.Context, host string, port int) (*cliSessionRe } req.Header.Set("Content-Type", "application/json") - resp, err := httpClient.Do(req) + resp, err := authClient().Do(req) if err != nil { return nil, err } @@ -229,7 +319,7 @@ func fetchCLIToken(ctx context.Context, host string, sessionID string, redemptio } req.Header.Set("Content-Type", "application/json") - resp, err := httpClient.Do(req) + resp, err := authClient().Do(req) if err != nil { return nil, err } @@ -249,16 +339,19 @@ func fetchCLIToken(ctx context.Context, host string, sessionID string, redemptio return &tokenResp, nil } -func refreshToken(ctx context.Context, host string, refreshTok string) (string, error) { - refreshURL := authBaseURL(host) + "/auth/refresh_token" +// refreshToken exchanges stored's refresh token at authHost and saves the new +// access token over it. The org scope the server returns is not saved: it is +// the server's default (the lowest org id), not a choice the user made. +func refreshToken(ctx context.Context, authHost string, stored HostCredentials) (string, error) { + refreshURL := authBaseURL(authHost) + "/auth/refresh_token" req, err := http.NewRequestWithContext(ctx, http.MethodGet, refreshURL, nil) if err != nil { return "", err } - req.AddCookie(&http.Cookie{Name: "refresh_token", Value: refreshTok}) + req.AddCookie(&http.Cookie{Name: "refresh_token", Value: stored.RefreshToken}) - resp, err := httpClient.Do(req) + resp, err := authClient().Do(req) if err != nil { return "", err } @@ -269,9 +362,8 @@ func refreshToken(ctx context.Context, host string, refreshTok string) (string, } var result struct { - AccessToken string `json:"accessToken"` - ExpiresIn int `json:"expiresIn"` - QueryParameters map[string]string `json:"queryParameters"` + AccessToken string `json:"accessToken"` + ExpiresIn int `json:"expiresIn"` } err = json.NewDecoder(resp.Body).Decode(&result) @@ -297,14 +389,11 @@ func refreshToken(ctx context.Context, host string, refreshTok string) (string, return "", fmt.Errorf("refresh returned empty access token") } - expiresAt := time.Now().Add(time.Duration(result.ExpiresIn) * time.Second).Unix() + refreshed := stored + refreshed.AccessToken = result.AccessToken + refreshed.ExpiresAt = time.Now().Add(time.Duration(result.ExpiresIn) * time.Second).Unix() - err = SaveHostCredentials(host, HostCredentials{ - AccessToken: result.AccessToken, - RefreshToken: refreshTok, - ExpiresAt: expiresAt, - QueryParameters: result.QueryParameters, - }) + err = SaveHostCredentials(authHost, refreshed) if err != nil { logger.L(ctx).Warn("failed to save refreshed credentials", logger.Err(err)) } @@ -338,6 +427,9 @@ func apiHost(host string) string { return "api." + host } +// openBrowser is OpenBrowser, a var so tests can stand in for the user. +var openBrowser = OpenBrowser + // OpenBrowser opens the given URL in the user's default browser. func OpenBrowser(url string) error { var cmd *exec.Cmd diff --git a/internal/auth/login_test.go b/internal/auth/login_test.go index 9c5b3e2..3649742 100644 --- a/internal/auth/login_test.go +++ b/internal/auth/login_test.go @@ -46,7 +46,7 @@ func refreshHandler(w http.ResponseWriter, _ *http.Request) { func TestRefreshTokenReadsSetCookie(t *testing.T) { newAuthServer(t, refreshHandler) - got, err := refreshToken(context.Background(), "example.nullify.ai", "stored-refresh") + got, err := refreshToken(context.Background(), "example.nullify.ai", HostCredentials{RefreshToken: "stored-refresh"}) if err != nil { t.Fatalf("refreshToken: %v", err) } @@ -70,8 +70,29 @@ func TestRefreshTokenReadsSetCookie(t *testing.T) { if d := saved.ExpiresAt - time.Now().Unix(); d < 3500 || d > 3600 { t.Errorf("ExpiresAt delta = %ds, want ~3600", d) } - if saved.QueryParameters["githubOwnerId"] != "42" { - t.Errorf("saved QueryParameters = %v, want githubOwnerId=42", saved.QueryParameters) +} + +// The refresh response carries the server's default org (the lowest id). It +// is not the user's pick, so it must not land in the store or replace the +// scope already there. +func TestRefreshTokenLeavesStoredOrgAlone(t *testing.T) { + newAuthServer(t, refreshHandler) + stored := HostCredentials{ + AccessToken: "old", + RefreshToken: "rt", + QueryParameters: map[string]string{"githubOwnerId": "7"}, + } + + if _, err := refreshToken(context.Background(), "example.nullify.ai", stored); err != nil { + t.Fatalf("refreshToken: %v", err) + } + + creds, err := LoadCredentials() + if err != nil { + t.Fatalf("LoadCredentials: %v", err) + } + if got := creds[CredentialKey("example.nullify.ai")].QueryParameters; got["githubOwnerId"] != "7" { + t.Errorf("saved QueryParameters = %v, want the stored githubOwnerId=7", got) } } @@ -81,7 +102,7 @@ func TestRefreshTokenBodyTokenWins(t *testing.T) { _, _ = w.Write([]byte(`{"accessToken":"body-token","expiresIn":60}`)) }) - got, err := refreshToken(context.Background(), "example.nullify.ai", "rt") + got, err := refreshToken(context.Background(), "example.nullify.ai", HostCredentials{RefreshToken: "rt"}) if err != nil { t.Fatalf("refreshToken: %v", err) } @@ -98,7 +119,7 @@ func TestRefreshTokenIgnoresEmptyCookie(t *testing.T) { _, _ = w.Write([]byte(`{"expiresIn":3600}`)) }) - got, err := refreshToken(context.Background(), "example.nullify.ai", "rt") + got, err := refreshToken(context.Background(), "example.nullify.ai", HostCredentials{RefreshToken: "rt"}) if err != nil { t.Fatalf("refreshToken: %v", err) } @@ -119,7 +140,7 @@ func TestRefreshTokenExpiresFallback(t *testing.T) { _, _ = w.Write([]byte(`{}`)) }) - if _, err := refreshToken(context.Background(), "example.nullify.ai", "rt"); err != nil { + if _, err := refreshToken(context.Background(), "example.nullify.ai", HostCredentials{RefreshToken: "rt"}); err != nil { t.Fatalf("refreshToken: %v", err) } @@ -137,7 +158,7 @@ func TestRefreshTokenNoTokenAnywhere(t *testing.T) { _, _ = w.Write([]byte(`{"expiresIn":3600}`)) }) - _, err := refreshToken(context.Background(), "example.nullify.ai", "rt") + _, err := refreshToken(context.Background(), "example.nullify.ai", HostCredentials{RefreshToken: "rt"}) if err == nil { t.Fatal("want an error when neither body nor cookie carries a token") } @@ -148,7 +169,7 @@ func TestRefreshTokenNon200(t *testing.T) { w.WriteHeader(http.StatusUnauthorized) }) - _, err := refreshToken(context.Background(), "example.nullify.ai", "rt") + _, err := refreshToken(context.Background(), "example.nullify.ai", HostCredentials{RefreshToken: "rt"}) if err == nil { t.Fatal("want an error on a non-200 refresh") } From 97f81cb8a4b06b3be9b169482a2d582f17330b4d Mon Sep 17 00:00:00 2001 From: whiskeylotus <0xpentest@proton.me> Date: Thu, 29 Oct 2026 17:06:09 -0400 Subject: [PATCH 2/6] fix(client): [AGI-6072] send the bearer only to the API host Both client transports set the Authorization header on every request they carried, redirect hops included, so a 3xx from the API sent the user's token to wherever it pointed. The bearer now goes only on an https request to the client's own API host that the caller made; redirects, other hosts, other ports and plain http go out without it. A 401 from such a request no longer triggers a forced refresh. Plan-Step: 2/5 Co-Authored-By: Claude Opus 5.5 --- internal/client/refreshing_transport.go | 12 +- internal/client/refreshing_transport_test.go | 22 +-- internal/client/transport.go | 33 ++++- internal/client/transport_test.go | 148 +++++++++++++++++++ 4 files changed, 202 insertions(+), 13 deletions(-) create mode 100644 internal/client/transport_test.go diff --git a/internal/client/refreshing_transport.go b/internal/client/refreshing_transport.go index d69bcd6..915153a 100644 --- a/internal/client/refreshing_transport.go +++ b/internal/client/refreshing_transport.go @@ -195,11 +195,21 @@ func (t *refreshingAuthTransport) RoundTrip(req *http.Request) (*http.Response, r.Body = io.NopCloser(bytes.NewReader(bodyBytes)) r.ContentLength = int64(len(bodyBytes)) } - r.Header.Set("Authorization", "Bearer "+token) r.Header.Set("User-Agent", "Nullify-CLI/mcp") + if token == "" { + r.Header.Del("Authorization") + } else { + r.Header.Set("Authorization", "Bearer "+token) + } return t.transport.RoundTrip(r) } + // Off the API host, or on a redirect hop, the request goes out bare and a + // 401 there is not ours to refresh. + if !bearerAllowed(t.nullifyHost, req) { + return attempt("") + } + token := t.getToken(req.Context()) resp, err := attempt(token) if err != nil { diff --git a/internal/client/refreshing_transport_test.go b/internal/client/refreshing_transport_test.go index d301d77..7d9b264 100644 --- a/internal/client/refreshing_transport_test.go +++ b/internal/client/refreshing_transport_test.go @@ -52,7 +52,7 @@ func TestRefreshOn401RetriesWithNewToken(t *testing.T) { inner, _ := bearerRouter(map[string]int{"stale": 401, "fresh": 200}) tr := newRefreshingTransport("stale", tp, inner) - req, _ := http.NewRequest(http.MethodGet, "http://x", nil) + req, _ := http.NewRequest(http.MethodGet, "https://api.acme.nullify.ai/", nil) resp, err := tr.RoundTrip(req) if err != nil { t.Fatalf("unexpected err: %v", err) @@ -71,7 +71,7 @@ func TestNo401MeansNoRefresh(t *testing.T) { inner, calls := bearerRouter(map[string]int{"good": 200}) tr := newRefreshingTransport("good", tp, inner) - req, _ := http.NewRequest(http.MethodGet, "http://x", nil) + req, _ := http.NewRequest(http.MethodGet, "https://api.acme.nullify.ai/", nil) resp, err := tr.RoundTrip(req) if err != nil { t.Fatalf("unexpected err: %v", err) @@ -92,7 +92,7 @@ func TestRefreshFailureSurfaces401(t *testing.T) { inner, _ := bearerRouter(map[string]int{"stale": 401}) tr := newRefreshingTransport("stale", tp, inner) - req, _ := http.NewRequest(http.MethodGet, "http://x", nil) + req, _ := http.NewRequest(http.MethodGet, "https://api.acme.nullify.ai/", nil) resp, err := tr.RoundTrip(req) if err != nil { t.Fatalf("unexpected err: %v", err) @@ -118,7 +118,7 @@ func TestRefreshReplaysBodyOn401(t *testing.T) { }) tr := newRefreshingTransport("stale", tp, inner) - req, _ := http.NewRequest(http.MethodPost, "http://x", strings.NewReader("payload")) + req, _ := http.NewRequest(http.MethodPost, "https://api.acme.nullify.ai/", strings.NewReader("payload")) resp, err := tr.RoundTrip(req) if err != nil { t.Fatalf("unexpected err: %v", err) @@ -152,7 +152,7 @@ func TestForceRefreshIsStormSafe(t *testing.T) { for range n { go func() { defer wg.Done() - req, _ := http.NewRequest(http.MethodGet, "http://x", nil) + req, _ := http.NewRequest(http.MethodGet, "https://api.acme.nullify.ai/", nil) resp, err := tr.RoundTrip(req) if err == nil { resp.Body.Close() @@ -196,7 +196,7 @@ func TestRefreshOn401UsesRefreshProviderWhenTokenLooksValid(t *testing.T) { inner, _ := bearerRouter(map[string]int{"revoked": 401, "fresh": 200}) tr := newRefreshingTransportWithRefresher("revoked", tp, rp, inner) - req, _ := http.NewRequest(http.MethodGet, "http://x", nil) + req, _ := http.NewRequest(http.MethodGet, "https://api.acme.nullify.ai/", nil) resp, err := tr.RoundTrip(req) if err != nil { t.Fatalf("unexpected err: %v", err) @@ -219,7 +219,7 @@ func TestRefreshOn401GivesUpWhenNotRefreshable(t *testing.T) { inner, calls := bearerRouter(map[string]int{"fixed": 401}) tr := newRefreshingTransportWithRefresher("fixed", func() (string, error) { return "fixed", nil }, rp, inner) - req, _ := http.NewRequest(http.MethodGet, "http://x", nil) + req, _ := http.NewRequest(http.MethodGet, "https://api.acme.nullify.ai/", nil) resp, err := tr.RoundTrip(req) if err != nil { t.Fatalf("unexpected err: %v", err) @@ -243,7 +243,7 @@ func TestFailingProviderIsBackedOff(t *testing.T) { tr.cachedAt = time.Now().Add(-time.Hour) for i := 0; i < 5; i++ { - req, _ := http.NewRequest(http.MethodGet, "http://x", nil) + req, _ := http.NewRequest(http.MethodGet, "https://api.acme.nullify.ai/", nil) resp, err := tr.RoundTrip(req) if err != nil { t.Fatalf("RoundTrip %d: %v", i, err) @@ -268,7 +268,7 @@ func TestForceRefreshIsBackedOffAfterFailure(t *testing.T) { tr := newRefreshingTransportWithRefresher("revoked", func() (string, error) { return "revoked", nil }, rp, inner) for i := range 5 { - req, _ := http.NewRequest(http.MethodGet, "http://x", nil) + req, _ := http.NewRequest(http.MethodGet, "https://api.acme.nullify.ai/", nil) resp, err := tr.RoundTrip(req) if err != nil { t.Fatalf("RoundTrip %d: %v", i, err) @@ -297,7 +297,7 @@ func TestForceRefreshTreatsUnchangedTokenAsFailure(t *testing.T) { tr := newRefreshingTransportWithRefresher("revoked", func() (string, error) { return "revoked", nil }, rp, inner) for i := range 2 { - req, _ := http.NewRequest(http.MethodGet, "http://x", nil) + req, _ := http.NewRequest(http.MethodGet, "https://api.acme.nullify.ai/", nil) resp, err := tr.RoundTrip(req) if err != nil { t.Fatalf("RoundTrip %d: %v", i, err) @@ -323,7 +323,7 @@ func TestForceRefreshKeepsCacheOnEmptyToken(t *testing.T) { inner, _ := bearerRouter(map[string]int{"revoked": 401}) tr := newRefreshingTransportWithRefresher("revoked", func() (string, error) { return "revoked", nil }, rp, inner) - req, _ := http.NewRequest(http.MethodGet, "http://x", nil) + req, _ := http.NewRequest(http.MethodGet, "https://api.acme.nullify.ai/", nil) resp, err := tr.RoundTrip(req) if err != nil { t.Fatalf("unexpected err: %v", err) diff --git a/internal/client/transport.go b/internal/client/transport.go index 2f8b9d6..2da10e5 100644 --- a/internal/client/transport.go +++ b/internal/client/transport.go @@ -2,6 +2,8 @@ package client import ( "net/http" + "net/url" + "strings" "github.com/nullify-platform/cli/internal/logger" ) @@ -14,7 +16,36 @@ type authTransport struct { func (t *authTransport) RoundTrip(req *http.Request) (*http.Response, error) { r := req.Clone(req.Context()) - r.Header.Set("Authorization", "Bearer "+t.token) r.Header.Set("User-Agent", "Nullify-CLI/"+logger.Version) + if !bearerAllowed(t.nullifyHost, req) { + r.Header.Del("Authorization") + return t.transport.RoundTrip(r) + } + r.Header.Set("Authorization", "Bearer "+t.token) return t.transport.RoundTrip(r) } + +// bearerAllowed reports whether req may carry the Nullify bearer: an https +// request to the API host the client was built for, made by the caller rather +// than by following a redirect. Anything else goes out without it, so a +// redirect, or a URL the server handed back, never receives the user's token. +func bearerAllowed(nullifyHost string, req *http.Request) bool { + if req.Response != nil { + return false + } + return isAPIOrigin(nullifyHost, req.URL) +} + +func isAPIOrigin(nullifyHost string, u *url.URL) bool { + if u == nil || u.Scheme != "https" || u.User != nil { + return false + } + if port := u.Port(); port != "" && port != "443" { + return false + } + apiHost := strings.ToLower(nullifyHost) + if !strings.HasPrefix(apiHost, "api.") { + apiHost = "api." + apiHost + } + return strings.ToLower(u.Hostname()) == apiHost +} diff --git a/internal/client/transport_test.go b/internal/client/transport_test.go new file mode 100644 index 0000000..0d88be8 --- /dev/null +++ b/internal/client/transport_test.go @@ -0,0 +1,148 @@ +package client + +import ( + "net/http" + "sync" + "testing" + "time" +) + +// originRecorder plays both servers: the API host answers each path from +// redirects (a 302 to that Location) or 200, and every other host answers +// 200. It records the Authorization header each host received. +type originRecorder struct { + redirects map[string]string + mu sync.Mutex + seen []seenRequest +} + +type seenRequest struct{ url, auth string } + +func (o *originRecorder) RoundTrip(r *http.Request) (*http.Response, error) { + o.mu.Lock() + o.seen = append(o.seen, seenRequest{r.URL.String(), r.Header.Get("Authorization")}) + o.mu.Unlock() + if r.URL.Host == "api.acme.nullify.ai" { + if loc, ok := o.redirects[r.URL.Path]; ok { + resp := newResp(http.StatusFound, map[string]string{"Location": loc}) + resp.Request = r + return resp, nil + } + } + resp := newResp(http.StatusOK, nil) + resp.Request = r + return resp, nil +} + +func (o *originRecorder) authFor(url string) (string, bool) { + o.mu.Lock() + defer o.mu.Unlock() + for _, s := range o.seen { + if s.url == url { + return s.auth, true + } + } + return "", false +} + +func newOriginRecorder() *originRecorder { + return &originRecorder{redirects: map[string]string{ + "/to-foreign": "https://bucket.s3.amazonaws.com/report?X-Amz-Signature=x", + "/to-same": "https://api.acme.nullify.ai/landing", + "/to-plain": "http://api.acme.nullify.ai/plain", + "/to-userinfo": "https://user:pw@api.acme.nullify.ai/userinfo", + }} +} + +// clients returns one http.Client per transport, wired as the CLI wires +// them: the retry transport over the bearer transport over origin. +func clients(origin http.RoundTripper) map[string]*http.Client { + static := &authTransport{nullifyHost: "acme.nullify.ai", token: "secret", transport: origin} + refreshing := &refreshingAuthTransport{ + nullifyHost: "acme.nullify.ai", + tokenProvider: func() (string, error) { return "secret", nil }, + refreshProvider: func() (string, error) { return "refreshed", nil }, + transport: origin, + cachedToken: "secret", + cachedAt: time.Now(), + cacheTTL: time.Hour, + failureBackoff: time.Minute, + } + return map[string]*http.Client{ + "authTransport": {Transport: NewRetryTransport(static)}, + "refreshingAuthTransport": {Transport: NewRetryTransport(refreshing)}, + } +} + +func TestBearerGoesOnlyToTheAPIHost(t *testing.T) { + tests := []struct { + name string + get string + seenURL string + wantAuth string + }{ + {"same-host request carries the bearer", "https://api.acme.nullify.ai/findings", "https://api.acme.nullify.ai/findings", "Bearer secret"}, + {"explicit port 443 is the API host", "https://api.acme.nullify.ai:443/findings", "https://api.acme.nullify.ai:443/findings", "Bearer secret"}, + {"redirect to another host carries none", "https://api.acme.nullify.ai/to-foreign", "https://bucket.s3.amazonaws.com/report?X-Amz-Signature=x", ""}, + {"redirect back to the API host carries none", "https://api.acme.nullify.ai/to-same", "https://api.acme.nullify.ai/landing", ""}, + {"downgrade redirect carries none", "https://api.acme.nullify.ai/to-plain", "http://api.acme.nullify.ai/plain", ""}, + {"redirect with userinfo carries none", "https://api.acme.nullify.ai/to-userinfo", "https://user:pw@api.acme.nullify.ai/userinfo", ""}, + {"direct request to another host carries none", "https://bucket.s3.amazonaws.com/upload", "https://bucket.s3.amazonaws.com/upload", ""}, + {"lookalike host carries none", "https://api.acme.nullify.ai.evil.example/x", "https://api.acme.nullify.ai.evil.example/x", ""}, + {"other port carries none", "https://api.acme.nullify.ai:8443/x", "https://api.acme.nullify.ai:8443/x", ""}, + {"plain http carries none", "http://api.acme.nullify.ai/x", "http://api.acme.nullify.ai/x", ""}, + } + for _, tt := range tests { + for _, name := range []string{"authTransport", "refreshingAuthTransport"} { + t.Run(name+"/"+tt.name, func(t *testing.T) { + origin := newOriginRecorder() + client := clients(origin)[name] + + req, _ := http.NewRequest(http.MethodGet, tt.get, nil) + // A caller-set header must not survive off the API host either. + req.Header.Set("Authorization", "Bearer caller-set") + resp, err := client.Do(req) + if err != nil { + t.Fatalf("Do: %v", err) + } + resp.Body.Close() + + auth, ok := origin.authFor(tt.seenURL) + if !ok { + t.Fatalf("no request reached %s; saw %v", tt.seenURL, origin.seen) + } + if auth != tt.wantAuth { + t.Errorf("Authorization at %s = %q, want %q", tt.seenURL, auth, tt.wantAuth) + } + }) + } + } +} + +// A 401 from a host that never got the bearer says nothing about the token, +// so it must not burn a forced refresh. +func TestRefreshingTransportDoesNotRefreshForForeignHost(t *testing.T) { + refreshes := 0 + origin := rtFunc(func(r *http.Request) (*http.Response, error) { + if r.Header.Get("Authorization") != "" { + t.Errorf("bearer sent to %s", r.URL) + } + return newResp(http.StatusUnauthorized, nil), nil + }) + tr := newRefreshingTransportWithRefresher("secret", + func() (string, error) { return "secret", nil }, + func() (string, error) { refreshes++; return "refreshed", nil }, + origin) + + req, _ := http.NewRequest(http.MethodGet, "https://bucket.s3.amazonaws.com/x", nil) + resp, err := tr.RoundTrip(req) + if err != nil { + t.Fatalf("RoundTrip: %v", err) + } + if resp.StatusCode != http.StatusUnauthorized { + t.Errorf("status = %d, want the foreign 401 passed through", resp.StatusCode) + } + if refreshes != 0 { + t.Errorf("refreshes = %d, want 0", refreshes) + } +} From e77091ed696eca4cfce72d4b4e46790ffcaa3a89 Mon Sep 17 00:00:00 2001 From: whiskeylotus <0xpentest@proton.me> Date: Thu, 29 Oct 2026 17:34:34 -0400 Subject: [PATCH 3/6] feat(cli): [AGI-6072] pick the org a command runs as A router login covers every org the user has on a host, and with none named, multi-org requests fail with ErrNoTenantID. Old logins also saved the server's default org (the lowest id) as if the user had chosen it. Commands now run as the first of: --org (name or id), NULLIFY_ORG, NULLIFY_GITHUB_OWNER_ID, the pick saved by `nullify org use` if the login still covers it, the git remote's owner if it names one org, the only org. Otherwise a terminal prompts and remembers the answer, and anything else fails with the choices. The saved query_parameters are never a pick. The org list is the router's (or a legacy tenant's) /admin/installations, kept to installations whose nullifyBaseUrl is this host; orgs not routed yet and orgs on other hosts are reported apart. An API token names no org unless asked, so the server takes the token's own. `auth status`, `auth config` and `whoami` show the effective host, auth host and org. Also from the wave 1 review: store read errors are surfaced instead of "not authenticated", `auth switch` drops hosts whose login is gone, and a router login removes the tenant's previous login. Plan-Step: 3/5 Co-Authored-By: Claude Opus 5.5 --- cmd/cli/cmd/auth.go | 56 +++++- cmd/cli/cmd/org.go | 252 +++++++++++++++++++++++ cmd/cli/cmd/org_test.go | 275 +++++++++++++++++++++++++ cmd/cli/cmd/root.go | 2 + cmd/cli/cmd/runtime_auth.go | 43 ++-- cmd/cli/cmd/runtime_auth_test.go | 19 +- cmd/cli/cmd/whoami.go | 4 +- internal/auth/auth_host_test.go | 137 ++++++++++++- internal/auth/hosts.go | 16 +- internal/auth/login.go | 14 +- internal/auth/orgs.go | 336 +++++++++++++++++++++++++++++++ internal/auth/orgs_test.go | 203 +++++++++++++++++++ internal/lib/git.go | 64 +++++- internal/lib/git_test.go | 42 ++++ 14 files changed, 1408 insertions(+), 55 deletions(-) create mode 100644 cmd/cli/cmd/org.go create mode 100644 cmd/cli/cmd/org_test.go create mode 100644 internal/auth/orgs.go create mode 100644 internal/auth/orgs_test.go diff --git a/cmd/cli/cmd/auth.go b/cmd/cli/cmd/auth.go index ff6133e..29cdb63 100644 --- a/cmd/cli/cmd/auth.go +++ b/cmd/cli/cmd/auth.go @@ -1,6 +1,7 @@ package cmd import ( + "context" "encoding/json" "fmt" "os" @@ -62,7 +63,8 @@ var loginCmd = &cobra.Command{ var logoutCmd = &cobra.Command{ Use: "logout", Short: "Log out of Nullify", - Long: "Clear stored credentials for the current or specified host.", + Long: "Clear stored credentials for the current or specified host. On an instance that signs in through the " + + "Nullify router, this removes the router login, which every instance signed in through it shares.", RunE: func(cmd *cobra.Command, args []string) error { logoutHost, err := resolveHostForAuth() if err != nil { @@ -88,8 +90,22 @@ var logoutCmd = &cobra.Command{ var statusCmd = &cobra.Command{ Use: "status", Short: "Show authentication status", - Long: "Display the current authentication state including host, user, and token expiry.", + Long: "Display the current authentication state including host, auth host, organization, and token expiry.", RunE: func(cmd *cobra.Command, args []string) error { + ctx := cmd.Context() + + if usesMachineToken() { + statusHost, err := resolveHostE(ctx) + if err != nil { + fmt.Println("No host configured. Pass --host or set NULLIFY_HOST.") + return nil + } + fmt.Printf("Host: %s\n", statusHost) + fmt.Println("Auth: API token") + fmt.Printf("Org: %s\n", describeOrg(ctx, statusHost)) + return nil + } + cfg, err := auth.LoadConfig() if err != nil { fmt.Println("Not configured. Run 'nullify auth login --host .nullify.ai' to get started.") @@ -122,6 +138,7 @@ var statusCmd = &cobra.Command{ } else { fmt.Println("Status: authenticated") } + fmt.Printf("Org: %s\n", describeOrg(ctx, cfg.Host)) return nil }, } @@ -202,20 +219,25 @@ var switchCmd = &cobra.Command{ var configShowCmd = &cobra.Command{ Use: "config", Short: "Show current configuration", - Long: "Display the current CLI configuration as JSON.", + Long: "Display the effective host, how it authenticates (login and its auth host, or API token), and the organization commands run as, as JSON.", RunE: func(cmd *cobra.Command, args []string) error { - cfg, err := auth.LoadConfig() + ctx := cmd.Context() + configHost, err := resolveHostE(ctx) if err != nil { fmt.Fprintf(os.Stderr, "No config found. Run 'nullify init' to set up.\n") return nil } - shown := map[string]any{"host": cfg.Host} - if cfg.Host != "" { - if authHost, err := auth.AuthHostFor(cfg.Host); err == nil { + shown := map[string]any{"host": configHost} + if usesMachineToken() { + shown["auth"] = "api_token" + } else { + shown["auth"] = "login" + if authHost, err := auth.AuthHostFor(configHost); err == nil { shown["auth_host"] = authHost } } + shown["org"] = describeOrg(ctx, configHost) data, _ := json.MarshalIndent(shown, "", " ") fmt.Println(string(data)) @@ -223,6 +245,26 @@ var configShowCmd = &cobra.Command{ }, } +// describeOrg says which org commands on host run as right now, without +// prompting. +func describeOrg(ctx context.Context, host string) string { + token, err := lib.GetNullifyToken(ctx, host, nullifyToken, githubToken) + if err != nil { + return "unknown (not authenticated)" + } + scope, err := resolveOrgScope(ctx, host, token, orgInputsFor(host, orgFlag, "--org"), false) + switch { + case err != nil: + return "not chosen: " + err.Error() + case scope.Org != nil: + return scope.Org.String() + case scope.Query["githubOwnerId"] != "": + return "githubOwnerId " + scope.Query["githubOwnerId"] + " (NULLIFY_GITHUB_OWNER_ID)" + default: + return "the API token's own (the server derives it)" + } +} + func init() { rootCmd.AddCommand(authCmd) authCmd.AddCommand(loginCmd) diff --git a/cmd/cli/cmd/org.go b/cmd/cli/cmd/org.go new file mode 100644 index 0000000..6657507 --- /dev/null +++ b/cmd/cli/cmd/org.go @@ -0,0 +1,252 @@ +package cmd + +import ( + "bufio" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "os" + "strconv" + "strings" + + "github.com/nullify-platform/cli/internal/auth" + "github.com/nullify-platform/cli/internal/lib" + "github.com/nullify-platform/cli/internal/output" + "github.com/spf13/cobra" +) + +// listOrgs, promptIn and promptOut are vars so tests can stand in for the +// router and the terminal. +var ( + listOrgs = auth.ListOrgs + promptIn io.Reader = os.Stdin + promptOut io.Writer = os.Stderr +) + +// orgScope is the org a command runs as and the query parameters that say +// so. Org is nil when the scope is left to the server (an API token's own +// org) or given only as a raw id. +type orgScope struct { + Org *auth.Org + Query map[string]string +} + +// usesMachineToken reports whether the command authenticates with an API or +// GitHub Actions token rather than a browser login, in lib.GetNullifyToken's +// precedence. Such a token carries its own single org. +func usesMachineToken() bool { + return nullifyToken != "" || os.Getenv("NULLIFY_TOKEN") != "" || + (os.Getenv("GITHUB_ACTIONS") == "true" && githubToken != "" && os.Getenv("GITHUB_REPOSITORY") != "") +} + +// orgInputsFor gathers every way the current invocation can name its org on +// host. flag and source are --org, or a caller's equivalent. +func orgInputsFor(host, flag, source string) auth.OrgInputs { + in := auth.OrgInputs{ + Flag: flag, + FlagSource: source, + EnvOrg: os.Getenv("NULLIFY_ORG"), + EnvGitHubOwnerID: os.Getenv("NULLIFY_GITHUB_OWNER_ID"), + GitOwner: lib.DetectOwnerFromGit(), + } + if hosts, err := auth.LoadHosts(); err == nil { + in.Stored = hosts[auth.CredentialKey(host)].Org + } + return in +} + +// resolveOrgScope picks the org a command on host runs as (auth.ResolveOrg). +// An API token names its org only when asked to: with no explicit choice +// the request goes without one and the server takes the token's own. When +// the choice is ambiguous and prompt is set, the user is asked and the +// answer is remembered for host. +func resolveOrgScope(ctx context.Context, host, token string, in auth.OrgInputs, prompt bool) (orgScope, error) { + listHost, err := auth.AuthHostFor(host) + if err != nil { + return orgScope{}, err + } + + if usesMachineToken() { + if strings.TrimSpace(in.Flag) == "" && strings.TrimSpace(in.EnvOrg) == "" { + if id := strings.TrimSpace(in.EnvGitHubOwnerID); id != "" { + return orgScope{Query: map[string]string{"githubOwnerId": id}}, nil + } + return orgScope{Query: map[string]string{}}, nil + } + // The token is the tenant's, so its orgs are read from the tenant. + listHost = host + in = auth.OrgInputs{Flag: in.Flag, FlagSource: in.FlagSource, EnvOrg: in.EnvOrg} + } + + list, err := listOrgs(ctx, host, listHost, token) + if err != nil { + return orgScope{}, err + } + + org, err := auth.ResolveOrg(host, list, in) + var notChosen *auth.ErrOrgNotChosen + if errors.As(err, ¬Chosen) && prompt && !usesMachineToken() { + org, err = promptForOrg(host, notChosen.Orgs) + if err == nil { + if saveErr := rememberOrg(host, org); saveErr != nil { + fmt.Fprintf(promptOut, "Could not remember %s for %s: %v\n", org.Name, host, saveErr) + } else { + fmt.Fprintf(promptOut, "Using %s on %s from now on; change it with 'nullify org use'.\n", org.Name, host) + } + } + } + if err != nil { + return orgScope{}, err + } + return orgScope{Org: &org, Query: org.Query()}, nil +} + +func promptForOrg(host string, orgs []auth.Org) (auth.Org, error) { + fmt.Fprintf(promptOut, "%s serves several of your organizations:\n", host) + for i, o := range orgs { + fmt.Fprintf(promptOut, " %d) %s\n", i+1, o) + } + fmt.Fprint(promptOut, "Choose one (number or name): ") + + line, err := bufio.NewReader(promptIn).ReadString('\n') + if err != nil && line == "" { + return auth.Org{}, fmt.Errorf("no organization chosen: %w", err) + } + answer := strings.TrimSpace(line) + if n, convErr := strconv.Atoi(answer); convErr == nil && n >= 1 && n <= len(orgs) { + return orgs[n-1], nil + } + for _, o := range orgs { + if o.Matches(answer) { + return o, nil + } + } + return auth.Org{}, fmt.Errorf("%q is not one of the listed organizations", answer) +} + +func rememberOrg(host string, org auth.Org) error { + return auth.UpdateHostConfig(host, func(hc auth.HostConfig) auth.HostConfig { + hc.Org = &auth.OrgPick{Name: org.Name, Key: org.Key, ID: org.ID} + return hc + }) +} + +var orgCmd = &cobra.Command{ + Use: "org", + Short: "Choose the organization commands run as", + Long: "A login can cover several organizations on one Nullify host. Commands run as the org chosen by " + + "--org, NULLIFY_ORG, NULLIFY_GITHUB_OWNER_ID, 'nullify org use', the git remote's owner, or the only org, in that order.", +} + +var orgListCmd = &cobra.Command{ + Use: "list", + Short: "List the organizations you can use on this host", + RunE: func(cmd *cobra.Command, args []string) error { + ctx := cmd.Context() + authCtx, err := resolveCommandToken(ctx) + if err != nil { + return err + } + + out, err := orgListing(ctx, authCtx.Host, authCtx.Token, orgFlag) + if err != nil { + return err + } + + data, err := json.MarshalIndent(out, "", " ") + if err != nil { + return err + } + return output.Print(cmd, data) + }, +} + +var orgUseCmd = &cobra.Command{ + Use: "use ", + Short: "Remember the organization to use on this host", + Args: cobra.ExactArgs(1), + RunE: func(cmd *cobra.Command, args []string) error { + ctx := cmd.Context() + if usesMachineToken() { + return fmt.Errorf("an API token carries its own organization; 'nullify org use' applies to 'nullify auth login' sessions") + } + authCtx, err := resolveCommandToken(ctx) + if err != nil { + return err + } + + listHost, err := auth.AuthHostFor(authCtx.Host) + if err != nil { + return err + } + list, err := listOrgs(ctx, authCtx.Host, listHost, authCtx.Token) + if err != nil { + return err + } + org, err := auth.ResolveOrg(authCtx.Host, list, auth.OrgInputs{Flag: args[0], FlagSource: "org"}) + if err != nil { + return err + } + if err := rememberOrg(authCtx.Host, org); err != nil { + return fmt.Errorf("failed to save the organization: %w", err) + } + + fmt.Printf("Using %s on %s\n", org, authCtx.Host) + return nil + }, +} + +// orgListing describes the orgs the token can use on host and which one a +// call naming none would run as, for `nullify org list` and the MCP +// list_organizations tool. flag is --org. +func orgListing(ctx context.Context, host, token, flag string) (map[string]any, error) { + listHost := host + if !usesMachineToken() { + var err error + if listHost, err = auth.AuthHostFor(host); err != nil { + return nil, err + } + } + list, err := listOrgs(ctx, host, listHost, token) + if err != nil { + return nil, err + } + + out := map[string]any{ + "host": host, + "organizations": nonNil(list.Orgs), + } + if len(list.Pending) > 0 { + out["pending"] = list.Pending + } + if len(list.Elsewhere) > 0 { + out["elsewhere"] = list.Elsewhere + } + scope, err := resolveOrgScope(ctx, host, token, orgInputsFor(host, flag, "--org"), false) + switch { + case err != nil: + out["selected"] = nil + out["selection"] = err.Error() + case scope.Org != nil: + out["selected"] = scope.Org.Name + default: + out["selected"] = nil + out["selection"] = "the API token's own organization" + } + return out, nil +} + +func nonNil(orgs []auth.Org) []auth.Org { + if orgs == nil { + return []auth.Org{} + } + return orgs +} + +func init() { + rootCmd.AddCommand(orgCmd) + orgCmd.AddCommand(orgListCmd) + orgCmd.AddCommand(orgUseCmd) +} diff --git a/cmd/cli/cmd/org_test.go b/cmd/cli/cmd/org_test.go new file mode 100644 index 0000000..f165d96 --- /dev/null +++ b/cmd/cli/cmd/org_test.go @@ -0,0 +1,275 @@ +package cmd + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "io" + "os" + "strings" + "testing" + "time" + + "github.com/nullify-platform/cli/internal/auth" + "github.com/stretchr/testify/require" +) + +var ( + testOrgA = auth.Org{Name: "acme-web", Provider: "github", Key: "githubOwnerId", ID: "1001001", Host: "api.acme.nullify.ai"} + testOrgB = auth.Org{Name: "acme-payments", Provider: "github", Key: "githubOwnerId", ID: "2002002", Host: "api.acme.nullify.ai"} +) + +// stubOrgList replaces the org lister; list receives the host it was asked +// to list on. +func stubOrgList(t *testing.T, list func(listHost string) (auth.OrgList, error)) { + t.Helper() + old := listOrgs + listOrgs = func(_ context.Context, _, listHost, _ string) (auth.OrgList, error) { return list(listHost) } + t.Cleanup(func() { listOrgs = old }) +} + +// loginEnv sets up a browser login on acme bound to the router, with no +// token or org in the environment and a working directory outside any git +// repository. +func loginEnv(t *testing.T) { + t.Helper() + t.Setenv("HOME", t.TempDir()) + t.Chdir(t.TempDir()) + t.Setenv("NULLIFY_HOST", "acme.nullify.ai") + t.Setenv("NULLIFY_TOKEN", "") + t.Setenv("NULLIFY_ORG", "") + t.Setenv("NULLIFY_GITHUB_OWNER_ID", "") + t.Setenv("GITHUB_ACTIONS", "") + + require.NoError(t, auth.SaveHostCredentials("api.nullify.ai", auth.HostCredentials{ + AccessToken: "router-token", + RefreshToken: "rt", + ExpiresAt: time.Now().Add(time.Hour).Unix(), + QueryParameters: map[string]string{"githubOwnerId": "1001001"}, + })) + require.NoError(t, auth.UpdateHostConfig("acme.nullify.ai", func(hc auth.HostConfig) auth.HostConfig { + hc.AuthHost = "api.nullify.ai" + return hc + })) + + oldHost, oldToken, oldGithub, oldOrg := host, nullifyToken, githubToken, orgFlag + host, nullifyToken, githubToken, orgFlag = "", "", "", "" + t.Cleanup(func() { host, nullifyToken, githubToken, orgFlag = oldHost, oldToken, oldGithub, oldOrg }) +} + +func twoTestOrgs(t *testing.T) *[]string { + listed := &[]string{} + stubOrgList(t, func(listHost string) (auth.OrgList, error) { + *listed = append(*listed, listHost) + return auth.OrgList{Orgs: []auth.Org{testOrgA, testOrgB}}, nil + }) + return listed +} + +func TestResolveCommandAuthListsOrgsOnTheRouter(t *testing.T) { + loginEnv(t) + listed := twoTestOrgs(t) + orgFlag = "acme-payments" + + authCtx, err := resolveCommandAuth(context.Background()) + require.NoError(t, err) + require.Equal(t, "router-token", authCtx.Token) + require.Equal(t, map[string]string{"githubOwnerId": "2002002"}, authCtx.QueryParams) + require.Equal(t, []string{"api.nullify.ai"}, *listed) +} + +// Two orgs, nothing chosen, no terminal: the command fails with the choices +// instead of falling back to the saved scope or the lowest id. +func TestResolveCommandAuthAmbiguousWithoutATerminal(t *testing.T) { + loginEnv(t) + twoTestOrgs(t) + + _, err := resolveCommandAuth(context.Background()) + var notChosen *auth.ErrOrgNotChosen + require.ErrorAs(t, err, ¬Chosen) + require.Contains(t, err.Error(), "acme-web") + require.Contains(t, err.Error(), "acme-payments") +} + +func TestOrgEnvVarsBeatTheStoredPick(t *testing.T) { + loginEnv(t) + twoTestOrgs(t) + require.NoError(t, rememberOrg("acme.nullify.ai", testOrgA)) + + authCtx, err := resolveCommandAuth(context.Background()) + require.NoError(t, err) + require.Equal(t, "1001001", authCtx.QueryParams["githubOwnerId"], "stored pick applies with nothing else set") + + t.Setenv("NULLIFY_GITHUB_OWNER_ID", "2002002") + authCtx, err = resolveCommandAuth(context.Background()) + require.NoError(t, err) + require.Equal(t, "2002002", authCtx.QueryParams["githubOwnerId"], "NULLIFY_GITHUB_OWNER_ID beats the stored pick") + + t.Setenv("NULLIFY_GITHUB_OWNER_ID", "") + t.Setenv("NULLIFY_ORG", "acme-payments") + authCtx, err = resolveCommandAuth(context.Background()) + require.NoError(t, err) + require.Equal(t, "2002002", authCtx.QueryParams["githubOwnerId"], "NULLIFY_ORG beats the stored pick") + + orgFlag = "acme-web" + authCtx, err = resolveCommandAuth(context.Background()) + require.NoError(t, err) + require.Equal(t, "1001001", authCtx.QueryParams["githubOwnerId"], "--org beats NULLIFY_ORG") +} + +func TestPromptAnswerIsRemembered(t *testing.T) { + loginEnv(t) + twoTestOrgs(t) + var shown bytes.Buffer + oldIn, oldOut := promptIn, promptOut + promptIn, promptOut = strings.NewReader("2\n"), &shown + t.Cleanup(func() { promptIn, promptOut = oldIn, oldOut }) + + scope, err := resolveOrgScope(context.Background(), "acme.nullify.ai", "router-token", orgInputsFor("acme.nullify.ai", "", "--org"), true) + require.NoError(t, err) + require.Equal(t, testOrgB, *scope.Org) + require.Contains(t, shown.String(), "1) acme-web") + + hosts, err := auth.LoadHosts() + require.NoError(t, err) + require.Equal(t, &auth.OrgPick{Name: "acme-payments", Key: "githubOwnerId", ID: "2002002"}, hosts["acme.nullify.ai"].Org) + require.Equal(t, "api.nullify.ai", hosts["acme.nullify.ai"].AuthHost, "remembering the org keeps the auth host") + + // The next command, with no terminal, uses the remembered answer. + authCtx, err := resolveCommandAuth(context.Background()) + require.NoError(t, err) + require.Equal(t, "2002002", authCtx.QueryParams["githubOwnerId"]) +} + +func TestPromptRejectsAnUnlistedAnswer(t *testing.T) { + loginEnv(t) + twoTestOrgs(t) + oldIn, oldOut := promptIn, promptOut + promptIn, promptOut = strings.NewReader("someone-else\n"), &bytes.Buffer{} + t.Cleanup(func() { promptIn, promptOut = oldIn, oldOut }) + + _, err := resolveOrgScope(context.Background(), "acme.nullify.ai", "router-token", orgInputsFor("acme.nullify.ai", "", "--org"), true) + require.Error(t, err) + hosts, _ := auth.LoadHosts() + require.Nil(t, hosts["acme.nullify.ai"].Org) +} + +func TestStoredPickNoLongerAuthorizedFallsThrough(t *testing.T) { + loginEnv(t) + stubOrgList(t, func(string) (auth.OrgList, error) { + return auth.OrgList{Orgs: []auth.Org{testOrgB}}, nil + }) + require.NoError(t, rememberOrg("acme.nullify.ai", testOrgA)) + + authCtx, err := resolveCommandAuth(context.Background()) + require.NoError(t, err) + require.Equal(t, "2002002", authCtx.QueryParams["githubOwnerId"]) +} + +func TestAPITokenOrgSelection(t *testing.T) { + loginEnv(t) + t.Setenv("NULLIFY_TOKEN", "api-token") + listed := twoTestOrgs(t) + + authCtx, err := resolveCommandAuth(context.Background()) + require.NoError(t, err) + require.Empty(t, authCtx.QueryParams, "the server takes an API token's own org") + require.Empty(t, *listed, "no org list without an explicit org") + + t.Setenv("NULLIFY_GITHUB_OWNER_ID", "999") + authCtx, err = resolveCommandAuth(context.Background()) + require.NoError(t, err) + require.Equal(t, map[string]string{"githubOwnerId": "999"}, authCtx.QueryParams) + require.Empty(t, *listed) + + t.Setenv("NULLIFY_GITHUB_OWNER_ID", "") + orgFlag = "acme-web" + authCtx, err = resolveCommandAuth(context.Background()) + require.NoError(t, err) + require.Equal(t, "1001001", authCtx.QueryParams["githubOwnerId"]) + require.Equal(t, []string{"acme.nullify.ai"}, *listed, "an API token lists on its own tenant, never the router") +} + +func TestOrgUseRemembersThePick(t *testing.T) { + loginEnv(t) + twoTestOrgs(t) + + orgUseCmd.SetContext(context.Background()) + require.NoError(t, orgUseCmd.RunE(orgUseCmd, []string{"2002002"})) + + hosts, err := auth.LoadHosts() + require.NoError(t, err) + require.Equal(t, "acme-payments", hosts["acme.nullify.ai"].Org.Name) + + err = orgUseCmd.RunE(orgUseCmd, []string{"nope"}) + require.Error(t, err) + require.Equal(t, "acme-payments", mustLoadHosts(t)["acme.nullify.ai"].Org.Name, "a failed use keeps the old pick") +} + +func TestOrgListFailureSurfaces(t *testing.T) { + loginEnv(t) + stubOrgList(t, func(string) (auth.OrgList, error) { return auth.OrgList{}, errors.New("router down") }) + + _, err := resolveCommandAuth(context.Background()) + require.ErrorContains(t, err, "router down") +} + +func mustLoadHosts(t *testing.T) auth.Hosts { + t.Helper() + hosts, err := auth.LoadHosts() + require.NoError(t, err) + return hosts +} + +func captureStdout(t *testing.T, run func()) string { + t.Helper() + r, w, err := os.Pipe() + require.NoError(t, err) + old := os.Stdout + os.Stdout = w + defer func() { os.Stdout = old }() + run() + require.NoError(t, w.Close()) + out, err := io.ReadAll(r) + require.NoError(t, err) + return string(out) +} + +func TestAuthConfigShowsHostAuthHostAndOrg(t *testing.T) { + loginEnv(t) + twoTestOrgs(t) + require.NoError(t, rememberOrg("acme.nullify.ai", testOrgB)) + + configShowCmd.SetContext(context.Background()) + out := captureStdout(t, func() { require.NoError(t, configShowCmd.RunE(configShowCmd, nil)) }) + var shown map[string]string + require.NoError(t, json.Unmarshal([]byte(out), &shown), out) + require.Equal(t, map[string]string{ + "host": "acme.nullify.ai", + "auth": "login", + "auth_host": "api.nullify.ai", + "org": "acme-payments (githubOwnerId 2002002)", + }, shown) + + t.Setenv("NULLIFY_TOKEN", "api-token") + out = captureStdout(t, func() { require.NoError(t, configShowCmd.RunE(configShowCmd, nil)) }) + shown = nil + require.NoError(t, json.Unmarshal([]byte(out), &shown), out) + require.Equal(t, "api_token", shown["auth"]) + require.Equal(t, "the API token's own (the server derives it)", shown["org"]) + require.NotContains(t, shown, "auth_host") +} + +func TestAuthStatusShowsTheOrg(t *testing.T) { + loginEnv(t) + twoTestOrgs(t) + require.NoError(t, auth.SaveConfig(&auth.Config{Host: "acme.nullify.ai"})) + + statusCmd.SetContext(context.Background()) + out := captureStdout(t, func() { require.NoError(t, statusCmd.RunE(statusCmd, nil)) }) + require.Contains(t, out, "Auth host: api.nullify.ai") + require.Contains(t, out, "Status: authenticated") + require.Contains(t, out, "Org: not chosen") + require.Contains(t, out, "acme-payments") +} diff --git a/cmd/cli/cmd/root.go b/cmd/cli/cmd/root.go index b931447..240178b 100644 --- a/cmd/cli/cmd/root.go +++ b/cmd/cli/cmd/root.go @@ -25,6 +25,7 @@ type commandOutput struct { var ( host string + orgFlag string verbose bool debug bool quiet bool @@ -121,6 +122,7 @@ func resetCommandOutput( func init() { rootCmd.PersistentFlags().StringVar(&host, "host", "", "The base URL of your Nullify API instance (e.g., acme.nullify.ai)") + rootCmd.PersistentFlags().StringVar(&orgFlag, "org", "", "Organization to run as, by name or id, when your login covers several (see 'nullify org list')") rootCmd.PersistentFlags().BoolVarP(&verbose, "verbose", "v", false, "Enable verbose logging") rootCmd.PersistentFlags().BoolVarP(&debug, "debug", "d", false, "Enable debug logging") rootCmd.PersistentFlags().StringVarP(&outputFmt, "output", "o", "json", "Output format (json, table, yaml, sarif)") diff --git a/cmd/cli/cmd/runtime_auth.go b/cmd/cli/cmd/runtime_auth.go index 33d2d0b..b2947b1 100644 --- a/cmd/cli/cmd/runtime_auth.go +++ b/cmd/cli/cmd/runtime_auth.go @@ -20,6 +20,9 @@ type commandAuthContext struct { Host string Token string QueryParams map[string]string + // Org is the org the command runs as, nil when the server takes it from + // an API token or it was given only as an id. + Org *auth.Org } // Client builds a NullifyClient for the resolved host and token. @@ -31,41 +34,43 @@ func (c *commandAuthContext) APIClient() *api.Client { return api.NewClient(c.Host, c.Token, c.QueryParams) } +// resolveCommandAuth resolves the host, token and org a command runs with. +// Stored credentials no longer carry the org: a login's saved +// query_parameters were the server's default, not the user's choice. func resolveCommandAuth(ctx context.Context) (*commandAuthContext, error) { - commandHost, err := resolveHostE(ctx) + authCtx, err := resolveCommandToken(ctx) if err != nil { return nil, err } - token, err := lib.GetNullifyToken(ctx, commandHost, nullifyToken, githubToken) + scope, err := resolveOrgScope(ctx, authCtx.Host, authCtx.Token, orgInputsFor(authCtx.Host, orgFlag, "--org"), stdinIsTTY()) if err != nil { - if errors.Is(err, lib.ErrNoToken) { - return nil, authError("not authenticated. Run 'nullify auth login' first") - } return nil, authError("%w", err) } + authCtx.Org = scope.Org + authCtx.QueryParams = scope.Query + return authCtx, nil +} - queryParams := map[string]string{} - creds, err := auth.LoadCredentials() - if err == nil { - if hostCreds, ok := creds[auth.CredentialKey(commandHost)]; ok && hostCreds.QueryParameters != nil { - for key, value := range hostCreds.QueryParameters { - queryParams[key] = value - } - } +// resolveCommandToken resolves the host and token only, for callers that +// choose the org themselves. +func resolveCommandToken(ctx context.Context) (*commandAuthContext, error) { + commandHost, err := resolveHostE(ctx) + if err != nil { + return nil, err } - // NULLIFY_GITHUB_OWNER_ID env var provides the githubOwnerId query param - // when using NULLIFY_TOKEN directly (no stored credentials). - if ownerID := os.Getenv("NULLIFY_GITHUB_OWNER_ID"); ownerID != "" { - if _, set := queryParams["githubOwnerId"]; !set { - queryParams["githubOwnerId"] = ownerID + token, err := lib.GetNullifyToken(ctx, commandHost, nullifyToken, githubToken) + if err != nil { + if errors.Is(err, lib.ErrNoToken) { + return nil, authError("not authenticated. Run 'nullify auth login' first") } + return nil, authError("%w", err) } return &commandAuthContext{ Host: commandHost, Token: token, - QueryParams: queryParams, + QueryParams: map[string]string{}, }, nil } diff --git a/cmd/cli/cmd/runtime_auth_test.go b/cmd/cli/cmd/runtime_auth_test.go index 3ff0ab2..2b7d760 100644 --- a/cmd/cli/cmd/runtime_auth_test.go +++ b/cmd/cli/cmd/runtime_auth_test.go @@ -33,10 +33,15 @@ func TestResolveCommandAuthUsesEnvTokenWithoutStoredCredentials(t *testing.T) { require.Empty(t, authCtx.QueryParams) } -func TestResolveCommandAuthClonesStoredQueryParams(t *testing.T) { +// Logins before this CLI saved the server's default org (the lowest id) as +// query_parameters. That was never the user's choice, so it must not scope +// requests: an API token leaves the org to the server instead. +func TestResolveCommandAuthIgnoresLegacySavedScope(t *testing.T) { t.Setenv("HOME", t.TempDir()) t.Setenv("NULLIFY_HOST", "acme.nullify.ai") t.Setenv("NULLIFY_TOKEN", "env-token") + t.Setenv("NULLIFY_ORG", "") + t.Setenv("NULLIFY_GITHUB_OWNER_ID", "") err := auth.SaveHostCredentials("acme.nullify.ai", auth.HostCredentials{ AccessToken: "stored-token", @@ -47,6 +52,10 @@ func TestResolveCommandAuthClonesStoredQueryParams(t *testing.T) { }, }) require.NoError(t, err) + stubOrgList(t, func(string) (auth.OrgList, error) { + t.Error("an API token with no explicit org must not list orgs") + return auth.OrgList{}, nil + }) originalHost := host originalNullifyToken := nullifyToken @@ -62,11 +71,5 @@ func TestResolveCommandAuthClonesStoredQueryParams(t *testing.T) { authCtx, err := resolveCommandAuth(context.Background()) require.NoError(t, err) - require.Equal(t, map[string]string{"githubOwnerId": "123"}, authCtx.QueryParams) - - authCtx.QueryParams["githubOwnerId"] = "456" - - creds, err := auth.LoadCredentials() - require.NoError(t, err) - require.Equal(t, "123", creds[auth.CredentialKey("acme.nullify.ai")].QueryParameters["githubOwnerId"]) + require.Empty(t, authCtx.QueryParams) } diff --git a/cmd/cli/cmd/whoami.go b/cmd/cli/cmd/whoami.go index 6ca9bd3..eb76252 100644 --- a/cmd/cli/cmd/whoami.go +++ b/cmd/cli/cmd/whoami.go @@ -36,9 +36,7 @@ var whoamiCmd = &cobra.Command{ info["error"] = fmt.Sprintf("no credentials for host %s", whoamiHost) } else { info["authenticated"] = true - if hostCreds.QueryParameters != nil { - info["query_parameters"] = hostCreds.QueryParameters - } + info["org"] = describeOrg(ctx, whoamiHost) if hostCreds.ExpiresAt > 0 { info["expires_at"] = hostCreds.ExpiresAt } diff --git a/internal/auth/auth_host_test.go b/internal/auth/auth_host_test.go index 97f0866..280a960 100644 --- a/internal/auth/auth_host_test.go +++ b/internal/auth/auth_host_test.go @@ -148,24 +148,32 @@ func fakeRouterLogin(t *testing.T) func(w http.ResponseWriter, r *http.Request, // that part by fetching the auth URL the session returned. func stubBrowser(t *testing.T) { t.Helper() + var wg sync.WaitGroup old := openBrowser openBrowser = func(url string) error { + wg.Add(1) go func() { - resp, err := http.Get(url) - if err != nil { - t.Errorf("callback: %v", err) - return + defer wg.Done() + // Login may close its callback server before this response is + // read; only a callback that never arrived fails the test, and + // it does so by Login timing out. + if resp, err := http.Get(url); err == nil { + resp.Body.Close() } - resp.Body.Close() }() return nil } - t.Cleanup(func() { openBrowser = old }) + t.Cleanup(func() { + wg.Wait() + openBrowser = old + }) } func TestLoginRunsSessionAndRedeemOnTheAuthHost(t *testing.T) { calls := newRoutedAuthServer(t, fakeRouterLogin(t)) stubBrowser(t) + // A login the tenant ran itself before it moved to the router. + seed(t, tenantHost, "tenant-pool-token", "tenant-pool-refresh", time.Now().Add(time.Hour).Unix()) if err := Login(context.Background(), tenantHost); err != nil { t.Fatalf("Login: %v", err) @@ -190,7 +198,7 @@ func TestLoginRunsSessionAndRedeemOnTheAuthHost(t *testing.T) { t.Fatalf("LoadCredentials: %v", err) } if _, ok := creds[CredentialKey(tenantHost)]; ok { - t.Error("router login was stored under the tenant's own key") + t.Error("the tenant-key login survived a router login") } if got := creds[CredentialKey(routerAPI)]; got.AccessToken != "router-access" || got.RefreshToken != "router-refresh" { t.Errorf("router credentials = %+v, want the redeemed tokens", got) @@ -342,6 +350,31 @@ func TestLogoutRemovesTheAuthHostLogin(t *testing.T) { } } +func TestKnownHostsDropsHostsAfterLogout(t *testing.T) { + newRoutedAuthServer(t, func(http.ResponseWriter, *http.Request, string, string) {}) + bindRouter(t, "live", time.Now().Add(time.Hour).Unix()) + if err := UpdateHostConfig("other.nullify.ai", func(hc HostConfig) HostConfig { + hc.AuthHost = routerAPI + return hc + }); err != nil { + t.Fatal(err) + } + + if err := Logout(tenantHost); err != nil { + t.Fatalf("Logout: %v", err) + } + got, err := KnownHosts() + if err != nil { + t.Fatalf("KnownHosts: %v", err) + } + if len(got) != 0 { + t.Errorf("KnownHosts = %v after the shared router login was removed, want none", got) + } + if _, err := GetValidToken(context.Background(), "other.nullify.ai"); err == nil { + t.Error("the other tenant on the same router login is still authenticated") + } +} + func TestKnownHostsListsTenantsNotTheRouter(t *testing.T) { newRoutedAuthServer(t, func(http.ResponseWriter, *http.Request, string, string) {}) bindRouter(t, "live", time.Now().Add(time.Hour).Unix()) @@ -356,3 +389,93 @@ func TestKnownHostsListsTenantsNotTheRouter(t *testing.T) { t.Errorf("KnownHosts = %v, want %v", got, want) } } + +// Redeem and refresh carry the redemption secret and the refresh cookie, so +// a redirect from the auth host must end the call, not be followed. +func TestAuthCallsDoNotFollowRedirects(t *testing.T) { + redirect := func(w http.ResponseWriter, r *http.Request, host, path string) { + if host == "api.evil.example" { + t.Errorf("followed a redirect to %s%s", host, path) + } + w.Header().Set("Location", "https://api.evil.example"+path) + w.WriteHeader(http.StatusTemporaryRedirect) + } + + t.Run("redeem", func(t *testing.T) { + calls := newRoutedAuthServer(t, redirect) + _, err := fetchCLIToken(context.Background(), routerAPI, "sid", strings.Repeat("ab", 32)) + if err == nil || !strings.Contains(err.Error(), "status 307") { + t.Fatalf("err = %v, want the 307 surfaced", err) + } + if len(*calls) != 1 { + t.Errorf("calls = %v, want one", *calls) + } + }) + t.Run("session", func(t *testing.T) { + calls := newRoutedAuthServer(t, redirect) + if _, err := createCLISession(context.Background(), routerAPI, 4242); err == nil { + t.Fatal("want an error on a redirected session") + } + if len(*calls) != 1 { + t.Errorf("calls = %v, want one", *calls) + } + }) + t.Run("refresh", func(t *testing.T) { + calls := newRoutedAuthServer(t, redirect) + if _, err := refreshToken(context.Background(), routerAPI, HostCredentials{RefreshToken: "rt"}); err == nil { + t.Fatal("want an error on a redirected refresh") + } + if len(*calls) != 1 { + t.Errorf("calls = %v, want one", *calls) + } + }) +} + +func TestLoginOnALegacyStackStaysOnTheTenant(t *testing.T) { + calls := newRoutedAuthServer(t, func(w http.ResponseWriter, r *http.Request, host, path string) { + if host != tenantAPI { + t.Errorf("call left the tenant: %s%s", host, path) + } + switch path { + case "/auth/cli/config": + http.NotFound(w, r) + case "/auth/cli/session": + var body struct{ Port int } + _ = json.NewDecoder(r.Body).Decode(&body) + _ = json.NewEncoder(w).Encode(cliSessionResponse{ + SessionID: "sid", + AuthURL: "http://127.0.0.1:" + strconv.Itoa(body.Port) + "/callback?session_id=sid&redemption_secret=" + strings.Repeat("ab", 32), + }) + case "/auth/cli/token": + _ = json.NewEncoder(w).Encode(cliTokenResponse{AccessToken: "tenant-access", RefreshToken: "tenant-refresh", ExpiresIn: 3600}) + } + }) + stubBrowser(t) + + if err := Login(context.Background(), tenantHost); err != nil { + t.Fatalf("Login: %v", err) + } + if len(*calls) != 3 { + t.Errorf("calls = %v, want config, session and redeem", *calls) + } + creds, _ := LoadCredentials() + if creds[CredentialKey(tenantHost)].AccessToken != "tenant-access" { + t.Errorf("credentials = %v, want the login under the tenant's own key", creds) + } + if authHost, _ := AuthHostFor(tenantHost); authHost != tenantAPI { + t.Errorf("auth host = %q, want the tenant", authHost) + } +} + +func TestCorruptHostsFileIsReportedNotTreatedAsLoggedOut(t *testing.T) { + newRoutedAuthServer(t, func(http.ResponseWriter, *http.Request, string, string) {}) + seed(t, tenantHost, "tok", "", time.Now().Add(time.Hour).Unix()) + if err := os.WriteFile(filepath.Join(os.Getenv("HOME"), ".nullify", "hosts.json"), []byte("{not json"), 0600); err != nil { + t.Fatal(err) + } + + _, err := GetValidToken(context.Background(), tenantHost) + if err == nil || !strings.Contains(err.Error(), "hosts.json") { + t.Fatalf("err = %v, want the hosts.json problem named", err) + } +} diff --git a/internal/auth/hosts.go b/internal/auth/hosts.go index 4008d4b..fd96c3e 100644 --- a/internal/auth/hosts.go +++ b/internal/auth/hosts.go @@ -12,9 +12,11 @@ import ( // HostConfig is what the CLI remembers about one API host. AuthHost is the // host that ran its login and serves its refresh, which on a router tenant is -// the router rather than the tenant itself. +// the router rather than the tenant itself. Org is the org the user picked +// for this host, if any. type HostConfig struct { - AuthHost string `json:"auth_host,omitempty"` + AuthHost string `json:"auth_host,omitempty"` + Org *OrgPick `json:"org,omitempty"` } // Hosts maps an API host, in CredentialKey form, to its HostConfig. It lives @@ -117,7 +119,7 @@ func LookupCredentials(host string) (hostCreds HostCredentials, authHost string, return hostCreds, authHost, ok, nil } -// KnownHosts lists the API hosts the user has signed in to, sorted. Router +// KnownHosts lists the API hosts the user is signed in to, sorted. Router // credentials are shared by the tenants bound to them, so the router's own key // is listed only when no tenant is. func KnownHosts() ([]string, error) { @@ -139,8 +141,12 @@ func KnownHosts() ([]string, error) { delete(known, authKey) } } - for h := range hosts { - known[h] = true + // A bound host is known while the login behind it is: a logout removes + // that login and, with it, every host signed in through it. + for h, hc := range hosts { + if _, ok := creds[CredentialKey(hc.AuthHost)]; ok && hc.AuthHost != "" { + known[h] = true + } } return slices.Sorted(maps.Keys(known)), nil diff --git a/internal/auth/login.go b/internal/auth/login.go index 6d16f7e..428e641 100644 --- a/internal/auth/login.go +++ b/internal/auth/login.go @@ -8,6 +8,7 @@ import ( "io" "net" "net/http" + "os" "os/exec" "runtime" "strings" @@ -146,6 +147,14 @@ waitLoop: return fmt.Errorf("failed to save host config: %w", err) } + // A login this tenant ran itself before it moved to the router is dead + // weight now; drop it so nothing mistakes it for the current one. + if CredentialKey(authHost) != CredentialKey(host) { + if err := DeleteHostCredentials(host); err != nil { + logger.L(ctx).Warn("failed to remove the tenant's previous login", logger.Err(err)) + } + } + err = SaveConfig(&Config{Host: host}) if err != nil { return fmt.Errorf("failed to save config: %w", err) @@ -173,9 +182,12 @@ func Logout(host string) error { // auth host that issued them, which is where they must be refreshed. func storedCredentials(host string) (HostCredentials, string, error) { hostCreds, authHost, ok, err := LookupCredentials(host) - if err != nil { + if errors.Is(err, os.ErrNotExist) { return HostCredentials{}, "", fmt.Errorf("not authenticated - run 'nullify auth login'") } + if err != nil { + return HostCredentials{}, "", fmt.Errorf("read stored credentials: %w", err) + } if !ok { return HostCredentials{}, "", fmt.Errorf("not authenticated for %s - run 'nullify auth login --host %s'", host, host) } diff --git a/internal/auth/orgs.go b/internal/auth/orgs.go new file mode 100644 index 0000000..812ccd4 --- /dev/null +++ b/internal/auth/orgs.go @@ -0,0 +1,336 @@ +package auth + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "net/url" + "sort" + "strings" +) + +// Org is one organization a request can be scoped to on an API host. Key and +// ID are the query parameter that scopes a request to it, e.g. +// githubOwnerId=1001001. +type Org struct { + Name string `json:"name"` + Provider string `json:"provider"` + Key string `json:"key"` + ID string `json:"id"` + // Host is the API host serving the org, empty when the router has not + // routed it yet. + Host string `json:"host,omitempty"` + // Status is the router's provisioning status: ready, pending or + // unavailable. Empty outside router mode. + Status string `json:"status,omitempty"` +} + +// Query returns the query parameters that scope a request to the org. +func (o Org) Query() map[string]string { + return map[string]string{o.Key: o.ID} +} + +func (o Org) String() string { + return fmt.Sprintf("%s (%s %s)", o.Name, o.Key, o.ID) +} + +// Matches reports whether ref names the org, by name (case-insensitive) or id. +func (o Org) Matches(ref string) bool { + ref = strings.TrimSpace(ref) + return ref != "" && (strings.EqualFold(o.Name, ref) || o.ID == ref) +} + +// OrgList is what the signed-in user can see from one API host. +type OrgList struct { + // Orgs are served by this host and can scope its requests. + Orgs []Org + // Pending are the user's orgs the router has not finished routing. + Pending []Org + // Elsewhere are the user's orgs served by another Nullify host. + Elsewhere []Org +} + +// providerKeys maps a provider family to the query parameter that scopes a +// request to one of its orgs. +var providerKeys = map[string]string{ + "github": "githubOwnerId", + "azure": "azureOrganizationId", + "bitbucket": "bitbucketWorkspaceId", + "gitlab": "gitlabGroupId", +} + +type installationsResponse struct { + Installations []struct { + Name string `json:"name"` + Provider struct { + ID string `json:"id"` + GitHub json.RawMessage `json:"github"` + Azure json.RawMessage `json:"azure"` + Bitbucket json.RawMessage `json:"bitbucket"` + GitLab json.RawMessage `json:"gitlab"` + } `json:"provider"` + NullifyBaseURL *string `json:"nullifyBaseUrl"` + QueryParameters map[string]string `json:"queryParameters"` + ProvisioningStatus string `json:"provisioningStatus"` + } `json:"installations"` +} + +// ListOrgs returns the orgs token can scope requests to on host, read from +// GET /admin/installations on listHost: the router for a router login, the +// tenant itself for a legacy login or an API token. An org counts for host +// when its nullifyBaseUrl names host; when listHost is host itself, an +// installation with no base URL is the tenant's own. +func ListOrgs(ctx context.Context, host, listHost, token string) (OrgList, error) { + req, err := http.NewRequestWithContext(ctx, http.MethodGet, authBaseURL(listHost)+"/admin/installations", nil) + if err != nil { + return OrgList{}, err + } + req.Header.Set("Authorization", "Bearer "+token) + + resp, err := authClient().Do(req) + if err != nil { + return OrgList{}, err + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + body, _ := io.ReadAll(io.LimitReader(resp.Body, 1024)) + return OrgList{}, fmt.Errorf("list organizations on %s: status %d: %s", apiHost(listHost), resp.StatusCode, string(body)) + } + + var out installationsResponse + if err := json.NewDecoder(io.LimitReader(resp.Body, 10<<20)).Decode(&out); err != nil { + return OrgList{}, fmt.Errorf("decode organizations: %w", err) + } + + target := strings.ToLower(apiHost(host)) + listIsTarget := strings.EqualFold(apiHost(listHost), target) + + var list OrgList + for _, inst := range out.Installations { + family := providerFamily(inst.Provider.ID, inst.Provider.GitHub, inst.Provider.Azure, inst.Provider.Bitbucket, inst.Provider.GitLab) + key := providerKeys[family] + id := inst.QueryParameters[key] + if key == "" || id == "" { + continue + } + org := Org{Name: inst.Name, Provider: family, Key: key, ID: id, Status: inst.ProvisioningStatus} + + if inst.NullifyBaseURL != nil { + org.Host = baseURLHost(*inst.NullifyBaseURL) + } + + switch { + case org.Host == target: + list.Orgs = append(list.Orgs, org) + case org.Host != "": + list.Elsewhere = append(list.Elsewhere, org) + case inst.NullifyBaseURL != nil: + // A base URL that is not a plain https origin is no route at all. + org.Status = "unavailable" + list.Pending = append(list.Pending, org) + case inst.ProvisioningStatus == "pending" || inst.ProvisioningStatus == "unavailable": + list.Pending = append(list.Pending, org) + case listIsTarget: + org.Host = target + list.Orgs = append(list.Orgs, org) + default: + list.Pending = append(list.Pending, org) + } + } + + for _, orgs := range [][]Org{list.Orgs, list.Pending, list.Elsewhere} { + sort.Slice(orgs, func(i, j int) bool { return strings.ToLower(orgs[i].Name) < strings.ToLower(orgs[j].Name) }) + } + return list, nil +} + +func providerFamily(id string, github, azure, bitbucket, gitlab json.RawMessage) string { + present := func(raw json.RawMessage) bool { return len(raw) > 0 && string(raw) != "null" } + switch { + case present(github): + return "github" + case present(azure): + return "azure" + case present(bitbucket): + return "bitbucket" + case present(gitlab): + return "gitlab" + } + return strings.ToLower(id) +} + +// baseURLHost returns the API host a tenant's nullifyBaseUrl names, lower +// case and without a default port, or "" when it is not a plain https origin. +// It accepts what the router's own route validation accepts. +func baseURLHost(raw string) string { + u, err := url.Parse(strings.TrimSpace(raw)) + if err != nil || u.Scheme != "https" || u.User != nil || u.RawQuery != "" || u.Fragment != "" { + return "" + } + if u.Path != "" && u.Path != "/" { + return "" + } + if port := u.Port(); port != "" && port != "443" { + return "" + } + return strings.ToLower(u.Hostname()) +} + +// OrgPick is an org the user chose explicitly for a host, with `nullify org +// use` or at the prompt. It is the only scope the CLI remembers. +type OrgPick struct { + Name string `json:"name"` + Key string `json:"key"` + ID string `json:"id"` +} + +// OrgInputs are the ways a command can name its org, in precedence order. +type OrgInputs struct { + // Flag is --org, or a caller's equivalent (an MCP tool's org argument). + Flag string + // FlagSource names Flag in errors, e.g. "--org". + FlagSource string + // EnvOrg is NULLIFY_ORG. + EnvOrg string + // EnvGitHubOwnerID is NULLIFY_GITHUB_OWNER_ID. + EnvGitHubOwnerID string + // Stored is the remembered pick for this host. + Stored *OrgPick + // GitOwner is the owner in the current directory's git remote. + GitOwner string +} + +// ErrOrgNotChosen reports that several orgs could serve the request and +// nothing chose one. Orgs lists the choices. +type ErrOrgNotChosen struct { + Host string + Orgs []Org +} + +func (e *ErrOrgNotChosen) Error() string { + names := make([]string, len(e.Orgs)) + for i, o := range e.Orgs { + names[i] = o.String() + } + return fmt.Sprintf("%s serves %d organizations; choose one with --org , NULLIFY_ORG, or 'nullify org use': %s", + e.Host, len(e.Orgs), strings.Join(names, ", ")) +} + +// ErrNoOrgs reports that the user has no org routed to this host. +var ErrNoOrgs = errors.New("no organization on this host is available to you") + +// ResolveOrg picks the org a request on host runs as, from list and in: +// an explicit name or id (--org, then NULLIFY_ORG, then +// NULLIFY_GITHUB_OWNER_ID), the stored pick if it is still in the list, the +// git remote owner if it names exactly one org, the only org. An explicit +// name that matches nothing is an error, never a fallthrough. Anything left +// ambiguous returns *ErrOrgNotChosen; the lowest id is never a default. +func ResolveOrg(host string, list OrgList, in OrgInputs) (Org, error) { + explicit := []struct{ ref, source string }{ + {in.Flag, in.FlagSource}, + {in.EnvOrg, "NULLIFY_ORG"}, + } + for _, e := range explicit { + if strings.TrimSpace(e.ref) == "" { + continue + } + return findOrg(host, list, e.ref, e.source) + } + if id := strings.TrimSpace(in.EnvGitHubOwnerID); id != "" { + for _, o := range list.Orgs { + if o.Key == "githubOwnerId" && o.ID == id { + return o, nil + } + } + return Org{}, notFound(host, list, id, "NULLIFY_GITHUB_OWNER_ID") + } + + if in.Stored != nil { + for _, o := range list.Orgs { + if o.Key == in.Stored.Key && o.ID == in.Stored.ID { + return o, nil + } + } + } + + if owner := strings.TrimSpace(in.GitOwner); owner != "" { + var match []Org + for _, o := range list.Orgs { + if strings.EqualFold(o.Name, owner) { + match = append(match, o) + } + } + if len(match) == 1 { + return match[0], nil + } + } + + switch len(list.Orgs) { + case 0: + return Org{}, noOrgs(host, list) + case 1: + return list.Orgs[0], nil + default: + return Org{}, &ErrOrgNotChosen{Host: host, Orgs: list.Orgs} + } +} + +func findOrg(host string, list OrgList, ref, source string) (Org, error) { + var match []Org + for _, o := range list.Orgs { + if o.Matches(ref) { + match = append(match, o) + } + } + switch len(match) { + case 1: + return match[0], nil + case 0: + return Org{}, notFound(host, list, ref, source) + default: + return Org{}, &ErrOrgNotChosen{Host: host, Orgs: match} + } +} + +func notFound(host string, list OrgList, ref, source string) error { + for _, o := range list.Pending { + if o.Matches(ref) { + return fmt.Errorf("%s %q: %s is not routed to a Nullify host yet (status %s); try again later", source, ref, o.Name, statusOr(o.Status, "pending")) + } + } + for _, o := range list.Elsewhere { + if o.Matches(ref) { + return fmt.Errorf("%s %q: %s is served by %s, not %s; use --host %s", source, ref, o.Name, o.Host, apiHost(host), strings.TrimPrefix(o.Host, "api.")) + } + } + return fmt.Errorf("%s %q is not one of your organizations on %s (%s)", source, ref, apiHost(host), describeOrgs(list.Orgs)) +} + +func noOrgs(host string, list OrgList) error { + if len(list.Pending) > 0 { + return fmt.Errorf("%w: %s; %d of your organizations are not routed yet", ErrNoOrgs, apiHost(host), len(list.Pending)) + } + return fmt.Errorf("%w: %s", ErrNoOrgs, apiHost(host)) +} + +func describeOrgs(orgs []Org) string { + if len(orgs) == 0 { + return "none available" + } + names := make([]string, len(orgs)) + for i, o := range orgs { + names[i] = o.Name + } + return "available: " + strings.Join(names, ", ") +} + +func statusOr(s, fallback string) string { + if s == "" { + return fallback + } + return s +} diff --git a/internal/auth/orgs_test.go b/internal/auth/orgs_test.go new file mode 100644 index 0000000..574682a --- /dev/null +++ b/internal/auth/orgs_test.go @@ -0,0 +1,203 @@ +package auth + +import ( + "context" + "errors" + "net/http" + "strings" + "testing" +) + +var ( + acmeWeb = Org{Name: "acme-web", Provider: "github", Key: "githubOwnerId", ID: "1001001", Host: tenantAPI} + acmePay = Org{Name: "acme-payments", Provider: "github", Key: "githubOwnerId", ID: "2002002", Host: tenantAPI} + pendingOrg = Org{Name: "newco", Provider: "github", Key: "githubOwnerId", ID: "555", Status: "pending"} + otherStack = Org{Name: "elsewhere-inc", Provider: "github", Key: "githubOwnerId", ID: "666", Host: "api.other.nullify.ai"} +) + +func twoOrgs() OrgList { + return OrgList{Orgs: []Org{acmeWeb, acmePay}, Pending: []Org{pendingOrg}, Elsewhere: []Org{otherStack}} +} + +func TestResolveOrgPrecedence(t *testing.T) { + stored := &OrgPick{Name: "acme-payments", Key: "githubOwnerId", ID: "2002002"} + stale := &OrgPick{Name: "gone", Key: "githubOwnerId", ID: "1"} + + tests := []struct { + name string + list OrgList + in OrgInputs + want Org + }{ + {"--org by name", twoOrgs(), OrgInputs{Flag: "ACME-WEB"}, acmeWeb}, + {"--org by id", twoOrgs(), OrgInputs{Flag: "2002002"}, acmePay}, + {"--org beats NULLIFY_ORG", twoOrgs(), OrgInputs{Flag: "acme-payments", EnvOrg: "acme-web"}, acmePay}, + {"NULLIFY_ORG beats NULLIFY_GITHUB_OWNER_ID", twoOrgs(), OrgInputs{EnvOrg: "acme-payments", EnvGitHubOwnerID: "1001001"}, acmePay}, + {"NULLIFY_GITHUB_OWNER_ID beats the stored pick", twoOrgs(), OrgInputs{EnvGitHubOwnerID: "1001001", Stored: stored}, acmeWeb}, + {"NULLIFY_ORG beats the stored pick", twoOrgs(), OrgInputs{EnvOrg: "acme-web", Stored: stored}, acmeWeb}, + {"stored pick beats the git owner", twoOrgs(), OrgInputs{Stored: stored, GitOwner: "acme-web"}, acmePay}, + {"stored pick no longer in the list falls through to the git owner", twoOrgs(), OrgInputs{Stored: stale, GitOwner: "ACME-WEB"}, acmeWeb}, + {"git owner names one org", twoOrgs(), OrgInputs{GitOwner: "acme-payments"}, acmePay}, + {"git owner matching nothing falls through to the sole org", OrgList{Orgs: []Org{acmeWeb}}, OrgInputs{GitOwner: "someone-else"}, acmeWeb}, + {"sole org", OrgList{Orgs: []Org{acmePay}, Pending: []Org{pendingOrg}}, OrgInputs{}, acmePay}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + tt.in.FlagSource = "--org" + got, err := ResolveOrg(tenantHost, tt.list, tt.in) + if err != nil { + t.Fatalf("ResolveOrg: %v", err) + } + if got != tt.want { + t.Errorf("org = %v, want %v", got, tt.want) + } + }) + } +} + +func TestResolveOrgNeverDefaultsToTheLowestID(t *testing.T) { + stale := &OrgPick{Name: "gone", Key: "githubOwnerId", ID: "1"} + _, err := ResolveOrg(tenantHost, twoOrgs(), OrgInputs{Stored: stale, GitOwner: "someone-else"}) + + var notChosen *ErrOrgNotChosen + if !errors.As(err, ¬Chosen) { + t.Fatalf("err = %v, want *ErrOrgNotChosen", err) + } + if len(notChosen.Orgs) != 2 { + t.Errorf("choices = %v, want both orgs", notChosen.Orgs) + } + for _, want := range []string{"acme-web", "acme-payments", "--org", "nullify org use"} { + if !strings.Contains(err.Error(), want) { + t.Errorf("error %q does not mention %q", err, want) + } + } +} + +func TestResolveOrgExplicitMissIsAnError(t *testing.T) { + tests := []struct { + name string + in OrgInputs + want string + }{ + {"unknown name does not fall through to the stored pick", OrgInputs{Flag: "nope", FlagSource: "--org", Stored: &OrgPick{Key: "githubOwnerId", ID: "2002002"}}, `--org "nope" is not one of your organizations`}, + {"unknown NULLIFY_ORG", OrgInputs{EnvOrg: "nope"}, `NULLIFY_ORG "nope"`}, + {"unknown NULLIFY_GITHUB_OWNER_ID", OrgInputs{EnvGitHubOwnerID: "999"}, `NULLIFY_GITHUB_OWNER_ID "999"`}, + {"pending route reported apart", OrgInputs{Flag: "newco", FlagSource: "--org"}, "not routed to a Nullify host yet (status pending)"}, + {"org on another host reported apart", OrgInputs{Flag: "elsewhere-inc", FlagSource: "--org"}, "served by api.other.nullify.ai"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + _, err := ResolveOrg(tenantHost, twoOrgs(), tt.in) + if err == nil || !strings.Contains(err.Error(), tt.want) { + t.Fatalf("err = %v, want one containing %q", err, tt.want) + } + }) + } +} + +func TestResolveOrgWithNoOrgs(t *testing.T) { + _, err := ResolveOrg(tenantHost, OrgList{Pending: []Org{pendingOrg}}, OrgInputs{}) + if !errors.Is(err, ErrNoOrgs) || !strings.Contains(err.Error(), "not routed yet") { + t.Fatalf("err = %v, want ErrNoOrgs mentioning the pending org", err) + } +} + +const routerInstallations = `{"installations": [ + {"name": "acme-payments", "avatarUrl": "", "provider": {"id": "GitHub", "github": {"ownerId": 2002002}}, + "nullifyBaseUrl": "https://api.acme.nullify.ai", "queryParameters": {"githubOwnerId": "2002002"}, "provisioningStatus": "ready"}, + {"name": "acme-web", "avatarUrl": "", "provider": {"id": "GitHub", "github": {"ownerId": 1001001}}, + "nullifyBaseUrl": "https://API.acme.nullify.ai:443/", "queryParameters": {"githubOwnerId": "1001001"}, "provisioningStatus": "ready"}, + {"name": "elsewhere-inc", "avatarUrl": "", "provider": {"id": "GitHub", "github": {"ownerId": 666}}, + "nullifyBaseUrl": "https://api.other.nullify.ai", "queryParameters": {"githubOwnerId": "666"}, "provisioningStatus": "ready"}, + {"name": "newco", "avatarUrl": "", "provider": {"id": "GitHub", "github": {"ownerId": 555}}, + "queryParameters": {"githubOwnerId": "555"}, "provisioningStatus": "pending"}, + {"name": "pathy", "avatarUrl": "", "provider": {"id": "GitHub", "github": {"ownerId": 777}}, + "nullifyBaseUrl": "https://api.acme.nullify.ai/evil", "queryParameters": {"githubOwnerId": "777"}, "provisioningStatus": "ready"}, + {"name": "contoso", "avatarUrl": "", "provider": {"id": "Azure", "azure": {"organizationId": "c-1"}}, + "nullifyBaseUrl": "https://api.acme.nullify.ai", "queryParameters": {"azureOrganizationId": "c-1"}, "provisioningStatus": "ready"}, + {"name": "ws", "avatarUrl": "", "provider": {"id": "Bitbucket", "bitbucket": {"workspaceId": "w-1"}}, + "nullifyBaseUrl": "https://api.acme.nullify.ai", "queryParameters": {"bitbucketWorkspaceId": "w-1"}, "provisioningStatus": "ready"}, + {"name": "grp", "avatarUrl": "", "provider": {"id": "GitLab", "gitlab": {"groupId": 9}}, + "nullifyBaseUrl": "https://api.acme.nullify.ai", "queryParameters": {"gitlabGroupId": "9"}, "provisioningStatus": "ready"}, + {"name": "slack-only", "avatarUrl": "", "provider": {"id": "Slack", "slack": {}}, + "nullifyBaseUrl": "https://api.acme.nullify.ai", "queryParameters": {"slackTeamId": "T1"}, "provisioningStatus": "ready"} +]}` + +func TestListOrgsFromTheRouter(t *testing.T) { + var gotAuth string + calls := newRoutedAuthServer(t, func(w http.ResponseWriter, r *http.Request, host, path string) { + gotAuth = r.Header.Get("Authorization") + _, _ = w.Write([]byte(routerInstallations)) + }) + + list, err := ListOrgs(context.Background(), tenantHost, routerAPI, "router-token") + if err != nil { + t.Fatalf("ListOrgs: %v", err) + } + if len(*calls) != 1 || (*calls)[0] != (hostCall{routerAPI, "/admin/installations"}) { + t.Errorf("calls = %v, want one installations read on the router", *calls) + } + if gotAuth != "Bearer router-token" { + t.Errorf("Authorization = %q, want the router token", gotAuth) + } + + names := func(orgs []Org) string { + var out []string + for _, o := range orgs { + out = append(out, o.Name+"="+o.Key+":"+o.ID) + } + return strings.Join(out, ",") + } + if got, want := names(list.Orgs), "acme-payments=githubOwnerId:2002002,acme-web=githubOwnerId:1001001,contoso=azureOrganizationId:c-1,grp=gitlabGroupId:9,ws=bitbucketWorkspaceId:w-1"; got != want { + t.Errorf("orgs = %s\nwant %s", got, want) + } + if got, want := names(list.Elsewhere), "elsewhere-inc=githubOwnerId:666"; got != want { + t.Errorf("elsewhere = %s, want %s", got, want) + } + if got, want := names(list.Pending), "newco=githubOwnerId:555,pathy=githubOwnerId:777"; got != want { + t.Errorf("pending = %s, want %s", got, want) + } +} + +// A legacy stack lists its own installations without base URLs. +func TestListOrgsFromTheTenantItself(t *testing.T) { + newRoutedAuthServer(t, func(w http.ResponseWriter, _ *http.Request, host, _ string) { + if host != tenantAPI { + t.Errorf("listed on %s, want the tenant", host) + } + _, _ = w.Write([]byte(`{"installations": [{"name": "acme", "provider": {"id": "GitHub", "github": {}},` + + ` "queryParameters": {"githubOwnerId": "42"}}]}`)) + }) + + list, err := ListOrgs(context.Background(), tenantHost, tenantHost, "tok") + if err != nil { + t.Fatalf("ListOrgs: %v", err) + } + if len(list.Orgs) != 1 || list.Orgs[0].ID != "42" || list.Orgs[0].Host != tenantAPI { + t.Errorf("orgs = %v, want the tenant's own org", list.Orgs) + } +} + +func TestListOrgsFailsClosed(t *testing.T) { + tests := []struct { + name string + handler func(w http.ResponseWriter) + want string + }{ + {"unauthorized", func(w http.ResponseWriter) { w.WriteHeader(http.StatusUnauthorized) }, "status 401"}, + {"redirect is not followed", func(w http.ResponseWriter) { + w.Header().Set("Location", "https://api.evil.example/admin/installations") + w.WriteHeader(http.StatusFound) + }, "status 302"}, + {"garbage", func(w http.ResponseWriter) { _, _ = w.Write([]byte("")) }, "decode organizations"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + newRoutedAuthServer(t, func(w http.ResponseWriter, _ *http.Request, _, _ string) { tt.handler(w) }) + _, err := ListOrgs(context.Background(), tenantHost, routerAPI, "tok") + if err == nil || !strings.Contains(err.Error(), tt.want) { + t.Fatalf("err = %v, want one containing %q", err, tt.want) + } + }) + } +} diff --git a/internal/lib/git.go b/internal/lib/git.go index cc15c82..2d42723 100644 --- a/internal/lib/git.go +++ b/internal/lib/git.go @@ -10,6 +10,22 @@ import ( // DetectRepoFromGit reads the git remote origin URL from the nearest .git/config // and extracts the repository name. func DetectRepoFromGit() string { + if gitConfig := findGitConfig(); gitConfig != "" { + return ParseRepoFromGitConfig(gitConfig) + } + return "" +} + +// DetectOwnerFromGit reads the git remote origin URL from the nearest +// .git/config and extracts the organization that owns the repository. +func DetectOwnerFromGit() string { + if gitConfig := findGitConfig(); gitConfig != "" { + return ExtractRepoOwner(originURL(gitConfig)) + } + return "" +} + +func findGitConfig() string { dir, err := os.Getwd() if err != nil { return "" @@ -18,21 +34,27 @@ func DetectRepoFromGit() string { for { gitConfig := filepath.Join(dir, ".git", "config") if _, err := os.Stat(gitConfig); err == nil { - return ParseRepoFromGitConfig(gitConfig) + return gitConfig } parent := filepath.Dir(dir) if parent == dir { - break + return "" } dir = parent } - - return "" } // ParseRepoFromGitConfig extracts the repo name from a .git/config file. func ParseRepoFromGitConfig(path string) string { + if url := originURL(path); url != "" { + return ExtractRepoName(url) + } + return "" +} + +// originURL returns the url of remote "origin" in a .git/config file. +func originURL(path string) string { f, err := os.Open(path) if err != nil { return "" @@ -57,7 +79,7 @@ func ParseRepoFromGitConfig(path string) string { if inOrigin && strings.HasPrefix(line, "url") { parts := strings.SplitN(line, "=", 2) if len(parts) == 2 { - return ExtractRepoName(strings.TrimSpace(parts[1])) + return strings.TrimSpace(parts[1]) } } } @@ -89,3 +111,35 @@ func ExtractRepoName(remoteURL string) string { return "" } + +// ExtractRepoOwner extracts the owning organization from a git remote URL: the +// first path segment after the host, which is the GitHub org, GitLab group, +// Bitbucket workspace or Azure DevOps organization. Handles SCP-style SSH +// (git@github.com:org/repo.git), ssh:// and https:// forms, and Azure's +// "v3/" SSH prefix. Returns "" when the URL names no owner. +func ExtractRepoOwner(remoteURL string) string { + remoteURL = strings.TrimSpace(remoteURL) + + var path string + if i := strings.Index(remoteURL, "://"); i >= 0 { + rest := remoteURL[i+3:] + slash := strings.Index(rest, "/") + if slash < 0 { + return "" + } + path = rest[slash+1:] + } else if at := strings.Index(remoteURL, ":"); at >= 0 && !strings.Contains(remoteURL[:at], "/") { + path = remoteURL[at+1:] + } else { + return "" + } + + segments := strings.Split(strings.Trim(path, "/"), "/") + if len(segments) > 1 && segments[0] == "v3" && strings.Contains(remoteURL, "dev.azure.com") { + segments = segments[1:] + } + if len(segments) < 2 || segments[0] == "" { + return "" + } + return segments[0] +} diff --git a/internal/lib/git_test.go b/internal/lib/git_test.go index 9bd21bf..6d6ab85 100644 --- a/internal/lib/git_test.go +++ b/internal/lib/git_test.go @@ -1,6 +1,8 @@ package lib import ( + "os" + "path/filepath" "testing" "github.com/stretchr/testify/require" @@ -56,3 +58,43 @@ func TestExtractRepoName(t *testing.T) { }) } } + +func TestExtractRepoOwner(t *testing.T) { + tests := []struct { + name string + remoteURL string + want string + }{ + {"GitHub SSH", "git@github.com:acme-web/api.git", "acme-web"}, + {"GitHub SSH without .git", "git@github.com:acme-payments/web", "acme-payments"}, + {"GitHub HTTPS", "https://github.com/nullify-platform/cli.git", "nullify-platform"}, + {"GitHub HTTPS with credentials", "https://x-access-token:abc@github.com/acme/repo.git", "acme"}, + {"ssh scheme", "ssh://git@github.com/acme/repo.git", "acme"}, + {"ssh scheme with port", "ssh://git@github.com:22/acme/repo.git", "acme"}, + {"GitLab nested group", "https://gitlab.com/group/subgroup/repo.git", "group"}, + {"Bitbucket HTTPS with user", "https://jane@bitbucket.org/workspace/repo.git", "workspace"}, + {"Azure DevOps HTTPS", "https://dev.azure.com/contoso/project/_git/repo", "contoso"}, + {"Azure DevOps SSH", "git@ssh.dev.azure.com:v3/contoso/project/repo", "contoso"}, + {"just a name", "my-repo", ""}, + {"local path", "/srv/git/repo.git", ""}, + {"host only", "https://github.com/", ""}, + {"owner without repo", "https://github.com/acme", ""}, + {"empty", "", ""}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + require.Equal(t, tt.want, ExtractRepoOwner(tt.remoteURL)) + }) + } +} + +func TestParseRepoFromGitConfigStillReadsOrigin(t *testing.T) { + path := filepath.Join(t.TempDir(), "config") + config := "[core]\n\tbare = false\n[remote \"upstream\"]\n\turl = git@github.com:other/fork.git\n" + + "[remote \"origin\"]\n\turl = git@github.com:acme/widgets.git\n\tfetch = +refs/heads/*:refs/remotes/origin/*\n" + require.NoError(t, os.WriteFile(path, []byte(config), 0600)) + + require.Equal(t, "widgets", ParseRepoFromGitConfig(path)) + require.Equal(t, "acme", ExtractRepoOwner(originURL(path))) +} From 74d27313fce4528a83b7bceec941df62ecfbce48 Mon Sep 17 00:00:00 2001 From: whiskeylotus <0xpentest@proton.me> Date: Thu, 29 Oct 2026 17:34:52 -0400 Subject: [PATCH 4/6] feat(mcp): [AGI-6072] scope each MCP call to an org `nullify mcp serve` resolved one org at startup and failed outright on a multi-org login. It now starts on auth alone. Every tool takes an optional org argument and list_organizations lists the choices; a call uses its org argument, then serve --org, then the CLI's chain, and returns the choices when that is still ambiguous. Resource reads follow the same rule. The scope rides in the request context and is set on each request by a transport, so concurrent calls for different orgs cannot cross. Plan-Step: 4/5 Co-Authored-By: Claude Opus 5.5 --- cmd/cli/cmd/mcp.go | 140 ++++++++++++------- cmd/cli/cmd/mcp_test.go | 111 ++++++++++++++++ internal/mcp/orgs.go | 134 +++++++++++++++++++ internal/mcp/orgs_test.go | 259 ++++++++++++++++++++++++++++++++++++ internal/mcp/server.go | 31 +++-- internal/mcp/server_test.go | 1 + 6 files changed, 620 insertions(+), 56 deletions(-) create mode 100644 cmd/cli/cmd/mcp_test.go create mode 100644 internal/mcp/orgs.go create mode 100644 internal/mcp/orgs_test.go diff --git a/cmd/cli/cmd/mcp.go b/cmd/cli/cmd/mcp.go index 11807d5..3e3fc6c 100644 --- a/cmd/cli/cmd/mcp.go +++ b/cmd/cli/cmd/mcp.go @@ -1,10 +1,15 @@ package cmd import ( + "context" + "errors" "fmt" + "net/http" + "slices" "strings" "github.com/nullify-platform/cli/internal/api" + "github.com/nullify-platform/cli/internal/auth" "github.com/nullify-platform/cli/internal/client" "github.com/nullify-platform/cli/internal/lib" "github.com/nullify-platform/cli/internal/mcp" @@ -20,66 +25,107 @@ var mcpCmd = &cobra.Command{ var mcpServeCmd = &cobra.Command{ Use: "serve", Short: "Start the MCP server", - Long: "Start the Nullify MCP server over stdio. Configure your AI tool to run 'nullify mcp serve'.", + Long: "Start the Nullify MCP server over stdio. Configure your AI tool to run 'nullify mcp serve'. " + + "When your login covers several organizations, tools take an optional org argument; --org sets the default.", RunE: func(cmd *cobra.Command, args []string) error { ctx := cmd.Context() - authCtx, err := resolveCommandAuth(ctx) - if err != nil { - return err - } - - queryParams := authCtx.QueryParams - - // Apply --repo flag or auto-detect from git repoFlag, _ := cmd.Flags().GetString("repo") - repo := resolveRepo(repoFlag) - if repo != "" { - queryParams["repository"] = repo - } - - // Resolve --tools flag toolsFlag, _ := cmd.Flags().GetString("tools") - toolSet := mcp.ToolSet(toolsFlag) - validSets := mcp.ValidToolSets() - validSet := false - for _, v := range validSets { - if v == toolsFlag { - validSet = true - break - } - } - if !validSet { - return fmt.Errorf("invalid --tools value %q. Valid values: %s", toolsFlag, strings.Join(validSets, ", ")) - } - - // Drive the generated API client through a refreshing+retrying transport - // so long-running MCP sessions keep working as the token rotates. Tenant - // scoping (owner/installation IDs, --repo) rides along as default params. - tokenProvider := func() (string, error) { - return lib.GetNullifyToken(ctx, authCtx.Host, nullifyToken, githubToken) - } - refreshProvider := func() (string, error) { - return lib.RefreshNullifyToken(ctx, authCtx.Host, nullifyToken, githubToken) - } - httpClient, clientErr := client.NewRefreshingHTTPClient(authCtx.Host, tokenProvider, refreshProvider) - if clientErr != nil { - return fmt.Errorf("failed to create client: %w", clientErr) + apiClient, toolSet, orgs, err := mcpServeConfig(ctx, repoFlag, toolsFlag) + if err != nil { + return err } - // Empty token literal is intentional: the refreshing transport's - // RoundTrip injects (and rotates) the Authorization header after - // api.Client.do sets it, so the static value in api.Client is - // overridden on every request. - apiClient := api.NewClient(authCtx.Host, "", queryParams, api.WithHTTPClient(httpClient)) - err = mcp.ServeWithClient(ctx, apiClient, toolSet) - if err != nil { + if err := mcp.ServeWithClient(ctx, apiClient, toolSet, orgs); err != nil { return fmt.Errorf("MCP server error: %w", err) } return nil }, } +// mcpServeConfig resolves auth, never an org: a login may cover several +// orgs, and each tool call picks one (its org argument, then serve --org, +// then the usual chain). +func mcpServeConfig(ctx context.Context, repoFlag, toolsFlag string) (*api.Client, mcp.ToolSet, mcp.Orgs, error) { + authCtx, err := resolveCommandToken(ctx) + if err != nil { + return nil, "", mcp.Orgs{}, err + } + + queryParams := map[string]string{} + if repo := resolveRepo(repoFlag); repo != "" { + queryParams["repository"] = repo + } + + toolSet := mcp.ToolSet(toolsFlag) + if !slices.Contains(mcp.ValidToolSets(), toolsFlag) { + return nil, "", mcp.Orgs{}, fmt.Errorf("invalid --tools value %q. Valid values: %s", toolsFlag, strings.Join(mcp.ValidToolSets(), ", ")) + } + + // Drive the generated API client through a refreshing+retrying transport + // so long-running MCP sessions keep working as the token rotates; the + // token is refreshed wherever the login lives (the router, for a router + // tenant). Each call's org scope is set by ScopeTransport. + tokenProvider := func() (string, error) { + return lib.GetNullifyToken(ctx, authCtx.Host, nullifyToken, githubToken) + } + refreshProvider := func() (string, error) { + return lib.RefreshNullifyToken(ctx, authCtx.Host, nullifyToken, githubToken) + } + refreshing, err := client.NewRefreshingHTTPClient(authCtx.Host, tokenProvider, refreshProvider) + if err != nil { + return nil, "", mcp.Orgs{}, fmt.Errorf("failed to create client: %w", err) + } + httpClient := &http.Client{ + Timeout: refreshing.Timeout, + Transport: mcp.ScopeTransport(refreshing.Transport), + CheckRedirect: refreshing.CheckRedirect, + } + // Empty token literal is intentional: the refreshing transport's + // RoundTrip injects (and rotates) the Authorization header after + // api.Client.do sets it, so the static value in api.Client is + // overridden on every request. + apiClient := api.NewClient(authCtx.Host, "", queryParams, api.WithHTTPClient(httpClient)) + + serveOrg := orgFlag + orgs := mcp.Orgs{ + Resolve: func(ctx context.Context, org string) (map[string]string, error) { + flag, source := serveOrg, "--org" + if org != "" { + flag, source = org, "org argument" + } + token, err := tokenProvider() + if err != nil { + return nil, err + } + scope, err := resolveOrgScope(ctx, authCtx.Host, token, orgInputsFor(authCtx.Host, flag, source), false) + var notChosen *auth.ErrOrgNotChosen + if errors.As(err, ¬Chosen) { + choices := make([]string, len(notChosen.Orgs)) + for i, o := range notChosen.Orgs { + choices[i] = o.String() + } + return nil, fmt.Errorf("%s serves several of your organizations; pass the org argument (name or id), "+ + "or start the server with --org: %s", authCtx.Host, strings.Join(choices, ", ")) + } + if err != nil { + return nil, err + } + return scope.Query, nil + }, + List: func(ctx context.Context) (any, error) { + token, err := tokenProvider() + if err != nil { + return nil, err + } + return orgListing(ctx, authCtx.Host, token, serveOrg) + }, + } + + return apiClient, toolSet, orgs, nil +} + func init() { rootCmd.AddCommand(mcpCmd) mcpCmd.AddCommand(mcpServeCmd) diff --git a/cmd/cli/cmd/mcp_test.go b/cmd/cli/cmd/mcp_test.go new file mode 100644 index 0000000..93fd238 --- /dev/null +++ b/cmd/cli/cmd/mcp_test.go @@ -0,0 +1,111 @@ +package cmd + +import ( + "context" + "encoding/json" + "io" + "net/http" + "strings" + "sync" + "testing" + + "github.com/nullify-platform/cli/internal/auth" + "github.com/nullify-platform/cli/internal/logger" + "github.com/nullify-platform/cli/internal/mcp" + "github.com/stretchr/testify/require" +) + +// fakeNullify stands in for the router and the acme tenant at the transport +// level, so the real auth, refresh and client code run unchanged. +type fakeNullify struct { + mu sync.Mutex + live string // the access token the tenant accepts + refreshes []string + apiCalls []string +} + +func (f *fakeNullify) RoundTrip(r *http.Request) (*http.Response, error) { + f.mu.Lock() + defer f.mu.Unlock() + reply := func(status int, body string, header http.Header) (*http.Response, error) { + if header == nil { + header = http.Header{} + } + return &http.Response{StatusCode: status, Header: header, Body: io.NopCloser(strings.NewReader(body)), Request: r}, nil + } + + switch { + case r.URL.Host == "api.nullify.ai" && r.URL.Path == "/auth/refresh_token": + f.refreshes = append(f.refreshes, r.URL.Host) + f.live = "refreshed-token" + h := http.Header{} + h.Add("Set-Cookie", "access_token=refreshed-token; Max-Age=3600; Path=/; Secure; HttpOnly") + return reply(200, `{"expiresIn":3600}`, h) + case r.URL.Host == "api.nullify.ai" && r.URL.Path == "/admin/installations": + return reply(200, `{"installations": [ + {"name": "acme-web", "provider": {"id": "GitHub", "github": {}}, "nullifyBaseUrl": "https://api.acme.nullify.ai", + "queryParameters": {"githubOwnerId": "1001001"}, "provisioningStatus": "ready"}, + {"name": "acme-payments", "provider": {"id": "GitHub", "github": {}}, "nullifyBaseUrl": "https://api.acme.nullify.ai", + "queryParameters": {"githubOwnerId": "2002002"}, "provisioningStatus": "ready"}]}`, nil) + case r.URL.Host == "api.acme.nullify.ai": + f.apiCalls = append(f.apiCalls, r.Header.Get("Authorization")+" "+r.URL.Query().Get("githubOwnerId")) + if r.Header.Get("Authorization") != "Bearer "+f.live { + return reply(401, `{"message":"Unauthorized"}`, nil) + } + return reply(200, `{}`, nil) + } + return reply(404, "", nil) +} + +func TestMCPServeOnARouterLogin(t *testing.T) { + loginEnv(t) + fake := &fakeNullify{live: "router-token"} + old := http.DefaultTransport + http.DefaultTransport = fake + t.Cleanup(func() { http.DefaultTransport = old }) + + ctx, err := logger.ConfigureDevelopmentLogger(context.Background(), "error") + require.NoError(t, err) + + apiClient, toolSet, orgs, err := mcpServeConfig(ctx, "widgets", "default") + require.NoError(t, err, "serve starts with two orgs and no org chosen") + s := mcp.NewServer(apiClient, toolSet, orgs) + + call := func(args map[string]any) string { + msg, err := json.Marshal(map[string]any{"jsonrpc": "2.0", "id": 1, "method": "tools/call", + "params": map[string]any{"name": "list_repositories", "arguments": args}}) + require.NoError(t, err) + out, err := json.Marshal(s.HandleMessage(ctx, msg)) + require.NoError(t, err) + return string(out) + } + + out := call(nil) + require.Contains(t, out, `"isError":true`, "two orgs and no org argument") + require.Contains(t, out, "acme-payments") + require.Contains(t, out, "pass the org argument") + + out = call(map[string]any{"org": "acme-payments"}) + require.NotContains(t, out, `"isError":true`, out) + + // The tenant stops accepting the token mid-session: the same process + // refreshes at the router, never the tenant, and retries. + fake.mu.Lock() + fake.live = "only-a-refresh-gets-this" + fake.mu.Unlock() + out = call(map[string]any{"org": "acme-web"}) + require.NotContains(t, out, `"isError":true`, out) + + fake.mu.Lock() + defer fake.mu.Unlock() + require.Equal(t, []string{"api.nullify.ai"}, fake.refreshes) + require.Equal(t, []string{ + "Bearer router-token 2002002", + "Bearer router-token 1001001", + "Bearer refreshed-token 1001001", + }, fake.apiCalls) + + creds, err := auth.LoadCredentials() + require.NoError(t, err) + require.Equal(t, "refreshed-token", creds["nullify.ai"].AccessToken, "the refreshed router login is saved") +} diff --git a/internal/mcp/orgs.go b/internal/mcp/orgs.go new file mode 100644 index 0000000..c23b900 --- /dev/null +++ b/internal/mcp/orgs.go @@ -0,0 +1,134 @@ +package mcp + +import ( + "context" + "encoding/json" + "maps" + "net/http" + + "github.com/mark3labs/mcp-go/mcp" + "github.com/mark3labs/mcp-go/server" +) + +// Orgs is how the server scopes each call to an organization. The server +// starts without one: a login can cover several orgs, and each tool call or +// resource read picks its own. +type Orgs struct { + // Resolve returns the query parameters that scope one call. org is the + // call's org argument, empty when it named none; Resolve then falls back + // to serve --org and the CLI's usual chain, and errors with the choices + // when that is still ambiguous. + Resolve func(ctx context.Context, org string) (map[string]string, error) + // List backs the list_organizations tool. Its result is returned as JSON. + List func(ctx context.Context) (any, error) +} + +const ( + orgArg = "org" + listOrganizationsTool = "list_organizations" +) + +type scopeKey struct{} + +func withScope(ctx context.Context, scope map[string]string) context.Context { + return context.WithValue(ctx, scopeKey{}, scope) +} + +// ScopeTransport sets the org scope a call resolved onto every request that +// call makes, over the client's default parameters. Scope travels in the +// request context, never in shared state, so concurrent calls for different +// orgs cannot cross. +func ScopeTransport(next http.RoundTripper) http.RoundTripper { + return scopeTransport{next: next} +} + +type scopeTransport struct{ next http.RoundTripper } + +func (t scopeTransport) RoundTrip(req *http.Request) (*http.Response, error) { + scope, _ := req.Context().Value(scopeKey{}).(map[string]string) + if len(scope) == 0 { + return t.next.RoundTrip(req) + } + r := req.Clone(req.Context()) + q := r.URL.Query() + for k, v := range scope { + q.Set(k, v) + } + r.URL.RawQuery = q.Encode() + return t.next.RoundTrip(r) +} + +// orgMiddleware resolves each tool call's org before the tool runs. +func orgMiddleware(orgs Orgs) server.ToolHandlerMiddleware { + return func(next server.ToolHandlerFunc) server.ToolHandlerFunc { + return func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) { + if req.Params.Name == listOrganizationsTool { + return next(ctx, req) + } + scope, err := orgs.Resolve(ctx, getStringArg(req.GetArguments(), orgArg)) + if err != nil { + return toolError(err), nil + } + return next(withScope(ctx, scope), req) + } + } +} + +// resourceOrgMiddleware gives resource reads, which take no arguments, the +// same org rule as a tool call with no org argument. +func resourceOrgMiddleware(orgs Orgs) server.ResourceHandlerMiddleware { + return func(next server.ResourceHandlerFunc) server.ResourceHandlerFunc { + return func(ctx context.Context, req mcp.ReadResourceRequest) ([]mcp.ResourceContents, error) { + scope, err := orgs.Resolve(ctx, "") + if err != nil { + return nil, err + } + return next(withScope(ctx, scope), req) + } + } +} + +// addOrgArgument adds the optional org argument to every registered tool. +func addOrgArgument(s *server.MCPServer) { + tools := s.ListTools() + updated := make([]server.ServerTool, 0, len(tools)) + for _, st := range tools { + tool := st.Tool + props := maps.Clone(tool.InputSchema.Properties) + if props == nil { + props = map[string]any{} + } + props[orgArg] = map[string]any{ + "type": "string", + "description": "Organization to run as, by name or id. Needed only when your login covers several " + + "organizations and no default is set; list_organizations shows them.", + } + tool.InputSchema.Properties = props + if tool.InputSchema.Type == "" { + tool.InputSchema.Type = "object" + } + updated = append(updated, server.ServerTool{Tool: tool, Handler: st.Handler}) + } + s.SetTools(updated...) +} + +func registerOrgTools(s *server.MCPServer, orgs Orgs) { + s.AddTool( + mcp.NewTool(listOrganizationsTool, + mcp.WithDescription("List the organizations your Nullify login can use on this host, which one calls use by default, "+ + "and any not routed yet. Pass one as the org argument of another tool to run it as that organization."), + mcp.WithReadOnlyHintAnnotation(true), + ), + func(ctx context.Context, _ mcp.CallToolRequest) (*mcp.CallToolResult, error) { + out, err := orgs.List(ctx) + if err != nil { + return toolError(err), nil + } + data, err := json.MarshalIndent(out, "", " ") + if err != nil { + return toolError(err), nil + } + return toolResult(string(data)), nil + }, + ) +} diff --git a/internal/mcp/orgs_test.go b/internal/mcp/orgs_test.go new file mode 100644 index 0000000..e2cb446 --- /dev/null +++ b/internal/mcp/orgs_test.go @@ -0,0 +1,259 @@ +package mcp + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "net/http" + "net/http/httptest" + "net/url" + "strings" + "sync" + "testing" + + "github.com/mark3labs/mcp-go/server" + "github.com/nullify-platform/cli/internal/api" + "github.com/stretchr/testify/require" +) + +// scopeRecorder is a fake Nullify API that answers every call with an empty +// object and records each request's query. +type scopeRecorder struct { + mu sync.Mutex + queries []url.Values +} + +func (r *scopeRecorder) handler(w http.ResponseWriter, req *http.Request) { + r.mu.Lock() + r.queries = append(r.queries, req.URL.Query()) + r.mu.Unlock() + _, _ = w.Write([]byte(`{}`)) +} + +func (r *scopeRecorder) seen() []url.Values { + r.mu.Lock() + defer r.mu.Unlock() + return append([]url.Values(nil), r.queries...) +} + +// orgsFor fakes the CLI's resolution: org names map to owner ids, and an +// empty name falls back to def, or fails with the choices when def is "". +func orgsFor(byName map[string]string, def string) Orgs { + names := make([]string, 0, len(byName)) + for n := range byName { + names = append(names, n) + } + return Orgs{ + Resolve: func(_ context.Context, org string) (map[string]string, error) { + if org == "" { + org = def + } + if org == "" { + return nil, fmt.Errorf("acme.nullify.ai serves %d organizations; choose one with --org: %s", len(byName), strings.Join(names, ", ")) + } + id, ok := byName[org] + if !ok { + return nil, fmt.Errorf("org argument %q is not one of your organizations", org) + } + return map[string]string{"githubOwnerId": id}, nil + }, + List: func(context.Context) (any, error) { + return map[string]any{"organizations": names}, nil + }, + } +} + +func newTestServer(t *testing.T, orgs Orgs) (*server.MCPServer, *scopeRecorder) { + t.Helper() + rec := &scopeRecorder{} + srv := httptest.NewServer(http.HandlerFunc(rec.handler)) + t.Cleanup(srv.Close) + + c := api.NewClient("acme.nullify.ai", "tok", map[string]string{"repository": "widgets"}, + api.WithHTTPClient(&http.Client{Transport: ScopeTransport(srv.Client().Transport)})) + c.BaseURL = srv.URL + return NewServer(c, ToolSetDefault, orgs), rec +} + +type rpcResult struct { + Result struct { + IsError bool `json:"isError"` + Content []struct { + Text string `json:"text"` + } `json:"content"` + Contents []struct { + Text string `json:"text"` + } `json:"contents"` + Tools []struct { + Name string `json:"name"` + InputSchema struct { + Properties map[string]any `json:"properties"` + } `json:"inputSchema"` + } `json:"tools"` + } `json:"result"` + Error *struct { + Message string `json:"message"` + } `json:"error"` +} + +func call(t *testing.T, s *server.MCPServer, method string, params any) rpcResult { + t.Helper() + msg, err := json.Marshal(map[string]any{"jsonrpc": "2.0", "id": 1, "method": method, "params": params}) + require.NoError(t, err) + out, err := json.Marshal(s.HandleMessage(context.Background(), msg)) + require.NoError(t, err) + var res rpcResult + require.NoError(t, json.Unmarshal(out, &res), string(out)) + return res +} + +func callTool(t *testing.T, s *server.MCPServer, name string, args map[string]any) rpcResult { + t.Helper() + return call(t, s, "tools/call", map[string]any{"name": name, "arguments": args}) +} + +func TestServerStartsWithSeveralOrgsAndNoDefault(t *testing.T) { + resolved := 0 + orgs := orgsFor(map[string]string{"acme-web": "1001001", "acme-payments": "2002002"}, "") + resolve := orgs.Resolve + orgs.Resolve = func(ctx context.Context, org string) (map[string]string, error) { + resolved++ + return resolve(ctx, org) + } + + s, rec := newTestServer(t, orgs) + res := call(t, s, "tools/list", map[string]any{}) + require.Nil(t, res.Error) + require.Zero(t, resolved, "listing tools must not need an org") + require.Empty(t, rec.seen(), "starting must not call the API") +} + +func TestEveryToolTakesAnOrgExceptListOrganizations(t *testing.T) { + s, _ := newTestServer(t, orgsFor(map[string]string{"a": "1"}, "a")) + res := call(t, s, "tools/list", map[string]any{}) + require.Nil(t, res.Error) + require.NotEmpty(t, res.Result.Tools) + + sawList := false + for _, tool := range res.Result.Tools { + _, hasOrg := tool.InputSchema.Properties["org"] + if tool.Name == "list_organizations" { + sawList = true + require.False(t, hasOrg, "list_organizations takes no org") + continue + } + require.True(t, hasOrg, "%s has no org argument", tool.Name) + } + require.True(t, sawList) +} + +func TestListOrganizationsTool(t *testing.T) { + s, rec := newTestServer(t, orgsFor(map[string]string{"acme-web": "1001001"}, "")) + res := callTool(t, s, "list_organizations", nil) + require.False(t, res.Result.IsError, res.Result.Content) + require.Contains(t, res.Result.Content[0].Text, "acme-web") + require.Empty(t, rec.seen()) +} + +func TestOneOrgNeedsNoArgument(t *testing.T) { + s, rec := newTestServer(t, orgsFor(map[string]string{"acme-payments": "2002002"}, "acme-payments")) + res := callTool(t, s, "list_repositories", nil) + require.False(t, res.Result.IsError, res.Result.Content) + + seen := rec.seen() + require.Len(t, seen, 1) + require.Equal(t, "2002002", seen[0].Get("githubOwnerId")) + require.Equal(t, "widgets", seen[0].Get("repository"), "fixed scope survives") +} + +func TestSeveralOrgsAndNoArgumentReturnsTheChoices(t *testing.T) { + s, rec := newTestServer(t, orgsFor(map[string]string{"acme-web": "1001001", "acme-payments": "2002002"}, "")) + res := callTool(t, s, "list_repositories", nil) + + require.True(t, res.Result.IsError) + require.Contains(t, res.Result.Content[0].Text, "acme-web") + require.Contains(t, res.Result.Content[0].Text, "acme-payments") + require.Empty(t, rec.seen(), "no API call without an org") +} + +func TestConcurrentCallsKeepTheirOwnOrg(t *testing.T) { + s, rec := newTestServer(t, orgsFor(map[string]string{"acme-web": "1001001", "acme-payments": "2002002"}, "")) + + const n = 20 + var wg sync.WaitGroup + errs := make(chan error, 2*n) + for i := 0; i < n; i++ { + for _, org := range []string{"acme-web", "acme-payments"} { + wg.Add(1) + go func(org string) { + defer wg.Done() + msg, _ := json.Marshal(map[string]any{"jsonrpc": "2.0", "id": 1, "method": "tools/call", + "params": map[string]any{"name": "list_repositories", "arguments": map[string]any{"org": org}}}) + out, _ := json.Marshal(s.HandleMessage(context.Background(), msg)) + if strings.Contains(string(out), `"isError":true`) { + errs <- errors.New(string(out)) + } + }(org) + } + } + wg.Wait() + close(errs) + for err := range errs { + t.Fatal(err) + } + + counts := map[string]int{} + for _, q := range rec.seen() { + require.Len(t, q["githubOwnerId"], 1, "one org per request: %v", q) + counts[q.Get("githubOwnerId")]++ + } + require.Equal(t, map[string]int{"1001001": n, "2002002": n}, counts) +} + +func TestUnknownOrgArgumentIsAToolError(t *testing.T) { + s, rec := newTestServer(t, orgsFor(map[string]string{"acme-payments": "2002002"}, "acme-payments")) + res := callTool(t, s, "list_repositories", map[string]any{"org": "nope"}) + require.True(t, res.Result.IsError) + require.Contains(t, res.Result.Content[0].Text, `"nope"`) + require.Empty(t, rec.seen()) +} + +func TestResourceReadFollowsTheSameRule(t *testing.T) { + ambiguous, rec := newTestServer(t, orgsFor(map[string]string{"acme-web": "1001001", "acme-payments": "2002002"}, "")) + res := call(t, ambiguous, "resources/read", map[string]any{"uri": "nullify://repos"}) + require.NotNil(t, res.Error, "ambiguous org must fail the read") + require.Contains(t, res.Error.Message, "acme-payments") + require.Empty(t, rec.seen()) + + single, rec := newTestServer(t, orgsFor(map[string]string{"acme-payments": "2002002"}, "acme-payments")) + res = call(t, single, "resources/read", map[string]any{"uri": "nullify://repos"}) + require.Nil(t, res.Error) + require.Len(t, rec.seen(), 1) + require.Equal(t, "2002002", rec.seen()[0].Get("githubOwnerId")) +} + +func TestScopeTransportLeavesUnscopedRequestsAlone(t *testing.T) { + var got string + tr := ScopeTransport(roundTripFunc(func(r *http.Request) (*http.Response, error) { + got = r.URL.RawQuery + return &http.Response{StatusCode: 200, Body: http.NoBody}, nil + })) + req, _ := http.NewRequest(http.MethodGet, "https://api.acme.nullify.ai/x?repository=w", nil) + _, err := tr.RoundTrip(req) + require.NoError(t, err) + require.Equal(t, "repository=w", got) + + req, _ = http.NewRequestWithContext(withScope(context.Background(), map[string]string{"githubOwnerId": "1"}), + http.MethodGet, "https://api.acme.nullify.ai/x?repository=w&githubOwnerId=9", nil) + _, err = tr.RoundTrip(req) + require.NoError(t, err) + q, _ := url.ParseQuery(got) + require.Equal(t, []string{"1"}, q["githubOwnerId"], "the call's org replaces any default") + require.Equal(t, "w", q.Get("repository")) + require.Equal(t, "https://api.acme.nullify.ai/x?repository=w&githubOwnerId=9", req.URL.String(), "the caller's request is not mutated") +} + +type roundTripFunc func(*http.Request) (*http.Response, error) + +func (f roundTripFunc) RoundTrip(r *http.Request) (*http.Response, error) { return f(r) } diff --git a/internal/mcp/server.go b/internal/mcp/server.go index f5cc2b8..bc418e4 100644 --- a/internal/mcp/server.go +++ b/internal/mcp/server.go @@ -14,28 +14,41 @@ import ( ) // ServeWithClient runs the MCP server over stdio, driving the generated -// api.Client. Tenant/owner scoping travels in the client's default query params. -func ServeWithClient(ctx context.Context, apiClient *api.Client, toolSet ToolSet) error { - return serveWithClientIO(ctx, apiClient, toolSet, os.Stdin, os.Stdout) +// api.Client. The client's default query params carry fixed scope such as +// --repo; each call's org comes from orgs and travels in the request context, +// so apiClient's HTTP client must route through ScopeTransport. +func ServeWithClient(ctx context.Context, apiClient *api.Client, toolSet ToolSet, orgs Orgs) error { + return serveWithClientIO(ctx, apiClient, toolSet, orgs, os.Stdin, os.Stdout) } -func serveWithClientIO(ctx context.Context, apiClient *api.Client, toolSet ToolSet, stdin io.Reader, stdout io.Writer) error { +func serveWithClientIO(ctx context.Context, apiClient *api.Client, toolSet ToolSet, orgs Orgs, stdin io.Reader, stdout io.Writer) error { + s := NewServer(apiClient, toolSet, orgs) + + logger.L(ctx).Debug("starting MCP server over stdio", logger.String("toolSet", string(toolSet))) + + stdioServer := server.NewStdioServer(s) + return stdioServer.Listen(ctx, stdin, stdout) +} + +// NewServer builds the MCP server: the tool set, list_organizations, an +// optional org argument on every other tool, and the resources and prompts. +func NewServer(apiClient *api.Client, toolSet ToolSet, orgs Orgs) *server.MCPServer { s := server.NewMCPServer( "Nullify", logger.Version, server.WithToolCapabilities(true), server.WithResourceCapabilities(true, false), server.WithPromptCapabilities(true), + server.WithToolHandlerMiddleware(orgMiddleware(orgs)), + server.WithResourceHandlerMiddleware(resourceOrgMiddleware(orgs)), ) registerTools(s, apiClient, toolSet) + addOrgArgument(s) + registerOrgTools(s, orgs) registerResources(s, apiClient) registerPrompts(s) - - logger.L(ctx).Debug("starting MCP server over stdio", logger.String("toolSet", string(toolSet))) - - stdioServer := server.NewStdioServer(s) - return stdioServer.Listen(ctx, stdin, stdout) + return s } // registerTools selects which tools to expose. The default set is deliberately diff --git a/internal/mcp/server_test.go b/internal/mcp/server_test.go index b72bc66..ede1930 100644 --- a/internal/mcp/server_test.go +++ b/internal/mcp/server_test.go @@ -54,6 +54,7 @@ func TestServeWithClientIOHandlesEOF(t *testing.T) { ctx, api.NewClient("acme.nullify.ai", "token", map[string]string{}), ToolSetDefault, + Orgs{}, strings.NewReader(""), io.Discard, ) From c23fb398da501e16c86498e201c4e6ed06c5faa9 Mon Sep 17 00:00:00 2001 From: whiskeylotus <0xpentest@proton.me> Date: Thu, 29 Oct 2026 17:34:52 -0400 Subject: [PATCH 5/6] docs: [AGI-6072] router login, --org and the MCP org argument Plan-Step: 5/5 Co-Authored-By: Claude Opus 5.5 --- README.md | 68 ++++++++++++++++++++++++++++++++++++++++++++++++++++--- 1 file changed, 65 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index 11d9b96..2a3b019 100644 --- a/README.md +++ b/README.md @@ -98,6 +98,34 @@ nullify auth login --host api.acme.nullify.ai Opens your browser for single sign-on. Tokens are stored locally and refreshed automatically. +Before it signs in, the CLI asks your instance where its login runs (`GET /auth/cli/config`). An instance attached to the Nullify router answers with the router's host (`api.nullify.ai`), and the CLI signs in and refreshes there; one login then covers every instance that shares it. An instance without that endpoint signs in on its own host. The CLI only follows an answer naming your instance or a Nullify router, and fails otherwise. + +### Organizations + +One login can cover several organizations on the same instance. Each command runs as the first of these that applies: + +1. `--org ` +2. `NULLIFY_ORG` (name or id) +3. `NULLIFY_GITHUB_OWNER_ID` +4. The organization you picked with `nullify org use`, if your login still covers it +5. The owner of the current directory's git remote, if it names one of your organizations +6. Your only organization on the instance + +When none applies, the CLI asks in a terminal and remembers the answer; elsewhere it fails and lists the choices. + +```sh +# See your organizations on this instance and which one is in use +nullify org list + +# Remember one for this instance +nullify org use acme-web + +# Run one command as another +nullify findings --org acme-payments +``` + +An API token (`NULLIFY_TOKEN`) belongs to one organization, and the server uses it when the command names none. + ### Environment Variables ```sh @@ -133,10 +161,12 @@ nullify auth switch |---------|-------------| | `nullify auth login` | Authenticate via browser SSO | | `nullify auth logout` | Clear stored credentials for a host | -| `nullify auth status` | Show auth state, host, and token expiry | +| `nullify auth status` | Show auth state, host, auth host, organization, and token expiry | | `nullify auth token` | Print raw access token to stdout (pipe-friendly) | | `nullify auth switch` | Switch active instance or list configured instances | -| `nullify auth config` | Print current CLI config as JSON | +| `nullify auth config` | Print the effective host, auth host, and organization as JSON | +| `nullify org list` | List the organizations you can use on this instance | +| `nullify org use ` | Remember the organization to use on this instance | ### Scanning @@ -274,6 +304,29 @@ Add the following to your MCP config file: } ``` +When your login covers several organizations, every tool takes an optional `org` argument (name or id), and the `list_organizations` tool lists them. Without the argument a call uses `--org`, then the same order as the CLI; if that still leaves several, the call returns the choices instead of guessing. To give each organization its own server, add one entry per organization: + +```json +{ + "mcpServers": { + "nullify-acme-web": { + "command": "nullify", + "args": ["mcp", "serve", "--org", "acme-web"] + }, + "nullify-acme-payments": { + "command": "nullify", + "args": ["mcp", "serve", "--org", "acme-payments"] + } + } +} +``` + +With Claude Code: + +```sh +claude mcp add nullify-acme-web -- nullify mcp serve --org acme-web +``` + Scope findings to a specific repository with `--repo`: ```json @@ -289,7 +342,15 @@ Scope findings to a specific repository with `--repo`: ## Configuration -The CLI stores configuration at `~/.nullify/config.json`. Host resolution priority: +The CLI stores configuration in `~/.nullify/`: + +| File | Holds | +|------|-------| +| `config.json` | The active instance | +| `credentials.json` | Login tokens, one entry per login host (the router's login is shared by the instances that use it) | +| `hosts.json` | Per instance: the host its login runs on, and the organization picked with `nullify org use` | + +Host resolution priority: 1. `--host` flag 2. Config file (`~/.nullify/config.json`) @@ -300,6 +361,7 @@ The CLI stores configuration at `~/.nullify/config.json`. Host resolution priori | Flag | Description | Default | |------|-------------|---------| | `--host` | Nullify API instance (e.g., `api.acme.nullify.ai`) | From config | +| `--org` | Organization to run as, by name or id (see [Organizations](#organizations)) | | | `--output, -o` | Output format (`json`, `table`, `yaml`) | `json` | | `--verbose, -v` | Enable verbose logging | `false` | | `--debug, -d` | Enable debug logging | `false` | From 18838501ad0cb9f4cd2f47e68ec0ec38144f78ca Mon Sep 17 00:00:00 2001 From: whiskeylotus <0xpentest@proton.me> Date: Thu, 29 Oct 2026 17:46:25 -0400 Subject: [PATCH 6/6] fix(cli,mcp): [AGI-6072] address wave 2 review - MCP refuses an org argument that is blank or not a name or id instead of running the call as the default org; a numeric id is taken as that id. - The org list retries once with a force-refreshed token when the router rejects the stored one, so a revoked or skewed token no longer breaks every MCP call until its stored expiry. - An ambiguous --org settled at the prompt no longer replaces the remembered pick, and API-token commands no longer read hosts.json. - `org list` reports a NULLIFY_GITHUB_OWNER_ID scope, and has help text. - README: NULLIFY_ORG, API-token org rules, and the real host order. - Tests for serve --org as the MCP default. Co-Authored-By: Claude Opus 5.5 --- README.md | 7 ++++--- cmd/cli/cmd/mcp_test.go | 40 ++++++++++++++++++++++++++++++++--- cmd/cli/cmd/org.go | 44 ++++++++++++++++++++++++++++++--------- cmd/cli/cmd/org_test.go | 32 ++++++++++++++++++++++++++++ internal/auth/orgs.go | 10 ++++++++- internal/mcp/orgs.go | 36 +++++++++++++++++++++++++++++++- internal/mcp/orgs_test.go | 19 +++++++++++++++++ 7 files changed, 170 insertions(+), 18 deletions(-) diff --git a/README.md b/README.md index 2a3b019..ab4ed69 100644 --- a/README.md +++ b/README.md @@ -124,13 +124,14 @@ nullify org use acme-web nullify findings --org acme-payments ``` -An API token (`NULLIFY_TOKEN`) belongs to one organization, and the server uses it when the command names none. +An API token (`NULLIFY_TOKEN`) belongs to one organization, and the server uses it when the command names none. With a token, only `--org`, `NULLIFY_ORG` and `NULLIFY_GITHUB_OWNER_ID` apply (the first two are looked up on the instance itself); `nullify org use`, the git remote and the prompt are for browser logins. ### Environment Variables ```sh export NULLIFY_HOST=api.acme.nullify.ai export NULLIFY_TOKEN=your-token-here +export NULLIFY_ORG=acme-web # optional: organization to run as, by name or id ``` ### Multiple Instances @@ -353,8 +354,8 @@ The CLI stores configuration in `~/.nullify/`: Host resolution priority: 1. `--host` flag -2. Config file (`~/.nullify/config.json`) -3. `NULLIFY_HOST` environment variable +2. `NULLIFY_HOST` environment variable +3. Config file (`~/.nullify/config.json`) ### Global Flags diff --git a/cmd/cli/cmd/mcp_test.go b/cmd/cli/cmd/mcp_test.go index 93fd238..3570ead 100644 --- a/cmd/cli/cmd/mcp_test.go +++ b/cmd/cli/cmd/mcp_test.go @@ -42,6 +42,9 @@ func (f *fakeNullify) RoundTrip(r *http.Request) (*http.Response, error) { h.Add("Set-Cookie", "access_token=refreshed-token; Max-Age=3600; Path=/; Secure; HttpOnly") return reply(200, `{"expiresIn":3600}`, h) case r.URL.Host == "api.nullify.ai" && r.URL.Path == "/admin/installations": + if r.Header.Get("Authorization") != "Bearer "+f.live { + return reply(401, `{"message":"Unauthorized"}`, nil) + } return reply(200, `{"installations": [ {"name": "acme-web", "provider": {"id": "GitHub", "github": {}}, "nullifyBaseUrl": "https://api.acme.nullify.ai", "queryParameters": {"githubOwnerId": "1001001"}, "provisioningStatus": "ready"}, @@ -88,8 +91,9 @@ func TestMCPServeOnARouterLogin(t *testing.T) { out = call(map[string]any{"org": "acme-payments"}) require.NotContains(t, out, `"isError":true`, out) - // The tenant stops accepting the token mid-session: the same process - // refreshes at the router, never the tenant, and retries. + // The token stops being accepted mid-session, before its stored expiry: + // the same process refreshes at the router, never the tenant, for the + // org list and for the API call, and both retry. fake.mu.Lock() fake.live = "only-a-refresh-gets-this" fake.mu.Unlock() @@ -98,7 +102,7 @@ func TestMCPServeOnARouterLogin(t *testing.T) { fake.mu.Lock() defer fake.mu.Unlock() - require.Equal(t, []string{"api.nullify.ai"}, fake.refreshes) + require.Equal(t, []string{"api.nullify.ai", "api.nullify.ai"}, fake.refreshes) require.Equal(t, []string{ "Bearer router-token 2002002", "Bearer router-token 1001001", @@ -109,3 +113,33 @@ func TestMCPServeOnARouterLogin(t *testing.T) { require.NoError(t, err) require.Equal(t, "refreshed-token", creds["nullify.ai"].AccessToken, "the refreshed router login is saved") } + +// serve --org is the default for calls that name no org; a call's own org +// argument still wins. +func TestMCPServeOrgFlagIsTheDefault(t *testing.T) { + loginEnv(t) + fake := &fakeNullify{live: "router-token"} + old := http.DefaultTransport + http.DefaultTransport = fake + t.Cleanup(func() { http.DefaultTransport = old }) + orgFlag = "acme-payments" + + ctx, err := logger.ConfigureDevelopmentLogger(context.Background(), "error") + require.NoError(t, err) + apiClient, toolSet, orgs, err := mcpServeConfig(ctx, "", "default") + require.NoError(t, err) + s := mcp.NewServer(apiClient, toolSet, orgs) + + for _, args := range []map[string]any{nil, {"org": "acme-web"}} { + msg, err := json.Marshal(map[string]any{"jsonrpc": "2.0", "id": 1, "method": "tools/call", + "params": map[string]any{"name": "list_repositories", "arguments": args}}) + require.NoError(t, err) + out, err := json.Marshal(s.HandleMessage(ctx, msg)) + require.NoError(t, err) + require.NotContains(t, string(out), `"isError":true`, string(out)) + } + + fake.mu.Lock() + defer fake.mu.Unlock() + require.Equal(t, []string{"Bearer router-token 2002002", "Bearer router-token 1001001"}, fake.apiCalls) +} diff --git a/cmd/cli/cmd/org.go b/cmd/cli/cmd/org.go index 6657507..0c30a79 100644 --- a/cmd/cli/cmd/org.go +++ b/cmd/cli/cmd/org.go @@ -63,11 +63,7 @@ func orgInputsFor(host, flag, source string) auth.OrgInputs { // the choice is ambiguous and prompt is set, the user is asked and the // answer is remembered for host. func resolveOrgScope(ctx context.Context, host, token string, in auth.OrgInputs, prompt bool) (orgScope, error) { - listHost, err := auth.AuthHostFor(host) - if err != nil { - return orgScope{}, err - } - + listHost := host if usesMachineToken() { if strings.TrimSpace(in.Flag) == "" && strings.TrimSpace(in.EnvOrg) == "" { if id := strings.TrimSpace(in.EnvGitHubOwnerID); id != "" { @@ -76,11 +72,15 @@ func resolveOrgScope(ctx context.Context, host, token string, in auth.OrgInputs, return orgScope{Query: map[string]string{}}, nil } // The token is the tenant's, so its orgs are read from the tenant. - listHost = host in = auth.OrgInputs{Flag: in.Flag, FlagSource: in.FlagSource, EnvOrg: in.EnvOrg} + } else { + var err error + if listHost, err = auth.AuthHostFor(host); err != nil { + return orgScope{}, err + } } - list, err := listOrgs(ctx, host, listHost, token) + list, err := listOrgsFresh(ctx, host, listHost, token) if err != nil { return orgScope{}, err } @@ -88,8 +88,11 @@ func resolveOrgScope(ctx context.Context, host, token string, in auth.OrgInputs, org, err := auth.ResolveOrg(host, list, in) var notChosen *auth.ErrOrgNotChosen if errors.As(err, ¬Chosen) && prompt && !usesMachineToken() { + explicit := strings.TrimSpace(in.Flag) != "" || strings.TrimSpace(in.EnvOrg) != "" || strings.TrimSpace(in.EnvGitHubOwnerID) != "" org, err = promptForOrg(host, notChosen.Orgs) - if err == nil { + // An ambiguous --org answers this command only; the remembered pick + // changes when nothing named an org. + if err == nil && !explicit { if saveErr := rememberOrg(host, org); saveErr != nil { fmt.Fprintf(promptOut, "Could not remember %s for %s: %v\n", org.Name, host, saveErr) } else { @@ -103,6 +106,21 @@ func resolveOrgScope(ctx context.Context, host, token string, in auth.OrgInputs, return orgScope{Org: &org, Query: org.Query()}, nil } +// listOrgsFresh lists orgs with token, and once more with a force-refreshed +// token if the list rejects it: a token can stop being accepted before its +// stored expiry, and the org list is often a command's first call. +func listOrgsFresh(ctx context.Context, host, listHost, token string) (auth.OrgList, error) { + list, err := listOrgs(ctx, host, listHost, token) + if !errors.Is(err, auth.ErrTokenRejected) { + return list, err + } + fresh, refreshErr := lib.RefreshNullifyToken(ctx, host, nullifyToken, githubToken) + if refreshErr != nil || fresh == "" || fresh == token { + return list, err + } + return listOrgs(ctx, host, listHost, fresh) +} + func promptForOrg(host string, orgs []auth.Org) (auth.Org, error) { fmt.Fprintf(promptOut, "%s serves several of your organizations:\n", host) for i, o := range orgs { @@ -143,6 +161,9 @@ var orgCmd = &cobra.Command{ var orgListCmd = &cobra.Command{ Use: "list", Short: "List the organizations you can use on this host", + Long: "List the organizations your login can use on this host, those not routed to a host yet, and those on " + + "other hosts. \"selected\" is the organization a command run now would use (--org changes it), and " + + "\"selection\" says why none is when that is the case.", RunE: func(cmd *cobra.Command, args []string) error { ctx := cmd.Context() authCtx, err := resolveCommandToken(ctx) @@ -181,7 +202,7 @@ var orgUseCmd = &cobra.Command{ if err != nil { return err } - list, err := listOrgs(ctx, authCtx.Host, listHost, authCtx.Token) + list, err := listOrgsFresh(ctx, authCtx.Host, listHost, authCtx.Token) if err != nil { return err } @@ -209,7 +230,7 @@ func orgListing(ctx context.Context, host, token, flag string) (map[string]any, return nil, err } } - list, err := listOrgs(ctx, host, listHost, token) + list, err := listOrgsFresh(ctx, host, listHost, token) if err != nil { return nil, err } @@ -231,6 +252,9 @@ func orgListing(ctx context.Context, host, token, flag string) (map[string]any, out["selection"] = err.Error() case scope.Org != nil: out["selected"] = scope.Org.Name + case scope.Query["githubOwnerId"] != "": + out["selected"] = nil + out["selection"] = "githubOwnerId " + scope.Query["githubOwnerId"] + " (NULLIFY_GITHUB_OWNER_ID)" default: out["selected"] = nil out["selection"] = "the API token's own organization" diff --git a/cmd/cli/cmd/org_test.go b/cmd/cli/cmd/org_test.go index f165d96..da8ed9c 100644 --- a/cmd/cli/cmd/org_test.go +++ b/cmd/cli/cmd/org_test.go @@ -273,3 +273,35 @@ func TestAuthStatusShowsTheOrg(t *testing.T) { require.Contains(t, out, "Org: not chosen") require.Contains(t, out, "acme-payments") } + +// An --org that matches two orgs is settled at the prompt for this command +// only; it must not replace the remembered pick. +func TestAmbiguousExplicitOrgDoesNotChangeTheRememberedPick(t *testing.T) { + loginEnv(t) + dupGitHub := auth.Org{Name: "dup", Provider: "github", Key: "githubOwnerId", ID: "1", Host: "api.acme.nullify.ai"} + dupGitLab := auth.Org{Name: "dup", Provider: "gitlab", Key: "gitlabGroupId", ID: "2", Host: "api.acme.nullify.ai"} + stubOrgList(t, func(string) (auth.OrgList, error) { + return auth.OrgList{Orgs: []auth.Org{dupGitHub, dupGitLab, testOrgA}}, nil + }) + require.NoError(t, rememberOrg("acme.nullify.ai", testOrgA)) + oldIn, oldOut := promptIn, promptOut + promptIn, promptOut = strings.NewReader("2\n"), &bytes.Buffer{} + t.Cleanup(func() { promptIn, promptOut = oldIn, oldOut }) + + scope, err := resolveOrgScope(context.Background(), "acme.nullify.ai", "router-token", orgInputsFor("acme.nullify.ai", "dup", "--org"), true) + require.NoError(t, err) + require.Equal(t, dupGitLab, *scope.Org) + require.Equal(t, "acme-web", mustLoadHosts(t)["acme.nullify.ai"].Org.Name) +} + +// A corrupt hosts.json must not break an API-token command, which never +// reads it. +func TestAPITokenIgnoresHostsFile(t *testing.T) { + loginEnv(t) + t.Setenv("NULLIFY_TOKEN", "api-token") + require.NoError(t, os.WriteFile(os.Getenv("HOME")+"/.nullify/hosts.json", []byte("{broken"), 0600)) + + authCtx, err := resolveCommandAuth(context.Background()) + require.NoError(t, err) + require.Empty(t, authCtx.QueryParams) +} diff --git a/internal/auth/orgs.go b/internal/auth/orgs.go index 812ccd4..65c575d 100644 --- a/internal/auth/orgs.go +++ b/internal/auth/orgs.go @@ -98,7 +98,11 @@ func ListOrgs(ctx context.Context, host, listHost, token string) (OrgList, error if resp.StatusCode != http.StatusOK { body, _ := io.ReadAll(io.LimitReader(resp.Body, 1024)) - return OrgList{}, fmt.Errorf("list organizations on %s: status %d: %s", apiHost(listHost), resp.StatusCode, string(body)) + err := fmt.Errorf("list organizations on %s: status %d: %s", apiHost(listHost), resp.StatusCode, string(body)) + if resp.StatusCode == http.StatusUnauthorized { + err = fmt.Errorf("%w: %w", ErrTokenRejected, err) + } + return OrgList{}, err } var out installationsResponse @@ -220,6 +224,10 @@ func (e *ErrOrgNotChosen) Error() string { e.Host, len(e.Orgs), strings.Join(names, ", ")) } +// ErrTokenRejected reports that the org list refused the token, which can +// happen before its stored expiry (revocation, clock skew). +var ErrTokenRejected = errors.New("token rejected") + // ErrNoOrgs reports that the user has no org routed to this host. var ErrNoOrgs = errors.New("no organization on this host is available to you") diff --git a/internal/mcp/orgs.go b/internal/mcp/orgs.go index c23b900..a6e1ef6 100644 --- a/internal/mcp/orgs.go +++ b/internal/mcp/orgs.go @@ -3,8 +3,12 @@ package mcp import ( "context" "encoding/json" + "fmt" "maps" + "math" "net/http" + "strconv" + "strings" "github.com/mark3labs/mcp-go/mcp" "github.com/mark3labs/mcp-go/server" @@ -65,7 +69,11 @@ func orgMiddleware(orgs Orgs) server.ToolHandlerMiddleware { if req.Params.Name == listOrganizationsTool { return next(ctx, req) } - scope, err := orgs.Resolve(ctx, getStringArg(req.GetArguments(), orgArg)) + org, err := orgArgument(req.GetArguments()) + if err != nil { + return toolError(err), nil + } + scope, err := orgs.Resolve(ctx, org) if err != nil { return toolError(err), nil } @@ -74,6 +82,32 @@ func orgMiddleware(orgs Orgs) server.ToolHandlerMiddleware { } } +// orgArgument returns the call's org argument, "" when it has none. An org +// the caller did name but that is blank or not a name or id is an error: the +// call must not quietly run as the default org instead. +func orgArgument(args map[string]any) (string, error) { + v, ok := args[orgArg] + if !ok || v == nil { + return "", nil + } + var org string + switch t := v.(type) { + case string: + org = strings.TrimSpace(t) + case float64: + if t != math.Trunc(t) || t < 0 { + return "", fmt.Errorf("org argument %v is not an organization id", t) + } + org = strconv.FormatFloat(t, 'f', -1, 64) + default: + return "", fmt.Errorf("org argument must be an organization name or id, got %T", v) + } + if org == "" { + return "", fmt.Errorf("org argument is empty; omit it to use the default organization") + } + return org, nil +} + // resourceOrgMiddleware gives resource reads, which take no arguments, the // same org rule as a tool call with no org argument. func resourceOrgMiddleware(orgs Orgs) server.ResourceHandlerMiddleware { diff --git a/internal/mcp/orgs_test.go b/internal/mcp/orgs_test.go index e2cb446..f90e114 100644 --- a/internal/mcp/orgs_test.go +++ b/internal/mcp/orgs_test.go @@ -219,6 +219,25 @@ func TestUnknownOrgArgumentIsAToolError(t *testing.T) { require.Empty(t, rec.seen()) } +// An org the caller named must never turn into the default org: a numeric +// id is that id, and anything blank or of another type is refused. +func TestOrgArgumentIsTakenAsGivenOrRefused(t *testing.T) { + orgs := orgsFor(map[string]string{"acme-web": "1001001", "1001001": "1001001", "acme-payments": "2002002"}, "acme-payments") + + s, rec := newTestServer(t, orgs) + res := callTool(t, s, "list_repositories", map[string]any{"org": 1001001}) + require.False(t, res.Result.IsError, res.Result.Content) + require.Equal(t, "1001001", rec.seen()[0].Get("githubOwnerId"), "a numeric id is that org, not the default") + + for _, bad := range []any{"", " ", true, 1.5, map[string]any{"name": "acme-web"}, []any{"acme-web"}} { + s, rec := newTestServer(t, orgs) + res := callTool(t, s, "list_repositories", map[string]any{"org": bad}) + require.True(t, res.Result.IsError, "org %#v", bad) + require.Contains(t, res.Result.Content[0].Text, "org argument") + require.Empty(t, rec.seen(), "org %#v reached the API", bad) + } +} + func TestResourceReadFollowsTheSameRule(t *testing.T) { ambiguous, rec := newTestServer(t, orgsFor(map[string]string{"acme-web": "1001001", "acme-payments": "2002002"}, "")) res := call(t, ambiguous, "resources/read", map[string]any{"uri": "nullify://repos"})