Merge cockroachdb#75575 cockroachdb#75583 cockroachdb#75592

75575: sql: check session_user can become user on deserialize_session r=rafiss a=otan

Release note (bug fix): `crdb_internal.deserialize_session` now checks
the session_user has the privilege to SET ROLE to the current_user before
changing the session settings.

75583: sql: de-flake TestRecordBatchSerializerDeserializeMemoryEstimate r=cucaroach a=cucaroach

Fixes: cockroachdb#73781

When BatchSize was low and each buffer was much smaller than capacity
we could end up with newMemorySize more than %33 smaller than original.
However the test is just making sure we aren't double counting memory
so testing that newMemorySize - originalMemory < 4/3 original.  I.e.
its okay if it shrinks but we don't want it to be more than %33 bigger.

Release note: None

75592: logictest: de-flake TestLogic/5node-disk/distsql_enum r=irfansharif a=irfansharif

Fixes cockroachdb#75570. Second attempt at 83f584d, though now with the right
syntax.

Release note: None

Co-authored-by: Oliver Tan <[email protected]>
Co-authored-by: Tommy Reilly <[email protected]>
Co-authored-by: irfan sharif <[email protected]>

pkg/col/colserde/record_batch_test.go

@@ -341,12 +341,13 @@ func TestRecordBatchSerializerDeserializeMemoryEstimate(t *testing.T) {
 	newMemorySize := colmem.GetBatchMemSize(dest)
 
 	// We expect that the original and the new memory sizes are relatively close
-	// to each other (do not differ by more than a third). We cannot guarantee
-	// more precise bound here because the capacities of the underlying []byte
-	// slices is unpredictable. However, this check is sufficient to ensure that
-	// we don't double count memory under `Bytes.data`.
+	// to each other, specifically newMemorySize must less than
+	// 4/3*originalMemorySize. We cannot guarantee more precise bound here because
+	// the capacities of the underlying []byte slices is unpredictable. However,
+	// this check is sufficient to ensure that we don't double count memory under
+	// `Bytes.data`.
 	const maxDeviation = float64(0.33)
-	deviation := math.Abs(float64(originalMemorySize-newMemorySize) / (float64(originalMemorySize)))
+	deviation := float64(newMemorySize-originalMemorySize) / float64(originalMemorySize)
 	require.GreaterOrEqualf(t, maxDeviation, deviation,
 		"new memory size %d is too far away from original %d", newMemorySize, originalMemorySize)
 }

pkg/sql/faketreeeval/evalctx.go

@@ -216,6 +216,13 @@ func (ep *DummyEvalPlanner) UserHasAdminRole(
 	return false, errors.WithStack(errEvalPlanner)
 }
 
+// CheckCanBecomeUser is part of the EvalPlanner interface.
+func (ep *DummyEvalPlanner) CheckCanBecomeUser(
+	ctx context.Context, becomeUser security.SQLUsername,
+) error {
+	return errors.WithStack(errEvalPlanner)
+}
+
 // MemberOfWithAdminOption is part of the EvalPlanner interface.
 func (ep *DummyEvalPlanner) MemberOfWithAdminOption(
 	ctx context.Context, member security.SQLUsername,

pkg/sql/logictest/testdata/logic_test/distsql_enum

@@ -65,7 +65,7 @@ ALTER TABLE t2 INJECT STATISTICS '[
   }
 ]'
 
-query T nodeidx=1 retry
+query T nodeidx=1,retry
 EXPLAIN (VEC)
 SELECT x from t1 WHERE EXISTS (SELECT x FROM t2 WHERE t1.x=t2.x AND t2.y='hello')
 ----

pkg/sql/sem/builtins/builtins.go

@@ -6384,6 +6384,9 @@ table's zone configuration this will return NULL.`,
 						"can only deserialize matching session users",
 					)
 				}
+				if err := evalCtx.Planner.CheckCanBecomeUser(evalCtx.Context, sd.User()); err != nil {
+					return nil, err
+				}
 				*evalCtx.SessionData() = *sd
 				return tree.MakeDBool(true), nil
 			},

pkg/sql/sem/tree/eval.go

