unsafesql: avoid panicking during query formatting

Part of the effort to guard access to the crdb_internal and system
namespaces includes auditing override access (and denied access) to
these unsafe internals. Included in this audit is the offending query
which attempted to pry into these namespaces.

In multiple locations however, this auditing caused the system to panic,
for different reasons. In one case, an incorrect number of annotations
on the query caused a panic. Another included a plan builder which had
no associated statement.

We see the process of going from plan -> query as a difficult one, and
thus guard this attempt to audit these accesses in a blanket panic
catcher, as it's not common that this will happen, and when it does we
don't want the system to wholesale fail the query.

Fixes: #153590 Epic: CRDB-24527

Release note: none

Author: angles-n-daemons

pkg/sql/unsafesql/BUILD.bazel

@@ -14,6 +14,7 @@ go_library(
         "//pkg/util/log/eventpb",
         "//pkg/util/log/severity",
         "//pkg/util/timeutil",
+        "@com_github_cockroachdb_redact//:redact",
     ],
 )
 
@@ -26,7 +27,9 @@ go_test(
         "//pkg/security/securityassets",
         "//pkg/security/securitytest",
         "//pkg/server",
+        "//pkg/settings",
         "//pkg/sql/isql",
+        "//pkg/sql/sem/tree",
         "//pkg/sql/sqlerrors",
         "//pkg/testutils/serverutils",
         "//pkg/testutils/sqlutils",

pkg/sql/unsafesql/unsafesql.go

@@ -17,6 +17,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/log/eventpb"
 	"github.com/cockroachdb/cockroach/pkg/util/log/severity"
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
+	"github.com/cockroachdb/redact"
 )
 
 // The accessedLogLimiter is used to limit the rate of logging unsafe internal access
@@ -53,7 +54,7 @@ func CheckInternalsAccess(
 		return nil
 	}
 
-	q := tree.FormatAstAsRedactableString(stmt, ann, sv)
+	q := SafeFormatQuery(stmt, ann, sv)
 	// If an override is set, allow access to this virtual table.
 	if sd.AllowUnsafeInternals {
 		// Log this access to the SENSITIVE_ACCESS channel since the override condition bypassed normal access controls.
@@ -69,3 +70,20 @@ func CheckInternalsAccess(
 	}
 	return sqlerrors.ErrUnsafeTableAccess
 }
+
+// SafeFormatQuery attempts to format the query for logging, but recovers from
+// any panics that may occur during formatting.
+func SafeFormatQuery(
+	stmt tree.Statement, ann *tree.Annotations, sv *settings.Values,
+) (s redact.RedactableString) {
+	if stmt == nil {
+		return "<nil statement>"
+	}
+	defer func() {
+		if r := recover(); r != nil {
+			log.Dev.Errorf(context.TODO(), "panic in SafeFormatQuery: %v", r)
+			s = "<panicked query format>"
+		}
+	}()
+	return tree.FormatAstAsRedactableString(stmt, ann, sv)
+}

pkg/sql/unsafesql/unsafesql_test.go

@@ -16,7 +16,9 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/security/securityassets"
 	"github.com/cockroachdb/cockroach/pkg/security/securitytest"
 	"github.com/cockroachdb/cockroach/pkg/server"
+	"github.com/cockroachdb/cockroach/pkg/settings"
 	"github.com/cockroachdb/cockroach/pkg/sql/isql"
+	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/sql/sqlerrors"
 	"github.com/cockroachdb/cockroach/pkg/sql/unsafesql"
 	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
@@ -255,3 +257,35 @@ func TestAccessCheckServer(t *testing.T) {
 		})
 	}
 }
+
+// panickingStatement is a mock Statement that panics during formatting
+type panickingStatement struct{}
+
+func (ps panickingStatement) String() string {
+	return "panicking statement"
+}
+
+func (ps panickingStatement) Format(ctx *tree.FmtCtx) {
+	panic("deliberate panic for testing")
+}
+
+func (ps panickingStatement) StatementReturnType() tree.StatementReturnType {
+	return tree.Unknown
+}
+
+func (ps panickingStatement) StatementType() tree.StatementType {
+	return tree.TypeDML
+}
+
+func (ps panickingStatement) StatementTag() string {
+	return "TEST"
+}
+
+func TestPanickingSQLFormat(t *testing.T) {
+	result := unsafesql.SafeFormatQuery(panickingStatement{}, nil, &settings.Values{})
+	require.Equal(
+		t,
+		"<panicked query format>",
+		string(result),
+	)
+}