Invalidate cached OAuth tokens when profiles change

.nextchanges/cli/profile-fingerprint.md

@@ -0,0 +1 @@
+* Require a new OAuth login before reusing cached credentials after the corresponding profile configuration changes. ([#6427](https://github.com/databricks/cli/pull/6427))

acceptance/cmd/auth/logout/stale-account-id-workspace-host/output.txt

@@ -19,9 +19,9 @@ Logged out of profile "logfood". Use --delete to also remove it from the config
 [DEFAULT]
 
 [logfood]
-host = [DATABRICKS_URL]
+host         = [DATABRICKS_URL]
 account_id = stale-account
-auth_type  = databricks-cli
+auth_type    = databricks-cli
 
 [__settings__]
 default_profile = logfood

acceptance/cmd/auth/logout/stale-account-id-workspace-host/script

@@ -1,44 +1,40 @@
 sethome "./home"
+export BROWSER="browser.py"
+export DATABRICKS_AUTH_STORAGE=plaintext
 
+# Setup: create a normal workspace profile and log in.
 cat > "./home/.databrickscfg" <<EOF
 ; The profile defined in the DEFAULT section is to be used as a fallback when no profile is explicitly specified.
 [DEFAULT]
 
 [logfood]
 host = ${DATABRICKS_HOST}
-account_id = stale-account
 auth_type  = databricks-cli
 
 [__settings__]
 default_profile = logfood
 EOF
 
-mkdir -p "./home/.databricks"
-cat > "./home/.databricks/token-cache.json" <<EOF
-{
-  "version": 1,
-  "tokens": {
-    "logfood": {
-      "access_token": "logfood-cached-token",
-      "token_type": "Bearer"
-    },
-    "${DATABRICKS_HOST}": {
-      "access_token": "logfood-host-token",
-      "token_type": "Bearer"
-    }
-  }
-}
-EOF
+$CLI auth login --host "$DATABRICKS_HOST" --profile logfood >/dev/null 2>&1
 
+# Confirm that login created a valid profile and both cached token entries.
 title "Profiles before logout — logfood should be valid"
 trace $CLI auth profiles
 
 title "Token cache keys before logout\n"
 jq -S '.tokens | keys' "./home/.databricks/token-cache.json"
 
+# Finish setup by adding the stale account routing that caused the regression.
+sed -i.bak '/^host/a\
+account_id = stale-account
+' "./home/.databrickscfg"
+sed -i.bak '/workspace_id/d' "./home/.databrickscfg"
+
+# Test: logout must still find the tokens using the workspace host.
 title "Logout without --delete"
 trace $CLI auth logout --profile logfood --auto-approve
 
+# Validation: logout preserves the profile and removes both cached entries.
 title "Config after logout — profile should still exist\n"
 cat "./home/.databrickscfg"
 

acceptance/cmd/auth/profile-change/output.txt

@@ -0,0 +1,6 @@
+
+>>> [CLI] auth login --host [DATABRICKS_URL] --profile fingerprint-test --scopes jobs
+Profile fingerprint-test was successfully saved
+
+>>> musterr [CLI] clusters list --profile fingerprint-test
+Error: error getting token: cache: profile "fingerprint-test" has changed since the last login; run `databricks auth login --profile "fingerprint-test"` to sign in again

acceptance/cmd/auth/profile-change/script

@@ -0,0 +1,13 @@
+# Ordinary workspace commands must reject a cached token after its profile changes.
+sethome "./home"
+export BROWSER="browser.py"
+export DATABRICKS_AUTH_STORAGE=plaintext
+
+# Login creates a cached token bound to the saved profile.
+trace $CLI auth login --host $DATABRICKS_HOST --profile fingerprint-test --scopes jobs
+
+# Simulate changing an authentication-related profile field by hand.
+sed -i.bak 's/scopes       = jobs/scopes       = all-apis,sql/' "./home/.databrickscfg"
+
+# An ordinary workspace command must reject the token bound to the old profile.
+trace musterr $CLI clusters list --profile fingerprint-test

acceptance/cmd/auth/profile-change/test.toml

@@ -0,0 +1,5 @@
+EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"]
+
+Ignore = [
+    "home",
+]

acceptance/cmd/auth/token/force-refresh-invalid-refresh-token/script

@@ -1,4 +1,11 @@
+# Create the profile and a valid profile-bound token that can be refreshed.
 setup_test_profile
 setup_test_token_cache
 
+# Replace the valid refresh token with one the fake OIDC server rejects.
+jq '.tokens["test-profile"].refresh_token = "invalid-refresh-token"' \
+    "./home/.databricks/token-cache.json" > "./token-cache.json"
+mv "./token-cache.json" "./home/.databricks/token-cache.json"
+
+# The CLI must report that the user needs to log in again.
 musterr $CLI auth token --profile test-profile --force-refresh

acceptance/cmd/auth/token/force-refresh-success/script

@@ -1,3 +1,4 @@
+# Create the profile and a valid profile-bound token that can be refreshed.
 setup_test_profile
 setup_test_token_cache
 

acceptance/cmd/auth/token/legacy-profile-fingerprint/output.txt

@@ -0,0 +1,3 @@
+
+>>> musterr [CLI] auth token --profile test-profile
+Error: cache: cached credentials for profile "test-profile" predate profile change detection; run `databricks auth login --profile "test-profile"` to sign in again

acceptance/cmd/auth/token/legacy-profile-fingerprint/script

@@ -0,0 +1,10 @@
+# A token without fingerprint metadata must require a new login.
+# Create the profile that the legacy cached token belongs to.
+setup_test_profile
+
+# Install a cache written by an older CLI that omits profile_fingerprint.
+mkdir -p "./home/.databricks"
+cp "./token-cache.json" "./home/.databricks/token-cache.json"
+
+# The CLI must reject a cached token whose profile binding cannot be verified.
+trace musterr $CLI auth token --profile test-profile

acceptance/cmd/auth/token/legacy-profile-fingerprint/test.toml

@@ -0,0 +1 @@
+EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"]

acceptance/cmd/auth/token/legacy-profile-fingerprint/token-cache.json

@@ -0,0 +1,11 @@
+{
+  "version": 1,
+  "tokens": {
+    "test-profile": {
+      "access_token": "cached-access-token",
+      "token_type": "Bearer",
+      "refresh_token": "test-refresh-token",
+      "expiry": "2099-01-01T00:00:00Z"
+    }
+  }
+}

acceptance/cmd/auth/token/profile-change/output.txt

@@ -0,0 +1,7 @@
+
+>>> [CLI] auth login --host [DATABRICKS_URL] --profile fingerprint-test --scopes jobs
+Profile fingerprint-test was successfully saved
+
+>>> musterr [CLI] auth token --profile fingerprint-test
+Error: cache: profile "fingerprint-test" has changed since the last login; run `databricks auth login --profile "fingerprint-test"` to sign in again
+Token cache unchanged

acceptance/cmd/auth/token/profile-change/script

@@ -0,0 +1,23 @@
+# A cached token must be rejected without mutation after its profile changes.
+sethome "./home"
+export BROWSER="browser.py"
+export DATABRICKS_AUTH_STORAGE=plaintext
+
+# Login creates a cached token bound to the saved profile.
+trace $CLI auth login --host $DATABRICKS_HOST --profile fingerprint-test --scopes jobs
+
+# Keep the original cache so we can verify that rejection does not modify it.
+cp "./home/.databricks/token-cache.json" "./token-cache.before.json"
+
+# Simulate changing an authentication-related profile field by hand.
+sed -i.bak 's/scopes       = jobs/scopes       = all-apis,sql/' "./home/.databrickscfg"
+
+# auth token must reject the token bound to the old profile.
+trace musterr $CLI auth token --profile fingerprint-test
+
+# Rejecting the stale token must leave the cached credentials unchanged.
+if cmp -s "./token-cache.before.json" "./home/.databricks/token-cache.json"; then
+    echo "Token cache unchanged"
+else
+    echo "Token cache changed"
+fi

acceptance/cmd/auth/token/profile-change/test.toml

@@ -0,0 +1,6 @@
+EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"]
+
+Ignore = [
+    "home",
+    "token-cache.before.json",
+]

acceptance/cmd/auth/token/script.prepare

@@ -1,11 +1,13 @@
 setup_test_profile() {
+    # Preserve the test server host before clearing ambient profile selection.
     export DATABRICKS_HOST_ORIG="$DATABRICKS_HOST"
 
     sethome "./home"
     unset DATABRICKS_HOST
     unset DATABRICKS_TOKEN
     unset DATABRICKS_CONFIG_PROFILE
 
+    # Write the minimal OAuth profile used by token-cache tests.
     cat > "./home/.databrickscfg" <<ENDCFG
 [test-profile]
 host = $DATABRICKS_HOST_ORIG
@@ -14,18 +16,17 @@ ENDCFG
 }
 
 setup_test_token_cache() {
-    mkdir -p "./home/.databricks"
-    cat > "./home/.databricks/token-cache.json" <<ENDCACHE
-{
-  "version": 1,
-  "tokens": {
-    "test-profile": {
-      "access_token": "cached-access-token",
-      "token_type": "Bearer",
-      "refresh_token": "test-refresh-token",
-      "expiry": "2099-01-01T00:00:00Z"
-    }
-  }
-}
-ENDCACHE
+    export BROWSER="browser.py"
+    export DATABRICKS_AUTH_STORAGE=plaintext
+
+    # Use login so the cached token contains a real binding to the saved profile.
+    $CLI auth login --host "$DATABRICKS_HOST_ORIG" --profile test-profile >/dev/null 2>&1
+
+    # Change only the token fields needed to control refresh behavior in each test.
+    jq '.tokens["test-profile"] += {
+        "access_token": "cached-access-token",
+        "refresh_token": "test-refresh-token",
+        "expiry": "2099-01-01T00:00:00Z"
+    }' "./home/.databricks/token-cache.json" > "./token-cache.json"
+    mv "./token-cache.json" "./home/.databricks/token-cache.json"
 }

cmd/auth/in_memory_test.go

@@ -6,7 +6,8 @@ import (
 )
 
 type inMemoryStore struct {
-	Tokens map[string]*oauth2.Token
+	Tokens       map[string]*oauth2.Token
+	Fingerprints map[string]string
 }
 
 // Lookup returns a copy to match real (file-backed) cache behavior, where
@@ -19,21 +20,32 @@ func (i *inMemoryStore) Lookup(key string) (storage.Entry, error) {
 		return storage.Entry{}, storage.ErrNotFound
 	}
 	cp := *token
-	return storage.Entry{Token: &cp}, nil
+	return storage.Entry{
+		Token:              &cp,
+		ProfileFingerprint: i.Fingerprints[key],
+	}, nil
 }
 
 // Put stores a copy to prevent callers from mutating cached entries after
 // put returns (mirrors file-backed cache semantics).
 func (i *inMemoryStore) Put(key string, e storage.Entry) error {
 	cp := *e.Token
 	i.Tokens[key] = &cp
+
+	if i.Fingerprints == nil {
+		i.Fingerprints = make(map[string]string)
+	}
+
+	i.Fingerprints[key] = e.ProfileFingerprint
+
 	return nil
 }
 
 // Delete deletes the entry under key. Deleting a missing entry is not
 // an error.
 func (i *inMemoryStore) Delete(key string) error {
 	delete(i.Tokens, key)
+	delete(i.Fingerprints, key)
 	return nil
 }
 

