[8.8] [Fleet] Fix usage of AsyncLocalStorage for audit log (#159983)

Backport https://github.com/elastic/kibana/pull/159807

Co-authored-by: Nicolas Chaulet <nicolas.chaulet@elastic.co>
Co-authored-by: kibanamachine <42973632+kibanamachine@users.noreply.github.com>

x-pack/plugins/fleet/server/plugin.ts

@@ -115,7 +115,6 @@ import type { PackagePolicyService } from './services/package_policy_service';
 import { PackagePolicyServiceImpl } from './services/package_policy';
 import { registerFleetUsageLogger, startFleetUsageLogger } from './services/fleet_usage_logger';
 import { CheckDeletedFilesTask } from './tasks/check_deleted_files_task';
-import { getRequestStore } from './services/request_store';
 
 export interface FleetSetupDeps {
   security: SecurityPluginSetup;
@@ -368,42 +367,38 @@ export class FleetPlugin
             .getSavedObjects()
             .getScopedClient(request, { excludedExtensions: [SECURITY_EXTENSION_ID] });
 
-        const requestStore = getRequestStore();
+        return {
+          get agentClient() {
+            const agentService = plugin.setupAgentService(esClient.asInternalUser, soClient);
 
-        return requestStore.run(request, () => {
-          return {
-            get agentClient() {
-              const agentService = plugin.setupAgentService(esClient.asInternalUser, soClient);
+            return {
+              asCurrentUser: agentService.asScoped(request),
+              asInternalUser: agentService.asInternalUser,
+            };
+          },
+          get packagePolicyService() {
+            const service = plugin.setupPackagePolicyService();
 
-              return {
-                asCurrentUser: agentService.asScoped(request),
-                asInternalUser: agentService.asInternalUser,
-              };
-            },
-            get packagePolicyService() {
-              const service = plugin.setupPackagePolicyService();
+            return {
+              asCurrentUser: service.asScoped(request),
+              asInternalUser: service.asInternalUser,
+            };
+          },
+          authz,
+          get internalSoClient() {
+            // Use a lazy getter to avoid constructing this client when not used by a request handler
+            return getInternalSoClient();
+          },
+          get spaceId() {
+            return deps.spaces?.spacesService?.getSpaceId(request) ?? DEFAULT_SPACE_ID;
+          },
 
-              return {
-                asCurrentUser: service.asScoped(request),
-                asInternalUser: service.asInternalUser,
-              };
-            },
-            authz,
-            get internalSoClient() {
-              // Use a lazy getter to avoid constructing this client when not used by a request handler
-              return getInternalSoClient();
-            },
-            get spaceId() {
-              return deps.spaces?.spacesService?.getSpaceId(request) ?? DEFAULT_SPACE_ID;
-            },
-
-            get limitedToPackages() {
-              if (routeAuthz && routeAuthz.granted) {
-                return routeAuthz.scopeDataToPackages;
-              }
-            },
-          };
-        });
+          get limitedToPackages() {
+            if (routeAuthz && routeAuthz.granted) {
+              return routeAuthz.scopeDataToPackages;
+            }
+          },
+        };
       }
     );
 

x-pack/plugins/fleet/server/services/request_store.ts

@@ -9,8 +9,8 @@ import { AsyncLocalStorage } from 'async_hooks';
 
 import type { KibanaRequest } from '@kbn/core-http-server';
 
-export function getRequestStore() {
-  const requestStore = new AsyncLocalStorage<KibanaRequest>();
+const requestStore = new AsyncLocalStorage<KibanaRequest>();
 
+export function getRequestStore() {
   return requestStore;
 }

x-pack/plugins/fleet/server/services/security/fleet_router.ts

@@ -17,6 +17,8 @@ import type {
 
 import type { FleetRequestHandlerContext } from '../..';
 
+import { getRequestStore } from '../request_store';
+
 import type { FleetAuthzRouteConfig, FleetAuthzRouter } from './types';
 import {
   checkSecurityEnabled,
@@ -60,30 +62,72 @@ export function makeRouterWithFleetAuthz<TContext extends FleetRequestHandlerCon
     return handler(context, request, response);
   };
 
+  const requestContextWrapper = async <R extends RouteMethod>({
+    context,
+    request,
+    response,
+    handler,
+  }: {
+    context: TContext;
+    request: KibanaRequest;
+    response: KibanaResponseFactory;
+    handler: RequestHandler<any, any, any, TContext, R, KibanaResponseFactory>;
+  }): Promise<IKibanaResponse<any>> => {
+    return getRequestStore().run(request, () => handler(context, request, response));
+  };
+
+  const fleetHandlerWrapper = async <R extends RouteMethod>({
+    context,
+    request,
+    response,
+    handler,
+    hasRequiredAuthz,
+  }: {
+    context: TContext;
+    request: KibanaRequest;
+    response: KibanaResponseFactory;
+    handler: RequestHandler<any, any, any, TContext, R, KibanaResponseFactory>;
+    hasRequiredAuthz?: FleetAuthzRouteConfig['fleetAuthz'];
+  }): Promise<IKibanaResponse<any>> => {
+    return requestContextWrapper({
+      context,
+      request,
+      response,
+      handler: (handlerContext, handlerRequest, handlerResponse) =>
+        routerAuthzWrapper({
+          context: handlerContext,
+          request: handlerRequest,
+          response: handlerResponse,
+          handler,
+          hasRequiredAuthz,
+        }),
+    });
+  };
+
   const fleetAuthzRouter: FleetAuthzRouter<TContext> = {
     get: ({ fleetAuthz: hasRequiredAuthz, ...options }, handler) => {
-      router.get(options, async (context, request, response) =>
-        routerAuthzWrapper({ context, request, response, handler, hasRequiredAuthz })
+      router.get(options, (context, request, response) =>
+        fleetHandlerWrapper({ context, request, response, handler, hasRequiredAuthz })
       );
     },
     delete: ({ fleetAuthz: hasRequiredAuthz, ...options }, handler) => {
-      router.delete(options, async (context, request, response) =>
-        routerAuthzWrapper({ context, request, response, handler, hasRequiredAuthz })
+      router.delete(options, (context, request, response) =>
+        fleetHandlerWrapper({ context, request, response, handler, hasRequiredAuthz })
       );
     },
     post: ({ fleetAuthz: hasRequiredAuthz, ...options }, handler) => {
-      router.post(options, async (context, request, response) =>
-        routerAuthzWrapper({ context, request, response, handler, hasRequiredAuthz })
+      router.post(options, (context, request, response) =>
+        fleetHandlerWrapper({ context, request, response, handler, hasRequiredAuthz })
       );
     },
     put: ({ fleetAuthz: hasRequiredAuthz, ...options }, handler) => {
-      router.put(options, async (context, request, response) =>
-        routerAuthzWrapper({ context, request, response, handler, hasRequiredAuthz })
+      router.put(options, (context, request, response) =>
+        fleetHandlerWrapper({ context, request, response, handler, hasRequiredAuthz })
       );
     },
     patch: ({ fleetAuthz: hasRequiredAuthz, ...options }, handler) => {
-      router.patch(options, async (context, request, response) =>
-        routerAuthzWrapper({ context, request, response, handler, hasRequiredAuthz })
+      router.patch(options, (context, request, response) =>
+        fleetHandlerWrapper({ context, request, response, handler, hasRequiredAuthz })
       );
     },
     handleLegacyErrors: (handler) => router.handleLegacyErrors(handler),