@@ -3232,6 +3232,10 @@ type EvalPlanner interface {
 	// the `system.users` table
 	UserHasAdminRole(ctx context.Context, user security.SQLUsername) (bool, error)
 
+	// CheckCanBecomeUser returns an error if the SessionUser cannot become the
+	// becomeUser.
+	CheckCanBecomeUser(ctx context.Context, becomeUser security.SQLUsername) error
+
 	// MemberOfWithAdminOption is used to collect a list of roles (direct and
 	// indirect) that the member is part of. See the comment on the planner
 	// implementation in authorization.go

pkg/sql/session_migration_test.go

@@ -14,6 +14,7 @@ import (
 	"context"
 	gosql "database/sql"
 	"fmt"
+	"net/url"
 	"strings"
 	"testing"
 
@@ -24,6 +25,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/testutils/testcluster"
 	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
+	"github.com/cockroachdb/cockroach/pkg/util/stop"
 	"github.com/cockroachdb/datadriven"
 	"github.com/stretchr/testify/require"
 )
@@ -57,6 +59,29 @@ func TestSessionMigration(t *testing.T) {
 		defer func() {
 			_ = dbConn.Close()
 		}()
+		_, err := dbConn.Exec("CREATE USER testuser")
+		require.NoError(t, err)
+
+		openUserConnFunc := func(user string) *gosql.DB {
+
+			pgURL, cleanupGoDB, err := sqlutils.PGUrlE(
+				tc.Server(0).ServingSQLAddr(),
+				"StartServer", /* prefix */
+				url.User(user),
+			)
+			require.NoError(t, err)
+			pgURL.Path = "defaultdb"
+
+			goDB, err := gosql.Open("postgres", pgURL.String())
+			require.NoError(t, err)
+
+			tc.Server(0).Stopper().AddCloser(
+				stop.CloserFn(func() {
+					cleanupGoDB()
+				}))
+
+			return goDB
+		}
 
 		vars := make(map[string]string)
 		datadriven.RunTest(t, path, func(t *testing.T, d *datadriven.TestData) string {
@@ -72,6 +97,10 @@ func TestSessionMigration(t *testing.T) {
 				require.NoError(t, dbConn.Close())
 				dbConn = openConnFunc()
 				return ""
+			case "user":
+				require.NoError(t, dbConn.Close())
+				dbConn = openUserConnFunc(d.Input)
+				return ""
 			case "exec":
 				_, err := dbConn.Exec(getQuery())
 				if err != nil {

pkg/sql/sessiondata/session_data.go

@@ -148,7 +148,7 @@ func (s *SessionData) GetDateStyle() pgdate.DateStyle {
 // SessionUser retrieves the session_user.
 // The SessionUser is the username that originally logged into the session.
 // If a user applies SET ROLE, the SessionUser remains the same whilst the
-// CurrentUser() changes.
+// User() changes.
 func (s *SessionData) SessionUser() security.SQLUsername {
 	if s.SessionUserProto == "" {
 		return s.User()

pkg/sql/testdata/session_migration/errors

@@ -67,3 +67,21 @@ SELECT crdb_internal.deserialize_session(
 )
 ----
 pq: crdb_internal.deserialize_session(): can only deserialize matching session users
+
+# We cannot deserialize into a current_user we do not match.
+user
+testuser
+----
+
+exec
+SELECT crdb_internal.deserialize_session(
+  decode(
+    -- minted by modifying `crdb_internal.serialize_session()` and setting
+    --   sd.SessionData.UserProto=root
+    --   sd.LocalOnlySessionData.SessionUser=testuser
+    '0a510a0964656661756c74646212102420636f636b726f6163682064656d6f1a04726f6f742204100222002802380842035554434a0524757365724a067075626c69635a0060808080207a0088010190018050124910904e3002380840026001680170017801880101d80101e00101f00101f80101900201b002808001c80201f202087465737475736572a00301a9030000000000408f40d00301e00301',
+    'hex'
+  )
+)
+----
+pq: crdb_internal.deserialize_session(): only root can become root

pkg/sql/user.go

@@ -528,7 +528,7 @@ func (p *planner) setRole(ctx context.Context, local bool, s security.SQLUsernam
 		}
 	}
 
-	if err := p.checkCanBecomeUser(ctx, becomeUser); err != nil {
+	if err := p.CheckCanBecomeUser(ctx, becomeUser); err != nil {
 		return err
 	}
 
@@ -573,7 +573,8 @@ func (p *planner) setRole(ctx context.Context, local bool, s security.SQLUsernam
 
 }
 
-func (p *planner) checkCanBecomeUser(ctx context.Context, becomeUser security.SQLUsername) error {
+// CheckCanBecomeUser implements the EvalPlanner interface.
+func (p *planner) CheckCanBecomeUser(ctx context.Context, becomeUser security.SQLUsername) error {
 	sessionUser := p.SessionData().SessionUser()
 
 	// Switching to None can always succeed.