Merge pull request #99055 from dhartunian/release-22.1-stmtdiag-fix

release-22.1: admin: statement diagnostics uses correct auth helpers

pkg/server/admin.go

@@ -197,7 +197,15 @@ func (s *adminServer) RegisterGateway(
 			http.Error(w, "invalid id", http.StatusBadRequest)
 			return
 		}
-		s.getStatementBundle(ctx, id, w)
+		// Add default user when running in Insecure mode because we don't
+		// retrieve the user from gRPC metadata (which falls back to `root`)
+		// but from HTTP metadata (which does not).
+		if s.server.cfg.Insecure {
+			ctx := req.Context()
+			ctx = context.WithValue(ctx, webSessionUserKey{}, security.RootUser)
+			req = req.WithContext(ctx)
+		}
+		s.getStatementBundle(req.Context(), id, w)
 	})
 
 	// Register the endpoints defined in the proto.
@@ -2323,14 +2331,13 @@ func (s *adminServer) QueryPlan(
 // getStatementBundle retrieves the statement bundle with the given id and
 // writes it out as an attachment.
 func (s *adminServer) getStatementBundle(ctx context.Context, id int64, w http.ResponseWriter) {
-	sessionUser, err := userFromContext(ctx)
-	if err != nil {
-		http.Error(w, err.Error(), http.StatusInternalServerError)
-		return
-	}
+	// This endpoint uses `getSQLUsername` to get its user metadata
+	// due to the fact that it is implemented as a raw HTTP handler
+	// instead of a gRPC handler.
+	sqlUsername := getSQLUsername(ctx)
 	row, err := s.server.sqlServer.internalExecutor.QueryRowEx(
 		ctx, "admin-stmt-bundle", nil, /* txn */
-		sessiondata.InternalExecutorOverride{User: sessionUser},
+		sessiondata.InternalExecutorOverride{User: sqlUsername},
 		"SELECT bundle_chunks FROM system.statement_diagnostics WHERE id=$1 AND bundle_chunks IS NOT NULL",
 		id,
 	)
@@ -2349,7 +2356,7 @@ func (s *adminServer) getStatementBundle(ctx context.Context, id int64, w http.R
 	for _, chunkID := range chunkIDs {
 		chunkRow, err := s.server.sqlServer.internalExecutor.QueryRowEx(
 			ctx, "admin-stmt-bundle", nil, /* txn */
-			sessiondata.InternalExecutorOverride{User: sessionUser},
+			sessiondata.InternalExecutorOverride{User: sqlUsername},
 			"SELECT data FROM system.statement_bundle_chunks WHERE id=$1",
 			chunkID,
 		)

pkg/server/admin_test.go

@@ -312,6 +312,43 @@ func TestAdminDebugRedirect(t *testing.T) {
 	}
 }
 
+func TestAdminAPIStatementDiagnosticsBundle(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+	s, db, _ := serverutils.StartServer(t, base.TestServerArgs{})
+	defer s.Stopper().Stop(context.Background())
+	ts := s.(*TestServer)
+
+	query := "EXPLAIN ANALYZE (DEBUG) SELECT 'secret'"
+	_, err := db.Exec(query)
+	require.NoError(t, err)
+
+	query = "SELECT id FROM system.statement_diagnostics LIMIT 1"
+	idRow, err := db.Query(query)
+	require.NoError(t, err)
+	var diagnosticRow string
+	if idRow.Next() {
+		err = idRow.Scan(&diagnosticRow)
+		require.NoError(t, err)
+	} else {
+		t.Fatal("no results")
+	}
+
+	client, err := ts.GetAuthenticatedHTTPClient(false)
+	require.NoError(t, err)
+	resp, err := client.Get(ts.AdminURL() + "/_admin/v1/stmtbundle/" + diagnosticRow)
+	require.NoError(t, err)
+	defer resp.Body.Close()
+	require.Equal(t, 500, resp.StatusCode)
+
+	adminClient, err := ts.GetAuthenticatedHTTPClient(true)
+	require.NoError(t, err)
+	adminResp, err := adminClient.Get(ts.AdminURL() + "/_admin/v1/stmtbundle/" + diagnosticRow)
+	require.NoError(t, err)
+	defer adminResp.Body.Close()
+	require.Equal(t, 200, adminResp.StatusCode)
+}
+
 func TestAdminAPIDatabases(t *testing.T) {
 	defer leaktest.AfterTest(t)()
 	defer log.Scope(t).Close(t)