cmd/auth/login.go

@@ -17,6 +17,7 @@ import (
 	"github.com/databricks/cli/libs/databrickscfg"
 	"github.com/databricks/cli/libs/databrickscfg/cfgpickers"
 	"github.com/databricks/cli/libs/databrickscfg/profile"
+	"github.com/databricks/cli/libs/databrickscfg/profilehash"
 	"github.com/databricks/cli/libs/env"
 	"github.com/databricks/cli/libs/log"
 	"github.com/databricks/databricks-sdk-go"
@@ -88,6 +89,29 @@ func (d *defaultDiscoveryClient) IntrospectToken(ctx context.Context, host, acce
 	return auth.IntrospectToken(ctx, host, accessToken, nil)
 }
 
+// setTokenProfileFingerprint runs after profile saving because OAuth-dependent
+// workspace and compute selection can change the final profile contents.
+func setTokenProfileFingerprint(ctx context.Context, profiler profile.Profiler, tokenStore storage.Store, profileName string) error {
+	savedProfile, err := loadProfileByName(ctx, profileName, profiler)
+	if err != nil {
+		return fmt.Errorf("load saved profile %q: %w", profileName, err)
+	}
+	if savedProfile == nil {
+		return fmt.Errorf("saved profile %q not found", profileName)
+	}
+
+	fingerprint, err := profilehash.Compute(*savedProfile)
+	if err != nil {
+		return fmt.Errorf("compute profile fingerprint: %w", err)
+	}
+
+	if err := storage.SetProfileFingerprint(tokenStore, profileName, fingerprint); err != nil {
+		return fmt.Errorf("save profile fingerprint: %w", err)
+	}
+
+	return nil
+}
+
 func newLoginCommand(authArguments *auth.AuthArguments) *cobra.Command {
 	defaultConfigPath := "~/.databrickscfg"
 	if runtime.GOOS == "windows" {
@@ -398,6 +422,10 @@ a new profile is created.
 				return err
 			}
 
+			if err := setTokenProfileFingerprint(ctx, profile.DefaultProfiler, tokenStore, profileName); err != nil {
+				return err
+			}
+
 			cmdio.LogString(ctx, fmt.Sprintf("Profile %s was successfully saved", profileName))
 		}
 
@@ -758,6 +786,10 @@ func discoveryLogin(ctx context.Context, in discoveryLoginInputs) error {
 		return fmt.Errorf("saving profile %q: %w", in.profileName, err)
 	}
 
+	if err := setTokenProfileFingerprint(ctx, profile.DefaultProfiler, in.tokenStore, in.profileName); err != nil {
+		return err
+	}
+
 	cmdio.LogString(ctx, fmt.Sprintf("Profile %s was successfully saved", in.profileName))
 	return nil
 }

