Merge #151887 #152021 #152121

151887: optbuilder: gate crdb_internal builtins from unauthorized access r=angles-n-daemons a=angles-n-daemons

The epic below outlines a set of work to prevent users from unauthorized
access to what we deem unsafe internals. These unsafe internals lie
mostly in the crdb_internal schema and the system database. This PR
makes it so crdb_internal builtins are disallowed from external use
unless the specified override is enabled.

Fixes: #151501
Epic: CRDB-24527
Release note (ops change): restricts access to all crdb_internal builtins
unless session variable `allow_unsafe_internals` is set to true, or if
the caller is internal.

152021: catalog/lease: fix assertion for no outstanding leases r=fqazi a=fqazi

Previously, the assertion confirming that no leases were outstanding was too strict because it checked for remaining descriptor versions immediately after setting the drain flag. While all user queries should be drained by this point, this check was too strict for internal operations. To address this, this patch moves the assertion into a closure and ensures that anything remaining has a reference count of zero.

Fixes: #151902

Release note: None

152121: authors: add jaywritescode to authors r=jaywritescode a=jaywritescode

Epic: None
Release note: None

Co-authored-by: Brian Dillmann <[email protected]>
Co-authored-by: Faizan Qazi <[email protected]>
Co-authored-by: Jay Harris <[email protected]>

Author: 4 people

AUTHORS

@@ -242,6 +242,7 @@ Jason Chan <[email protected]> <[email protected]>
 Jason Chen <[email protected]>
 Jason E. Aten <[email protected]>
 Jason Young <[email protected]>
+Jay Harris <[email protected]> <[email protected]>
 Jay Kominek <[email protected]>
 Jay Lim <[email protected]> <[email protected]> <[email protected]>
 Jay Rauchenstein <[email protected]>
@@ -576,4 +577,4 @@ Zachary Smith <[email protected]> Zachary.smith <[email protected]>
 Zakir Hussain Shaik <[email protected]> shaikzakiriitm <[email protected]>
 Zane Teh <[email protected]>
 何羿宏 <[email protected]>
-智雅楠 <[email protected]>
+智雅楠 <[email protected]>

pkg/server/api_v2_test.go

