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.

acceptance/bin/profile_fingerprint.py

@@ -0,0 +1,26 @@
+#!/usr/bin/env python3
+
+import hashlib
+import sys
+
+
+def append_uvarint(data, value):
+    while value >= 0x80:
+        data.append((value & 0x7F) | 0x80)
+        value >>= 7
+    data.append(value)
+
+
+values = {
+    "auth_type": "databricks-cli",
+    "host": sys.argv[1],
+}
+serialized = bytearray()
+for key in sorted(values):
+    value = values[key]
+    append_uvarint(serialized, len(key.encode()))
+    serialized.extend(key.encode())
+    append_uvarint(serialized, len(value.encode()))
+    serialized.extend(value.encode())
+
+print(hashlib.sha256(serialized).hexdigest())

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.
+setup_test_profile
+setup_test_token_cache
+
+# Tokens created before profile fingerprinting require one new interactive login.
+jq 'del(.tokens["test-profile"].profile_fingerprint)' \
+    "./home/.databricks/token-cache.json" > "./token-cache.json"
+mv "./token-cache.json" "./home/.databricks/token-cache.json"
+
+trace musterr $CLI auth token --profile test-profile

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

@@ -0,0 +1,3 @@
+Ignore = [
+    "home",
+]

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,18 @@
+# A cached token must be rejected without mutation after its profile changes.
+sethome "./home"
+export BROWSER="browser.py"
+export DATABRICKS_AUTH_STORAGE=plaintext
+
+trace $CLI auth login --host $DATABRICKS_HOST --profile fingerprint-test --scopes jobs
+
+cp "./home/.databricks/token-cache.json" "./token-cache.before.json"
+
+# Editing any parsed profile value must block cached-token reuse.
+sed -i 's/scopes       = jobs/scopes       = all-apis,sql/' "./home/.databrickscfg"
+trace musterr $CLI auth token --profile fingerprint-test
+
+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,4 @@
+Ignore = [
+    "home",
+    "token-cache.before.json",
+]

acceptance/cmd/auth/token/script.prepare

@@ -15,6 +15,8 @@ ENDCFG
 
 setup_test_token_cache() {
     mkdir -p "./home/.databricks"
+    # Match the cache metadata that login writes for the two-field test profile.
+    profile_fingerprint=$(profile_fingerprint.py "$DATABRICKS_HOST_ORIG")
     cat > "./home/.databricks/token-cache.json" <<ENDCACHE
 {
   "version": 1,
@@ -23,7 +25,8 @@ setup_test_token_cache() {
       "access_token": "cached-access-token",
       "token_type": "Bearer",
       "refresh_token": "test-refresh-token",
-      "expiry": "2099-01-01T00:00:00Z"
+      "expiry": "2099-01-01T00:00:00Z",
+      "profile_fingerprint": "$profile_fingerprint"
     }
   }
 }

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

@@ -89,6 +89,24 @@ 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)
+	}
+
+	if err := storage.SetProfileFingerprint(tokenStore, profileName, savedProfile.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" {
@@ -384,7 +402,7 @@ a new profile is created.
 			// experimental_is_unified_host is no longer written to new profiles.
 			// Routing now comes from .well-known discovery; stale keys on existing
 			// profiles are cleaned up via clearKeys above.
-			err := databrickscfg.SaveToProfile(ctx, &config.Config{
+			profileConfig := &config.Config{
 				Profile:             profileName,
 				Host:                authArguments.Host,
 				AuthType:            authTypeDatabricksCLI,
@@ -394,11 +412,16 @@ a new profile is created.
 				ConfigFile:          env.Get(ctx, "DATABRICKS_CONFIG_FILE"),
 				ServerlessComputeID: serverlessComputeID,
 				Scopes:              scopesList,
-			}, clearKeys...)
+			}
+			err := databrickscfg.SaveToProfile(ctx, profileConfig, clearKeys...)
 			if err != nil {
 				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))
 		}
 
@@ -742,22 +765,28 @@ func discoveryLogin(ctx context.Context, in discoveryLoginInputs) error {
 		"cluster_id",
 		"serverless_compute_id",
 	)