cmd/auth/login_test.go

@@ -19,6 +19,7 @@ import (
 	"github.com/databricks/cli/libs/auth/u2m"
 	"github.com/databricks/cli/libs/cmdio"
 	"github.com/databricks/cli/libs/databrickscfg/profile"
+	"github.com/databricks/cli/libs/databrickscfg/profilehash"
 	"github.com/databricks/cli/libs/env"
 	"github.com/databricks/cli/libs/log"
 	"github.com/spf13/cobra"
@@ -30,7 +31,10 @@ import (
 // newTestStore returns an in-memory token cache for tests so that
 // discoveryLogin and other login helpers don't touch ~/.databricks/token-cache.json.
 func newTestStore() storage.Store {
-	return &inMemoryStore{Tokens: map[string]*oauth2.Token{}}
+	// Prepopulate the entry because the fake Challenge does not perform the real OAuth cache write.
+	return &inMemoryStore{Tokens: map[string]*oauth2.Token{
+		"DISCOVERY": {AccessToken: "test-token"},
+	}}
 }
 
 // logBuffer is a thread-safe bytes.Buffer for capturing log output in tests.
@@ -806,6 +810,29 @@ func TestDiscoveryLogin_IntrospectionFailureStillSavesProfile(t *testing.T) {
 	assert.Empty(t, savedProfile.WorkspaceID)
 }
 
+// TestSetTokenProfileFingerprint verifies that a cached token is bound to the
+// profile saved by login.
+func TestSetTokenProfileFingerprint(t *testing.T) {
+	savedProfile := profile.Profile{
+		Name: "DISCOVERY",
+		Host: "https://workspace.example.test",
+	}
+	profiler := profile.InMemoryProfiler{
+		Profiles: profile.Profiles{savedProfile},
+	}
+	tokenStore := newTestStore()
+
+	err := setTokenProfileFingerprint(t.Context(), profiler, tokenStore, savedProfile.Name)
+	require.NoError(t, err)
+
+	want, err := profilehash.Compute(savedProfile)
+	require.NoError(t, err)
+	entry, err := tokenStore.Lookup(savedProfile.Name)
+	require.NoError(t, err)
+
+	assert.Equal(t, want, entry.ProfileFingerprint)
+}
+
 func TestDiscoveryLogin_AccountIDMismatchWarning(t *testing.T) {
 	tmpDir := t.TempDir()
 	configPath := filepath.Join(tmpDir, ".databrickscfg")