@@ -547,7 +547,7 @@ func drain(ctx context.Context, ts1 serverutils.TestServerInterface, t *testing.
 	timeoutCtx, cancel := context.WithTimeout(ctx, testutils.SucceedsSoonDuration())
 	defer cancel()
 
-	err := ts1.DrainClientsWithoutLeakCheck(ctx)
+	err := ts1.DrainClients(ctx)
 	require.NoError(t, err)
 
 	for timeoutCtx.Err() == nil {

pkg/server/drain.go

@@ -492,7 +492,7 @@ func (s *drainServer) drainClientsInternal(
 	// Drain all SQL table leases. This must be done after the pgServer has
 	// given sessions a chance to finish ongoing work and after the background
 	// tasks that may issue SQL statements have shut down.
-	s.sqlServer.leaseMgr.SetDraining(ctx, true /* drain */, reporter, assertOnLeakedDescriptor)
+	s.sqlServer.leaseMgr.SetDraining(ctx, true, reporter)
 
 	session, err := s.sqlServer.sqlLivenessProvider.Release(ctx)
 	if err != nil {

pkg/server/testserver.go

@@ -1242,11 +1242,6 @@ func (t *testTenant) DrainClients(ctx context.Context) error {
 	return t.drain.drainClients(ctx, nil /* reporter */)
 }
 
-// DrainClients exports the drainClientsInternal() method for use by tests.
-func (t *testTenant) DrainClientsWithoutLeakCheck(ctx context.Context) error {
-	return t.drain.drainClientsInternal(ctx, nil /* reporter */, false)
-}
-
 // Readiness is part of the serverutils.ApplicationLayerInterface.
 func (t *testTenant) Readiness(ctx context.Context) error {
 	return t.t.admin.checkReadinessForHealthCheck(ctx)
@@ -1948,13 +1943,6 @@ func (ts *testServer) DrainClients(ctx context.Context) error {
 	return ts.drain.drainClients(ctx, nil /* reporter */)
 }
 
-// DrainClientsWithoutLeakCheck exports the drainClients() method for use by tests
-// with descriptor leak checking disabled. This is useful for tests that focus on
-// restart safety logic rather than lease management correctness.
-func (ts *testServer) DrainClientsWithoutLeakCheck(ctx context.Context) error {
-	return ts.drain.drainClientsInternal(ctx, nil, false /* assertOnLeakedDescriptor */)
-}
-
 // Readiness is part of the serverutils.ApplicationLayerInterface.
 func (ts *testServer) Readiness(ctx context.Context) error {
 	return ts.admin.checkReadinessForHealthCheck(ctx)

pkg/sql/authorization_test.go

@@ -311,7 +311,7 @@ func checkUnsafeErr(t *testing.T, err error) {
 	t.Fatal("expected unsafe access error, got", err)
 }
 
-func TestCheckUnsafeSystemAccess(t *testing.T) {
+func TestCheckUnsafeInternalsAccess(t *testing.T) {
 	defer leaktest.AfterTest(t)()
 	defer log.Scope(t).Close(t)
 
@@ -321,92 +321,105 @@ func TestCheckUnsafeSystemAccess(t *testing.T) {
 	})
 	defer s.Stopper().Stop(ctx)
 
-	q := "SELECT * FROM system.namespace"
-
-	t.Run("as an external querier", func(t *testing.T) {
-		for _, test := range []struct {
-			AllowUnsafeInternals bool
-			Passes               bool
-		}{
-			{AllowUnsafeInternals: false, Passes: false},
-			{AllowUnsafeInternals: true, Passes: true},
-		} {
-			t.Run(fmt.Sprintf("%t", test), func(t *testing.T) {
-				conn := s.SQLConn(t)
-				_, err := conn.Exec("SET allow_unsafe_internals = $1", test.AllowUnsafeInternals)
-				require.NoError(t, err)
-
-				_, err = conn.Query(q)
-				if test.Passes {
+	t.Run("accessing the system database", func(t *testing.T) {
+		q := "SELECT * FROM system.namespace"
+
+		t.Run("as an external querier", func(t *testing.T) {
+			for _, test := range []struct {
+				AllowUnsafeInternals bool
+				Passes               bool
+			}{
+				{AllowUnsafeInternals: false, Passes: false},
+				{AllowUnsafeInternals: true, Passes: true},
+			} {
+				t.Run(fmt.Sprintf("%t", test), func(t *testing.T) {
+					conn := s.SQLConn(t)
+					_, err := conn.Exec("SET allow_unsafe_internals = $1", test.AllowUnsafeInternals)
 					require.NoError(t, err)
-				} else {
-					checkUnsafeErr(t, err)
-				}
-			})
-		}
-	})
 
-	t.Run("as an internal querier", func(t *testing.T) {
-		for _, test := range []struct {
-			AllowUnsafeInternals bool
-			Passes               bool
-		}{
-			{AllowUnsafeInternals: false, Passes: true},
-			{AllowUnsafeInternals: true, Passes: true},
-		} {
-			t.Run(fmt.Sprintf("%t", test), func(t *testing.T) {
-				idb := s.InternalDB().(isql.DB)
-				err := idb.Txn(ctx, func(ctx context.Context, txn isql.Txn) error {
-					txn.SessionData().LocalOnlySessionData.AllowUnsafeInternals = test.AllowUnsafeInternals
-
-					_, err := txn.QueryBuffered(ctx, "internal-query", txn.KV(), q)
-					return err
+					_, err = conn.Query(q)
+					if test.Passes {
+						require.NoError(t, err)
+					} else {
+						checkUnsafeErr(t, err)
+					}
 				})
+			}
+		})
 
-				require.NoError(t, err)
+		t.Run("as an internal querier", func(t *testing.T) {
+			for _, test := range []struct {
+				AllowUnsafeInternals bool
+				Passes               bool
+			}{
+				{AllowUnsafeInternals: false, Passes: true},
+				{AllowUnsafeInternals: true, Passes: true},
+			} {
+				t.Run(fmt.Sprintf("%t", test), func(t *testing.T) {
+					idb := s.InternalDB().(isql.DB)
+					err := idb.Txn(ctx, func(ctx context.Context, txn isql.Txn) error {
+						txn.SessionData().LocalOnlySessionData.AllowUnsafeInternals = test.AllowUnsafeInternals
+
+						_, err := txn.QueryBuffered(ctx, "internal-query", txn.KV(), q)
+						return err
+					})
 
-				if test.Passes {
 					require.NoError(t, err)
-				} else {
-					checkUnsafeErr(t, err)
-				}
-			})
-		}
-	})
-
-}
-
-func TestCheckUnsafeCRDBInternalAccess(t *testing.T) {
-	defer leaktest.AfterTest(t)()
-	defer log.Scope(t).Close(t)
 
-	ctx := context.Background()
-	s := serverutils.StartServerOnly(t, base.TestServerArgs{
-		DefaultTestTenant: base.TestControlsTenantsExplicitly,
+					if test.Passes {
+						require.NoError(t, err)
+					} else {
+						checkUnsafeErr(t, err)
+					}
+				})
+			}
+		})
 	})
-	defer s.Stopper().Stop(ctx)
 
-	// Test that with allow_unsafe_internals = false:
-	// - Supported tables (zones) are allowed
-	// - Unsupported tables (gossip_alerts) are denied
+	t.Run("accessing the crdb_internal schema", func(t *testing.T) {
+		t.Run("supported table allowed", func(t *testing.T) {
+			conn := s.SQLConn(t)
+			_, err := conn.Exec("SET allow_unsafe_internals = false")
+			require.NoError(t, err)
 
-	t.Run("supported table allowed", func(t *testing.T) {
-		conn := s.SQLConn(t)
-		_, err := conn.Exec("SET allow_unsafe_internals = false")
-		require.NoError(t, err)
+			// Supported crdb_internal tables should be allowed even when allow_unsafe_internals = false
+			_, err = conn.Query("SELECT * FROM crdb_internal.zones")
+			require.NoError(t, err, "supported crdb_internal table (zones) should be accessible when allow_unsafe_internals = false")
+		})
 
-		// Supported crdb_internal tables should be allowed even when allow_unsafe_internals = false
-		_, err = conn.Query("SELECT * FROM crdb_internal.zones")
-		require.NoError(t, err, "supported crdb_internal table (zones) should be accessible when allow_unsafe_internals = false")
+		t.Run("unsupported table denied", func(t *testing.T) {
+			conn := s.SQLConn(t)
+			_, err := conn.Exec("SET allow_unsafe_internals = false")
+			require.NoError(t, err)
+
+			// Unsupported crdb_internal tables should be denied when allow_unsafe_internals = false
+			_, err = conn.Query("SELECT * FROM crdb_internal.gossip_alerts")
+			checkUnsafeErr(t, err)
+		})
 	})
 
-	t.Run("unsupported table denied", func(t *testing.T) {
-		conn := s.SQLConn(t)
-		_, err := conn.Exec("SET allow_unsafe_internals = false")
-		require.NoError(t, err)
+	// The functionality for this lies in the pkg/sql/optbuilder/scalar.go
+	// file, but it is tested here as that package does not setup a test
+	// server.
+	t.Run("accessing crdb_internal builtins", func(t *testing.T) {
+		t.Run("non crdb_internal builtin allowed", func(t *testing.T) {
+			conn := s.SQLConn(t)
+			_, err := conn.Exec("SET allow_unsafe_internals = false")
+			require.NoError(t, err)
+
+			// Non crdb_internal tables should be allowed.
+			_, err = conn.Query("SELECT * FROM generate_series(1,5)")
+			require.NoError(t, err)
+		})
+
+		t.Run("crdb_internal builtin not allowed", func(t *testing.T) {
+			conn := s.SQLConn(t)
+			_, err := conn.Exec("SET allow_unsafe_internals = false")
+			require.NoError(t, err)
 
-		// Unsupported crdb_internal tables should be denied when allow_unsafe_internals = false
-		_, err = conn.Query("SELECT * FROM crdb_internal.gossip_alerts")
-		checkUnsafeErr(t, err)
+			// Unsupported crdb_internal builtins should be denied.
+			_, err = conn.Query("SELECT * FROM crdb_internal.tenant_span_stats()")
+			checkUnsafeErr(t, err)
+		})
 	})
 }

pkg/sql/catalog/lease/lease.go

@@ -1267,7 +1267,7 @@ type Manager struct {
 	// that we have all the updates for.
 	closeTimestamp atomic.Value
 
-	draining atomic.Value
+	draining atomic.Bool
 
 	// names is a cache for name -> id mappings. A mapping for the cache
 	// should only be used if we currently have an active lease on the respective
@@ -1417,6 +1417,7 @@ func NewLeaseManager(
 	lm.storage.regionPrefix.Store(enum.One)
 	lm.storage.writer = newKVWriter(codec, db.KV(), keys.LeaseTableID, settingsWatcher)
 	lm.stopper.AddCloser(lm.sem.Closer("stopper"))
+	lm.stopper.AddCloser(stop.CloserFn(lm.AssertAllLeasesAreReleasedAfterDrain))
 	lm.mu.descriptors = make(map[descpb.ID]*descriptorState)
 	lm.waitForInit = make(chan struct{})
 	// We are going to start the range feed later when StartRefreshLeasesTask
@@ -1731,7 +1732,7 @@ func (m *Manager) removeOnceDereferenced() bool {
 
 // IsDraining returns true if this node's lease manager is draining.
 func (m *Manager) IsDraining() bool {
-	return m.draining.Load().(bool)
+	return m.draining.Load()
 }
 
 // SetDraining (when called with 'true') removes all inactive leases. Any leases
@@ -1742,10 +1743,7 @@ func (m *Manager) IsDraining() bool {
 // been done by the time this call returns. See the explanation in
 // pkg/server/drain.go for details.
 func (m *Manager) SetDraining(
-	ctx context.Context,
-	drain bool,
-	reporter func(int, redact.SafeString),
-	assertOnLeakedDescriptor bool,
+	ctx context.Context, drain bool, reporter func(int, redact.SafeString),
 ) {
 	m.draining.Store(drain)
 	if !drain {
@@ -1759,15 +1757,6 @@ func (m *Manager) SetDraining(
 			t.mu.Lock()
 			defer t.mu.Unlock()
 			leasesToRelease := t.removeInactiveVersions(ctx)
-			// Ensure that all leases are released at this time. Ignore any descriptors
-			// that have acquisitions in progress.
-			if buildutil.CrdbTestBuild &&
-				assertOnLeakedDescriptor &&
-				len(t.mu.active.data) > 0 &&
-				t.mu.acquisitionsInProgress == 0 {
-				// Panic that a descriptor may have leaked.
-				panic(errors.AssertionFailedf("descriptor leak was detected for: %d (%s)", t.id, t.mu.active))
-			}
 			return leasesToRelease
 		}()
 		for _, l := range leases {
@@ -1780,6 +1769,45 @@ func (m *Manager) SetDraining(
 	}
 }
 
+// AssertAllLeasesAreReleasedAfterDrain asserts that all leases are released after
+// draining.
+func (m *Manager) AssertAllLeasesAreReleasedAfterDrain() {
+	if !buildutil.CrdbTestBuild || !m.draining.Load() {
+		return
+	}
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	for _, t := range m.mu.descriptors {
+		func() {
+			descriptorStr := strings.Builder{}
+			panicWithErr := false
+			t.mu.Lock()
+			defer t.mu.Unlock()
+			// Ensure that all leases are released at this time.
+			if len(t.mu.active.data) > 0 {
+				// Check if any of these have a non-zero ref count indicating some type of
+				// leak. It may be possible for the entry to exist if we were interrupted
+				// mid-acquisition by the draining process. But the reference count should
+				// *never* be non-zero.
+				for _, l := range t.mu.active.data {
+					if l.refcount.Load() == 0 {
+						continue
+					}
+					panicWithErr = true
+					if descriptorStr.Len() > 0 {
+						descriptorStr.WriteString(",")
+					}
+					descriptorStr.WriteString(fmt.Sprintf("{%s}", l.String()))
+				}
+				if panicWithErr {
+					panic(errors.AssertionFailedf("descriptor leak was detected for ID: %d, "+
+						"with versions [%s]", t.id, descriptorStr.String()))
+				}
+			}
+		}()
+	}
+}
+
 // isDescriptorStateEmpty determines if a descriptor state exists and
 // has any active versions inside it.
 func (m *Manager) isDescriptorStateEmpty(id descpb.ID) bool {

pkg/sql/catalog/lease/lease_test.go

@@ -522,8 +522,8 @@ func TestLeaseManagerDrain(testingT *testing.T) {
 		// starts draining.
 		l1RemovalTracker := leaseRemovalTracker.TrackRemoval(l1.Underlying())
 
-		t.node(1).SetDraining(ctx, true, nil /* reporter */, false)
-		t.node(2).SetDraining(ctx, true, nil /* reporter */, false)
+		t.node(1).SetDraining(ctx, true, nil /* reporter */)
+		t.node(2).SetDraining(ctx, true, nil)
 
 		// Leases cannot be acquired when in draining mode.
 		if _, err := t.acquire(1, descID); !testutils.IsError(err, "cannot acquire lease when draining") {
@@ -546,7 +546,7 @@ func TestLeaseManagerDrain(testingT *testing.T) {
 	{
 		// Check that leases with a refcount of 0 are correctly kept in the
 		// store once the drain mode has been exited.
-		t.node(1).SetDraining(ctx, false, nil /* reporter */, false /* assertOnLeakedDescriptor */)
+		t.node(1).SetDraining(ctx, false, nil)
 		l1 := t.mustAcquire(1, descID)
 		t.mustRelease(1, l1, nil)
 		t.expectLeases(descID, "/1/1")

pkg/sql/opt/optbuilder/BUILD.bazel

@@ -101,6 +101,7 @@ go_library(
         "//pkg/sql/sqltelemetry",
         "//pkg/sql/syntheticprivilege",
         "//pkg/sql/types",
+        "//pkg/sql/unsafesql",
         "//pkg/util",
         "//pkg/util/buildutil",
         "//pkg/util/errorutil",