-	err = databrickscfg.SaveToProfile(ctx, &config.Config{
+	profileConfig := &config.Config{
 		Profile:     in.profileName,
 		Host:        discoveredHost,
 		AuthType:    authTypeDatabricksCLI,
 		AccountID:   accountID,
 		WorkspaceID: workspaceID,
 		Scopes:      scopesList,
 		ConfigFile:  configFile,
-	}, clearKeys...)
+	}
+
+	err = databrickscfg.SaveToProfile(ctx, profileConfig, clearKeys...)
 	if err != nil {
 		if configFile != "" {
 			return fmt.Errorf("saving profile %q to %s: %w", in.profileName, configFile, err)
 		}
 		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

@@ -18,6 +18,7 @@ import (
 	"github.com/databricks/cli/libs/auth/storage"
 	"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/databricks/databricks-sdk-go/credentials/u2m"
@@ -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,46 @@ func TestDiscoveryLogin_IntrospectionFailureStillSavesProfile(t *testing.T) {
 	assert.Empty(t, savedProfile.WorkspaceID)
 }
 
+// TestDiscoveryLoginStoresSavedProfileFingerprint verifies that discovery login binds
+// its cached token to the profile values ultimately written to disk.
+func TestDiscoveryLoginStoresSavedProfileFingerprint(t *testing.T) {
+	home := t.TempDir()
+	configPath := filepath.Join(home, ".databrickscfg")
+	t.Setenv("HOME", home)
+	t.Setenv("USERPROFILE", home)
+	t.Setenv("DATABRICKS_CONFIG_FILE", "")
+
+	oauthArg, err := u2m.NewBasicDiscoveryOAuthArgument("DISCOVERY")
+	require.NoError(t, err)
+	oauthArg.SetDiscoveredHost("https://workspace.example.com")
+
+	tokenStore := newTestStore()
+	dc := &fakeDiscoveryClient{
+		oauthArg: oauthArg,
+		persistentAuth: &fakeDiscoveryPersistentAuth{
+			token: &oauth2.Token{AccessToken: "test-token"},
+		},
+		introspection: &auth.IntrospectionResult{},
+	}
+
+	ctx, _ := cmdio.NewTestContextWithStdout(t.Context())
+	err = discoveryLogin(ctx, discoveryLoginInputs{
+		dc:          dc,
+		profileName: "DISCOVERY",
+		timeout:     time.Second,
+		browserFunc: func(string) error { return nil },
+		tokenStore:  tokenStore,
+	})
+	require.NoError(t, err)
+
+	fingerprint, err := profilehash.FromFile(configPath, "DISCOVERY")
+	require.NoError(t, err)
+	entry, err := tokenStore.Lookup("DISCOVERY")
+	require.NoError(t, err)
+
+	assert.Equal(t, fingerprint, entry.ProfileFingerprint)
+}
+
 func TestDiscoveryLogin_AccountIDMismatchWarning(t *testing.T) {
 	tmpDir := t.TempDir()
 	configPath := filepath.Join(tmpDir, ".databrickscfg")

cmd/auth/token.go

@@ -264,20 +264,34 @@ func loadToken(ctx context.Context, args loadTokenArgs) (*oauth2.Token, error) {
 	if err != nil {
 		return nil, err
 	}
-	allArgs := append([]u2m.PersistentAuthOption{u2m.WithTokenCache(storage.OAuthTokenCache(ctx, args.tokenStore, args.mode))}, args.persistentAuthOpts...)
+
+	tokenStore := args.tokenStore
+	if existingProfile != nil {
+		if fingerprint := existingProfile.Fingerprint(); fingerprint != "" {
+			tokenStore = storage.NewProfileFingerprintStore(tokenStore, existingProfile.Name, fingerprint)
+		}
+	}
+
+	allArgs := append([]u2m.PersistentAuthOption{u2m.WithTokenCache(storage.OAuthTokenCache(ctx, tokenStore, args.mode))}, args.persistentAuthOpts...)
 	allArgs = append(allArgs, u2m.WithOAuthArgument(oauthArgument))
 	persistentAuth, err := u2m.NewPersistentAuth(ctx, allArgs...)
 	if err != nil {
 		helpMsg := helpfulError(ctx, args.profileName, oauthArgument)
 		return nil, fmt.Errorf("%w. %s", err, helpMsg)
 	}
+
 	var t *oauth2.Token
 	if args.forceRefresh {
 		t, err = persistentAuth.ForceRefreshToken()
 	} else {
 		t, err = persistentAuth.Token()
 	}
 	if err != nil {
+		// Fingerprint errors already include the exact login command needed to
+		// replace the stale grant, so the generic recovery suffix would duplicate it.
+		if errors.Is(err, storage.ErrProfileChanged) {
+			return nil, err
+		}
 		if errors.Is(err, cache.ErrNotFound) {
 			// The error returned by the SDK when the token cache doesn't exist or doesn't contain a token
 			// for the given host changed in SDK v0.77.0: https://github.com/databricks/databricks-sdk-go/pull/1250.
@@ -450,19 +464,24 @@ func runInlineLogin(ctx context.Context, profiler profile.Profiler, tokenStore s
 	clearKeys := oauthLoginClearKeys()
 	clearKeys = append(clearKeys, databrickscfg.ExperimentalIsUnifiedHostKey)
 
-	err = databrickscfg.SaveToProfile(ctx, &config.Config{
+	profileConfig := &config.Config{
 		Profile:     profileName,
 		Host:        loginArgs.Host,
 		AuthType:    authTypeDatabricksCLI,
 		AccountID:   loginArgs.AccountID,
 		WorkspaceID: loginArgs.WorkspaceID,
 		ConfigFile:  env.Get(ctx, "DATABRICKS_CONFIG_FILE"),
 		Scopes:      scopesList,
-	}, clearKeys...)
+	}
+	err = databrickscfg.SaveToProfile(ctx, profileConfig, clearKeys...)
 	if err != nil {
 		return "", nil, err
 	}
 
+	if err := setTokenProfileFingerprint(ctx, profiler, tokenStore, profileName); err != nil {
+		return "", nil, err
+	}
+
 	cmdio.LogString(ctx, fmt.Sprintf("Profile %s was successfully saved", profileName))
 
 	p, err := loadProfileByName(ctx, profileName, profiler)