From e1ac4869cb68c925a8b1dc6ef8e9f884e4e8bf2a Mon Sep 17 00:00:00 2001
From: David Hartunian <davidh@cockroachlabs.com>
Date: Sun, 19 Mar 2023 11:46:36 -0400
Subject: [PATCH] server: only set default tenant if login successful

Previously, we would always set the default tenant cookie to the
default tenant cluster setting regardless of what tenants the user
logged-in to successfully.

This change ensures that the default tenant selection is only used
when the successful logins include that tenant. Otherwise, we select
the first tenant from the list of successful logins.

Epic: CRDB-12100
Release note: None
---
 pkg/ccl/serverccl/server_controller_test.go | 43 +++++++++++++++++++++
 pkg/server/server_controller_http.go        | 14 ++++++-
 2 files changed, 56 insertions(+), 1 deletion(-)

diff --git a/pkg/ccl/serverccl/server_controller_test.go b/pkg/ccl/serverccl/server_controller_test.go
index 75ad77d59a9a..fac401ea8f69 100644
--- a/pkg/ccl/serverccl/server_controller_test.go
+++ b/pkg/ccl/serverccl/server_controller_test.go
@@ -9,6 +9,7 @@
 package serverccl
 
 import (
+	"bytes"
 	"context"
 	"fmt"
 	"io"
@@ -19,6 +20,7 @@ import (
 	"time"
 
 	"github.com/cockroachdb/cockroach/pkg/base"
+	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/server"
 	"github.com/cockroachdb/cockroach/pkg/server/serverpb"
 	"github.com/cockroachdb/cockroach/pkg/sql/lexbase"
@@ -183,6 +185,47 @@ VALUES($1, $2, $3, $4, $5)`, id, secret, username, created, expires)
 	require.Equal(t, body.Sessions[0].ApplicationName, "hello system")
 }
 
+// TestServerControllerDefaultHTTPTenant ensures that the default
+// tenant selected in the cookie does not use the default tenant
+// from the cluster setting *unless* the user successfully logged
+// in to that tenant.
+func TestServerControllerDefaultHTTPTenant(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	ctx := context.Background()
+
+	s, _, _ := serverutils.StartServer(t, base.TestServerArgs{DisableDefaultTestTenant: true})
+	defer s.Stopper().Stop(ctx)
+
+	_, sql, err := s.StartSharedProcessTenant(ctx, base.TestSharedProcessTenantArgs{
+		TenantName: "hello",
+		TenantID:   roachpb.MustMakeTenantID(10),
+	})
+	require.NoError(t, err)
+
+	_, err = sql.Exec("CREATE user foo with password 'cockroach'")
+	require.NoError(t, err)
+
+	client, err := s.GetUnauthenticatedHTTPClient()
+	require.NoError(t, err)
+
+	resp, err := client.Post(s.AdminURL()+"/login",
+		"application/json",
+		bytes.NewBuffer([]byte("{\"username\":\"foo\",\"password\":\"cockroach\"})")),
+	)
+	require.NoError(t, err)
+	defer resp.Body.Close()
+
+	tenantCookie := ""
+	for _, c := range resp.Cookies() {
+		if c.Name == server.TenantSelectCookieName {
+			tenantCookie = c.Value
+		}
+	}
+	require.Equal(t, "hello", tenantCookie)
+}
+
 // TestServerControllerBadHTTPCookies tests the controller's proxy
 // layer for correct behavior under scenarios where the client has
 // stale or invalid cookies. This helps ensure that we continue to
diff --git a/pkg/server/server_controller_http.go b/pkg/server/server_controller_http.go
index 730702a267da..89c0a860e530 100644
--- a/pkg/server/server_controller_http.go
+++ b/pkg/server/server_controller_http.go
@@ -201,9 +201,21 @@ func (c *serverController) attemptLoginToAllTenants() http.Handler {
 			http.SetCookie(w, &cookie)
 			// The tenant cookie needs to be set at some point in order for
 			// the dropdown to have a current selection on first load.
+
+			// We only set the default selection from the cluster setting
+			// if it's one of the valid logins. Otherwise, we just use the
+			// first one in the list.
+			tenantSelection := tenantNameToSetCookieSlice[0].name
+			defaultName := defaultTenantSelect.Get(&c.st.SV)
+			for _, t := range tenantNameToSetCookieSlice {
+				if t.name == defaultName {
+					tenantSelection = t.name
+					break
+				}
+			}
 			cookie = http.Cookie{
 				Name:     TenantSelectCookieName,
-				Value:    defaultTenantSelect.Get(&c.st.SV),
+				Value:    tenantSelection,
 				Path:     "/",
 				HttpOnly: false,
 			}