feat(BRE2-1210): support new kas flow

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...")
 }

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 {

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")
+}