Keep separate environment logins from breaking each other

internal/authutil/context.go

@@ -18,6 +18,9 @@ func GetUserKeyForCurrentSession() (string, *datumconfig.Session, error) {
 	if err != nil {
 		return "", nil, err
 	}
+	if err := EnsureUserKeysMigrated(cfg); err != nil {
+		return "", nil, err
+	}
 
 	if session := cfg.ActiveSessionEntry(); session != nil && session.UserKey != "" {
 		return session.UserKey, session, nil
@@ -49,6 +52,9 @@ func GetUserKeyForSession(sessionName string) (string, error) {
 	if err != nil {
 		return "", err
 	}
+	if err := EnsureUserKeysMigrated(cfg); err != nil {
+		return "", err
+	}
 	session := cfg.SessionByName(sessionName)
 	if session == nil {
 		return "", fmt.Errorf("no session named %q — run 'datumctl login' or 'datumctl auth update-kubeconfig'", sessionName)

internal/authutil/login.go

@@ -270,9 +270,10 @@ func completeLogin(ctx context.Context, provider *oidc.Provider, clientID, authH
 
 	fmt.Fprintf(out, "\nAuthenticated as: %s (%s)\n", claims.Name, claims.Email)
 
-	// Use email as the keyring key for interactive logins, matching the
-	// behaviour of the original cmd/auth/login.go implementation.
-	userKey := claims.Email
+	// Key the keyring entry by email *and* auth hostname so that the same email
+	// logged into different environments (e.g. production and staging) does not
+	// share one entry, where the second login would clobber the first.
+	userKey := userKeyFor(claims.Email, authHostname)
 
 	creds := StoredCredentials{
 		Hostname:         authHostname,

internal/authutil/migrate.go

@@ -0,0 +1,130 @@
+package authutil
+
+import (
+	"encoding/json"
+	"errors"
+	"strings"
+
+	"go.datum.net/datumctl/internal/datumconfig"
+	"go.datum.net/datumctl/internal/keyring"
+)
+
+// userKeyFor returns the keyring key that identifies a stored credential.
+// Credentials are keyed by identity *and* auth hostname so that logging in with
+// the same identity (e.g. email) in two environments — production and staging —
+// does not collapse onto a single keyring entry where the second login silently
+// overwrites the first.
+func userKeyFor(identity, authHostname string) string {
+	if identity == "" || authHostname == "" {
+		return identity
+	}
+	return identity + "@" + authHostname
+}
+
+// EnsureUserKeysMigrated runs MigrateUserKeys on cfg and persists the config when
+// anything changed. It is safe to call on every auth resolution: when there is
+// nothing to migrate it performs no keyring writes and no disk writes.
+func EnsureUserKeysMigrated(cfg *datumconfig.ConfigV1Beta1) error {
+	changed, err := MigrateUserKeys(cfg)
+	if err != nil {
+		return err
+	}
+	if changed {
+		return datumconfig.SaveV1Beta1(cfg)
+	}
+	return nil
+}
+
+// MigrateUserKeys upgrades legacy sessions whose keyring key was scoped only by
+// identity (an email address, historically) to the environment-scoped scheme
+// "<identity>@<auth-hostname>".
+//
+// Before this scheme, two logins with the same email in different environments
+// shared one keyring entry, so the second login silently clobbered the first —
+// and switching between the sessions presented a token minted for the wrong
+// environment, producing spurious "unauthorized" errors.
+//
+// For each legacy session it computes the new key and, when the single stored
+// blob was minted for that session's auth hostname, copies the blob to the new
+// key. When two sessions collided only the environment whose token currently
+// occupies the blob is recovered; the other session is left without a stored
+// token and must be re-authenticated with `datumctl login`.
+//
+// The old keyring entries are left untouched — they become inert once no session
+// references them — so the migration never destroys a token it could not first
+// copy. It reports whether cfg was modified; callers persist cfg when true.
+func MigrateUserKeys(cfg *datumconfig.ConfigV1Beta1) (bool, error) {
+	if cfg == nil {
+		return false, nil
+	}
+
+	changed := false
+	for i := range cfg.Sessions {
+		s := &cfg.Sessions[i]
+		authHost := s.Endpoint.AuthHostname
+		if authHost == "" || s.UserKey == "" {
+			continue
+		}
+		if strings.HasSuffix(s.UserKey, "@"+authHost) {
+			continue // already environment-scoped
+		}
+
+		identity := s.UserEmail
+		if identity == "" {
+			identity = s.UserKey
+		}
+		newKey := userKeyFor(identity, authHost)
+		if newKey == s.UserKey {
+			continue
+		}
+
+		if err := rehomeCredential(s.UserKey, newKey, authHost); err != nil {
+			return changed, err
+		}
+		s.UserKey = newKey
+		changed = true
+	}
+
+	return changed, nil
+}
+
+// rehomeCredential copies the credential stored under oldKey to newKey when the
+// stored blob was minted for authHost and newKey does not already hold a
+// credential. A missing, unparseable, or mismatched blob is not an error: the
+// session simply starts out needing a fresh login.
+func rehomeCredential(oldKey, newKey, authHost string) error {
+	raw, err := keyring.Get(ServiceName, oldKey)
+	if err != nil {
+		if errors.Is(err, keyring.ErrNotFound) {
+			return nil
+		}
+		return err
+	}
+	if raw == "" {
+		return nil
+	}
+
+	var creds StoredCredentials
+	if err := json.Unmarshal([]byte(raw), &creds); err != nil {
+		// Unparseable legacy blob: don't block migration, just skip the copy.
+		return nil
+	}
+
+	// Only claim the blob when it belongs to this session's environment. With a
+	// cross-environment collision exactly one session matches; the others need a
+	// fresh login. Blobs with no recorded hostname are assumed to match.
+	if creds.Hostname != "" && creds.Hostname != authHost {
+		return nil
+	}
+
+	if existing, err := keyring.Get(ServiceName, newKey); err == nil && existing != "" {
+		return nil // already migrated
+	} else if err != nil && !errors.Is(err, keyring.ErrNotFound) {
+		return err
+	}
+
+	if err := keyring.Set(ServiceName, newKey, raw); err != nil {
+		return err
+	}
+	return AddKnownUserKey(newKey)
+}

internal/authutil/migrate_test.go

@@ -0,0 +1,138 @@
+package authutil
+
+import (
+	"encoding/json"
+	"testing"
+
+	"go.datum.net/datumctl/internal/datumconfig"
+	"go.datum.net/datumctl/internal/keyring"
+	"golang.org/x/oauth2"
+)
+
+func TestUserKeyFor(t *testing.T) {
+	cases := []struct {
+		identity, host, want string
+	}{
+		{"swells@datum.net", "auth.datum.net", "swells@datum.net@auth.datum.net"},
+		{"swells@datum.net", "auth.staging.env.datum.net", "swells@datum.net@auth.staging.env.datum.net"},
+		{"swells@datum.net", "", "swells@datum.net"}, // no host to qualify with
+		{"", "auth.datum.net", ""}, // no identity
+	}
+	for _, c := range cases {
+		if got := userKeyFor(c.identity, c.host); got != c.want {
+			t.Errorf("userKeyFor(%q, %q) = %q, want %q", c.identity, c.host, got, c.want)
+		}
+	}
+}
+
+// storeCreds writes a StoredCredentials blob to the keyring under key.
+func storeCreds(t *testing.T, key, authHostname string) {
+	t.Helper()
+	blob, err := json.Marshal(StoredCredentials{
+		Hostname:  authHostname,
+		UserEmail: "swells@datum.net",
+		Subject:   "sub-" + authHostname,
+		Token:     &oauth2.Token{AccessToken: "tok-" + authHostname},
+	})
+	if err != nil {
+		t.Fatalf("marshal creds: %v", err)
+	}
+	if err := keyring.Set(ServiceName, key, string(blob)); err != nil {
+		t.Fatalf("seed keyring: %v", err)
+	}
+}
+
+// collidingConfig returns a config with a staging and production session that
+// share the legacy bare-email user key, mirroring the real-world bug.
+func collidingConfig() *datumconfig.ConfigV1Beta1 {
+	return &datumconfig.ConfigV1Beta1{
+		Sessions: []datumconfig.Session{
+			{
+				Name:      "swells@datum.net@api.staging.env.datum.net",
+				UserKey:   "swells@datum.net",
+				UserEmail: "swells@datum.net",
+				Endpoint:  datumconfig.Endpoint{AuthHostname: "auth.staging.env.datum.net"},
+			},
+			{
+				Name:      "swells@datum.net@api.datum.net",
+				UserKey:   "swells@datum.net",
+				UserEmail: "swells@datum.net",
+				Endpoint:  datumconfig.Endpoint{AuthHostname: "auth.datum.net"},
+			},
+		},
+	}
+}
+
+func TestMigrateUserKeys_CollisionRecoversMatchingEnvironment(t *testing.T) {
+	keyring.MockInit()
+
+	// The single shared blob currently holds the staging token (last written).
+	storeCreds(t, "swells@datum.net", "auth.staging.env.datum.net")
+
+	cfg := collidingConfig()
+	changed, err := MigrateUserKeys(cfg)
+	if err != nil {
+		t.Fatalf("MigrateUserKeys: %v", err)
+	}
+	if !changed {
+		t.Fatal("expected changed=true")
+	}
+
+	stagingKey := "swells@datum.net@auth.staging.env.datum.net"
+	prodKey := "swells@datum.net@auth.datum.net"
+
+	if got := cfg.Sessions[0].UserKey; got != stagingKey {
+		t.Errorf("staging session UserKey = %q, want %q", got, stagingKey)
+	}
+	if got := cfg.Sessions[1].UserKey; got != prodKey {
+		t.Errorf("prod session UserKey = %q, want %q", got, prodKey)
+	}
+
+	// The surviving (staging) token is re-homed under the new key...
+	if _, err := GetStoredCredentials(stagingKey); err != nil {
+		t.Errorf("expected staging creds under new key, got error: %v", err)
+	}
+	// ...while the clobbered (production) session has no token and must re-login.
+	if _, err := keyring.Get(ServiceName, prodKey); err == nil {
+		t.Error("expected no production credentials under new key")
+	}
+}
+
+func TestMigrateUserKeys_Idempotent(t *testing.T) {
+	keyring.MockInit()
+	storeCreds(t, "swells@datum.net", "auth.staging.env.datum.net")
+
+	cfg := collidingConfig()
+	if _, err := MigrateUserKeys(cfg); err != nil {
+		t.Fatalf("first migrate: %v", err)
+	}
+
+	changed, err := MigrateUserKeys(cfg)
+	if err != nil {
+		t.Fatalf("second migrate: %v", err)
+	}
+	if changed {
+		t.Error("second migration should be a no-op")
+	}
+}
+
+func TestMigrateUserKeys_AlreadyQualifiedNoop(t *testing.T) {
+	keyring.MockInit()
+
+	cfg := &datumconfig.ConfigV1Beta1{
+		Sessions: []datumconfig.Session{{
+			Name:      "swells@datum.net@api.datum.net",
+			UserKey:   "swells@datum.net@auth.datum.net",
+			UserEmail: "swells@datum.net",
+			Endpoint:  datumconfig.Endpoint{AuthHostname: "auth.datum.net"},
+		}},
+	}
+
+	changed, err := MigrateUserKeys(cfg)
+	if err != nil {
+		t.Fatalf("MigrateUserKeys: %v", err)
+	}
+	if changed {
+		t.Error("already-qualified session should not be migrated")
+	}
+}

internal/authutil/service_account_login.go

@@ -95,10 +95,13 @@ func RunServiceAccountLogin(ctx context.Context, credentialsPath, hostname, apiH
 		displayName = creds.ClientID
 	}
 
-	userKey := creds.ClientEmail
-	if userKey == "" {
-		userKey = creds.ClientID
+	identity := creds.ClientEmail
+	if identity == "" {
+		identity = creds.ClientID
 	}
+	// Scope the keyring entry by auth hostname so the same service account used
+	// against different environments does not share (and overwrite) one entry.
+	userKey := userKeyFor(identity, hostname)
 
 	keyFilePath, err := WriteServiceAccountKeyFile(userKey, creds.PrivateKey)
 	if err != nil {

internal/client/factory.go

@@ -243,6 +243,9 @@ func (c *CustomConfigFlags) loadDatumContext() (*datumconfig.DiscoveredContext,
 	if err != nil {
 		return nil, nil, err
 	}
+	if err := authutil.EnsureUserKeysMigrated(cfg); err != nil {
+		return nil, nil, err
+	}
 	ctxEntry := cfg.CurrentContextEntry()
 	if ctxEntry == nil {
 		return nil, nil, nil

internal/cmd/auth/switch.go

@@ -54,7 +54,7 @@ func runSwitch(_ *cobra.Command, args []string) error {
 		for i := range cfg.Sessions {
 			allSessions[i] = &cfg.Sessions[i]
 		}
-		sessionName, err = picker.SelectSession(allSessions)
+		sessionName, err = picker.SelectSession(allSessions, cfg.ActiveSession)
 		if err != nil {
 			return err
 		}
@@ -70,7 +70,7 @@ func runSwitch(_ *cobra.Command, args []string) error {
 		if len(sessions) == 1 {
 			sessionName = sessions[0].Name
 		} else {
-			sessionName, err = picker.SelectSession(sessions)
+			sessionName, err = picker.SelectSession(sessions, cfg.ActiveSession)
 			if err != nil {
 				return err
 			}

internal/cmd/whoami/whoami.go

@@ -25,6 +25,9 @@ func runWhoami(_ *cobra.Command, _ []string) error {
 	if err != nil {
 		return err
 	}
+	if err := authutil.EnsureUserKeysMigrated(cfg); err != nil {
+		return err
+	}
 
 	session := cfg.ActiveSessionEntry()
 	if session == nil {

internal/picker/picker.go

@@ -124,8 +124,11 @@ func buildContextOptions(contexts []datumconfig.DiscoveredContext, cfg *datumcon
 }
 
 // SelectSession presents an interactive picker for disambiguating between
-// sessions that share the same email. Returns the session name.
-func SelectSession(sessions []*datumconfig.Session) (string, error) {
+// sessions that share the same email. The session whose name matches
+// activeSessionName is labelled "(active)" and pre-selected so the picker opens
+// on the currently active session. Pass "" when no session is active. Returns
+// the selected session name.
+func SelectSession(sessions []*datumconfig.Session, activeSessionName string) (string, error) {
 	if len(sessions) == 0 {
 		return "", fmt.Errorf("no sessions to select from")
 	}
@@ -159,10 +162,14 @@ func SelectSession(sessions []*datumconfig.Session) (string, error) {
 		if showEndpoint {
 			label = fmt.Sprintf("%s  (%s)", s.UserEmail, datumconfig.StripScheme(s.Endpoint.Server))
 		}
+		if s.Name == activeSessionName {
+			label += "  (active)"
+		}
 		options[i] = huh.NewOption(label, s.Name)
 	}
 
-	var selected string
+	// Pre-select the active session so the picker opens on it.
+	selected := activeSessionName
 	form := huh.NewForm(
 		huh.NewGroup(
 			huh.NewSelect[string]().