From f7a28c2a59b031e9b20e59a563767a6b2a64ffd2 Mon Sep 17 00:00:00 2001 From: Pratik Patel Date: Fri, 9 Oct 2026 16:20:10 -0700 Subject: [PATCH] feat(BRE2-1210): support new kas flow --- pkg/auth/auth.go | 18 +++++--- pkg/auth/kas.go | 60 +++++++++++++++++++------ pkg/auth/kas_test.go | 103 +++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 162 insertions(+), 19 deletions(-) create mode 100644 pkg/auth/kas_test.go diff --git a/pkg/auth/auth.go b/pkg/auth/auth.go index e4f85a81d..dd3f9f365 100644 --- a/pkg/auth/auth.go +++ b/pkg/auth/auth.go @@ -370,12 +370,17 @@ func showLoginURL(url string) { fmt.Println(urlType(url)) } -func defaultAuthFunc(url, code string) { - codeType := color.New(color.FgWhite, color.Bold).SprintFunc() - if code != "" { - fmt.Println("Your Device Confirmation Code is 👉", codeType(code), "👈") - fmt.Print("\n") +func printDeviceCode(code string) { + if code == "" { + return } + codeType := color.New(color.FgWhite, color.Bold).SprintFunc() + fmt.Println("Your Device Confirmation Code is 👉", codeType(code), "👈") + fmt.Print("\n") +} + +func defaultAuthFunc(url, code string) { + printDeviceCode(code) // Best-effort: try to open the browser, but always show the URL below so // the user is never stranded if it doesn't open. @@ -384,7 +389,8 @@ func defaultAuthFunc(url, code string) { fmt.Println("\nWaiting for login to complete...") } -func skipBrowserAuthFunc(url, _ string) { +func skipBrowserAuthFunc(url, code string) { + printDeviceCode(code) showLoginURL(url) fmt.Println("\nWaiting for login to complete...") } diff --git a/pkg/auth/kas.go b/pkg/auth/kas.go index 4af3240dc..c2d550dab 100644 --- a/pkg/auth/kas.go +++ b/pkg/auth/kas.go @@ -75,12 +75,13 @@ func (a KasAuthenticator) GetNewAuthTokensWithRefresh(refreshToken string) (*ent } type LoginCallResponse struct { - LoginURL string `json:"loginUrl"` - SessionKey string `json:"sessionKey"` + UserCode string `json:"user_code"` + SessionKey string `json:"session_key"` + VerificationURL string `json:"verification_url"` } func (a KasAuthenticator) MakeLoginCall(id, email string) (LoginCallResponse, error) { - url := fmt.Sprintf("%s/device/login", a.BaseURL) + url := fmt.Sprintf("%s/v2/device/login", a.BaseURL) payload := map[string]string{ "email": email, "deviceId": id, @@ -137,9 +138,9 @@ func (a KasAuthenticator) DoDeviceAuthFlow(userLoginFlow func(url string, code s return nil, breverrors.WrapAndTrace(err) } - userLoginFlow(loginResp.LoginURL, "") + userLoginFlow(loginResp.VerificationURL, loginResp.UserCode) - return a.pollForTokens(loginResp.SessionKey, id) + return a.pollForAuthentication(loginResp.SessionKey) } func (a KasAuthenticator) maybePromptForEmail() (string, error) { @@ -169,8 +170,11 @@ func (a KasAuthenticator) maybePromptForEmail() (string, error) { return email, nil } -func (a KasAuthenticator) pollForTokens(sessionKey, id string) (*LoginTokens, error) { - // Try to retrieve tokens for up to 5 minutes +// pollForAuthentication waits until the user completes login in the browser. +// While the session is pending, HEAD /v2/ping returns 401; once the user +// authenticates, it returns 200 and the session_key becomes the bearer token. +func (a KasAuthenticator) pollForAuthentication(sessionKey string) (*LoginTokens, error) { + // Try to authenticate for up to PollTimeout timeout := time.After(a.PollTimeout) ticker := time.NewTicker(2 * time.Second) defer ticker.Stop() @@ -179,21 +183,51 @@ func (a KasAuthenticator) pollForTokens(sessionKey, id string) (*LoginTokens, er case <-timeout: return nil, breverrors.WrapAndTrace(fmt.Errorf("timed out waiting for login")) case <-ticker.C: - idToken, err := a.retrieveIDToken(sessionKey, id) - if err == nil { + authenticated, err := a.isSessionAuthenticated(sessionKey) + if err != nil { + // Transient errors (network blips, 5xx) are retried until timeout. + continue + } + if authenticated { return &LoginTokens{ AuthTokens: entity.AuthTokens{ - AccessToken: idToken, - RefreshToken: fmt.Sprintf("%s:%s", sessionKey, id), + AccessToken: sessionKey, }, - IDToken: idToken, + IDToken: sessionKey, }, nil } - // Continue polling on error } } } +// isSessionAuthenticated reports whether the session_key is now a valid bearer +// token by calling HEAD /v2/ping. 200 means authenticated; 401 means the user +// has not completed login yet. +func (a KasAuthenticator) isSessionAuthenticated(sessionKey string) (bool, error) { + pingURL := fmt.Sprintf("%s/v2/ping", a.BaseURL) + req, err := http.NewRequest(http.MethodHead, pingURL, nil) //nolint:noctx // fine + if err != nil { + return false, breverrors.WrapAndTrace(err) + } + req.Header.Set("Authorization", "Bearer "+sessionKey) + + client := &http.Client{} + resp, err := client.Do(req) + if err != nil { + return false, breverrors.WrapAndTrace(err) + } + defer resp.Body.Close() //nolint:errcheck // fine + + switch resp.StatusCode { + case http.StatusOK: + return true, nil + case http.StatusUnauthorized: + return false, nil + default: + return false, fmt.Errorf("unexpected ping status code: %d", resp.StatusCode) + } +} + type RetrieveIDTokenResponse struct { IDToken string `json:"token"` RequestStatus struct { diff --git a/pkg/auth/kas_test.go b/pkg/auth/kas_test.go new file mode 100644 index 000000000..72a9b2f9e --- /dev/null +++ b/pkg/auth/kas_test.go @@ -0,0 +1,103 @@ +package auth + +import ( + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestKasDoDeviceAuthFlow_CallsV2LoginThenPollsPing(t *testing.T) { + const ( + userCode = "ABCD-1234" + session = "session-key-abc" + verify = "https://brev.nvidia.com/cli-login?code=xyz" + ) + + var loginCalls, pingCalls int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch { + case r.Method == http.MethodPost && r.URL.Path == "/v2/device/login": + atomic.AddInt32(&loginCalls, 1) + var body map[string]string + require.NoError(t, json.NewDecoder(r.Body).Decode(&body)) + assert.Equal(t, "user@example.com", body["email"]) + assert.Equal(t, "https://brev.nvidia.com/cli-login", body["redirectUri"]) + assert.NotEmpty(t, body["deviceId"]) + + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]string{ + "user_code": userCode, + "session_key": session, + "verification_url": verify, + }) + case r.Method == http.MethodHead && r.URL.Path == "/v2/ping": + call := atomic.AddInt32(&pingCalls, 1) + assert.Equal(t, "Bearer "+session, r.Header.Get("Authorization")) + if call < 2 { + w.WriteHeader(http.StatusUnauthorized) + return + } + w.WriteHeader(http.StatusOK) + default: + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + w.WriteHeader(http.StatusNotFound) + } + })) + defer srv.Close() + + a := NewKasAuthenticator("user@example.com", srv.URL, "https://login.nvidia.com", false, "https://brev.nvidia.com/cli-login") + + var gotURL, gotCode string + tokens, err := a.DoDeviceAuthFlow(func(url, code string) { + gotURL, gotCode = url, code + }) + require.NoError(t, err) + + assert.Equal(t, verify, gotURL) + assert.Equal(t, userCode, gotCode) + assert.Equal(t, session, tokens.AccessToken, "session_key must be usable as the bearer token") + assert.Empty(t, tokens.RefreshToken, "v2 login does not issue a refresh token") + assert.Equal(t, int32(1), atomic.LoadInt32(&loginCalls)) + assert.GreaterOrEqual(t, atomic.LoadInt32(&pingCalls), int32(2), "must keep polling until ping succeeds") +} + +func TestKasDoDeviceAuthFlow_TimesOutWhenPingNeverSucceeds(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodPost && r.URL.Path == "/v2/device/login" { + w.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(w, `{"user_code":"c","session_key":"s","verification_url":"u"}`) + return + } + w.WriteHeader(http.StatusUnauthorized) + })) + defer srv.Close() + + a := NewKasAuthenticator("user@example.com", srv.URL, "https://login.nvidia.com", false, "") + a.PollTimeout = 50 * time.Millisecond + + _, err := a.DoDeviceAuthFlow(func(_, _ string) {}) + require.Error(t, err) + assert.Contains(t, err.Error(), "timed out") +} + +func TestKasDoDeviceAuthFlow_LoginErrorPropagates(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusInternalServerError) + _, _ = io.WriteString(w, "boom") + })) + defer srv.Close() + + a := NewKasAuthenticator("user@example.com", srv.URL, "https://login.nvidia.com", false, "") + + flowCalled := false + _, err := a.DoDeviceAuthFlow(func(_, _ string) { flowCalled = true }) + require.Error(t, err) + assert.False(t, flowCalled, "login flow must not start when the device login call fails") +}