[Fleet] Fetch only relevant assets for package policies operation

x-pack/platform/plugins/shared/fleet/common/types/models/epm.ts

@@ -126,6 +126,8 @@ export type InstallablePackage = RegistryPackage | ArchivePackage;
 
 export type AssetsMap = Map<string, Buffer | undefined>;
 
+export type PackagePolicyAssetsMap = AssetsMap & { __brand: 'PackagePolicyAssetsMap' };
+
 export interface ArchiveEntry {
   path: string;
   buffer?: Buffer;
@@ -723,6 +725,7 @@ export interface EsAssetReference {
 
 export interface PackageAssetReference {
   id: string;
+  path?: string; // Package installed prior to 9.1.0 will not have that property
   type: typeof ASSETS_SAVED_OBJECT_TYPE;
 }
 

x-pack/platform/plugins/shared/fleet/server/routes/package_policy/handlers.ts

@@ -60,7 +60,7 @@ import type {
   SimplifiedInputs,
   SimplifiedPackagePolicy,
 } from '../../../common/services/simplified_package_policy_helper';
-
+import { runWithCache } from '../../services/epm/packages/cache';
 import { validateAgentlessInputs } from '../../../common/services/agentless_policy_helper';
 
 import {
@@ -544,11 +544,12 @@ export const dryRunUpgradePackagePolicyHandler: RequestHandler<
 
   const body: UpgradePackagePolicyDryRunResponse = [];
   const { packagePolicyIds } = request.body;
-
-  for (const id of packagePolicyIds) {
-    const result = await packagePolicyService.getUpgradeDryRunDiff(soClient, id);
-    body.push(result);
-  }
+  await runWithCache(async () => {
+    for (const id of packagePolicyIds) {
+      const result = await packagePolicyService.getUpgradeDryRunDiff(soClient, id);
+      body.push(result);
+    }
+  });
 
   const firstFatalError = body.find((item) => item.statusCode && item.statusCode !== 200);
 

x-pack/platform/plugins/shared/fleet/server/services/epm/agent/agent.ts

@@ -79,7 +79,7 @@ function buildTemplateVariables(logger: Logger, variables: PackagePolicyConfigRe
     // support variables with . like key.patterns
     const keyParts = key.split('.');
     const lastKeyPart = keyParts.pop();
-    logger.debug(`Building agent template variables`);
+    // logger.debug(`Building agent template variables`);
 
     if (!lastKeyPart || !isValidKey(lastKeyPart)) {
       throw new PackageInvalidArchiveError(

x-pack/platform/plugins/shared/fleet/server/services/epm/archive/storage.ts

@@ -121,7 +121,11 @@ export async function saveArchiveEntriesFromAssetsMap(opts: {
     })
   );
 
-  const results = await savedObjectsClient.bulkCreate<PackageAsset>(bulkBody, { refresh: false });
+  const results = await savedObjectsClient.bulkCreate<PackageAsset>(bulkBody, {
+    refresh: false,
+    overwrite: true,
+  });
+
   return results;
 }
 

x-pack/platform/plugins/shared/fleet/server/services/epm/packages/cache.ts

@@ -17,12 +17,15 @@ const cacheStore = new AsyncLocalStorage<CacheSession>();
 
 const PACKAGE_INFO_CACHE_SIZE = 20;
 const PACKAGE_ASSETS_MAP_CACHE_SIZE = 1;
+const AGENT_TEMPLATE_ASSETS_MAP_CACHE_SIZE = 5;
 
 class CacheSession {
   private _packageInfoCache?: LRUCache<string, PackageInfo>;
 
   private _packageAssetsMap?: LRUCache<string, AssetsMap>;
 
+  private _agentTemplateAssetsMap?: LRUCache<string, AssetsMap>;
+
   getPackageInfoCache() {
     if (!this._packageInfoCache) {
       this._packageInfoCache = new LRUCache<string, PackageInfo>({
@@ -40,6 +43,15 @@ class CacheSession {
     }
     return this._packageAssetsMap;
   }
+
+  getAgentTemplateAssetsMapCache() {
+    if (!this._agentTemplateAssetsMap) {
+      this._agentTemplateAssetsMap = new LRUCache<string, AssetsMap>({
+        max: AGENT_TEMPLATE_ASSETS_MAP_CACHE_SIZE,
+      });
+    }
+    return this._agentTemplateAssetsMap;
+  }
 }
 
 export function getPackageInfoCache(pkgName: string, pkgVersion: string) {
@@ -65,6 +77,21 @@ export function setPackageAssetsMapCache(
     ?.set(`${pkgName}:${pkgVersion}`, assetsMap);
 }
 
+export function getAgentTemplateAssetsMapCache(pkgName: string, pkgVersion: string) {
+  return cacheStore.getStore()?.getAgentTemplateAssetsMapCache()?.get(`${pkgName}:${pkgVersion}`);
+}
+
+export function setAgentTemplateAssetsMapCache(
+  pkgName: string,
+  pkgVersion: string,
+  assetsMap: AssetsMap
+) {
+  return cacheStore
+    .getStore()
+    ?.getAgentTemplateAssetsMapCache()
+    ?.set(`${pkgName}:${pkgVersion}`, assetsMap);
+}
+
 export async function runWithCache<T = any>(cb: () => Promise<T>): Promise<T> {
   const cache = new CacheSession();
 

x-pack/platform/plugins/shared/fleet/server/services/epm/packages/get.ts

@@ -38,6 +38,7 @@ import type {
   InstalledPackage,
   PackageSpecManifest,
   AssetsMap,
+  PackagePolicyAssetsMap,
 } from '../../../../common/types';
 import {
   PACKAGES_SAVED_OBJECT_TYPE,
@@ -67,7 +68,7 @@ import { getPackagePolicySavedObjectType } from '../../package_policy';
 import { auditLoggingService } from '../../audit_logging';
 
 import { getFilteredSearchPackages } from '../filtered_packages';
-
+import { filterAssetPathForParseAndVerifyArchive } from '../archive/parse';
 import { airGappedUtils } from '../airgapped';
 
 import { createInstallableFrom } from '.';
@@ -76,6 +77,8 @@ import {
   setPackageAssetsMapCache,
   getPackageInfoCache,
   setPackageInfoCache,
+  getAgentTemplateAssetsMapCache,
+  setAgentTemplateAssetsMapCache,
 } from './cache';
 
 export { getFile } from '../registry';
@@ -715,15 +718,23 @@ export async function getInstalledPackageWithAssets(options: {
   pkgName: string;
   logger?: Logger;
   ignoreUnverified?: boolean;
+  assetsFilter?: (path: string) => boolean;
 }) {
   const installation = await getInstallation(options);
   if (!installation) {
     return;
   }
+  const assetsReference =
+    (typeof options.assetsFilter !== 'undefined'
+      ? installation.package_assets?.filter(({ path }) =>
+          typeof path !== 'undefined' ? options.assetsFilter!(path) : true
+        )
+      : installation.package_assets) ?? [];
+
   const esPackage = await getEsPackage(
     installation.name,
     installation.version,
-    installation.package_assets ?? [],
+    assetsReference,
     options.savedObjectsClient
   );
 
@@ -800,3 +811,56 @@ export async function getPackageAssetsMap({
     throw error;
   }
 }
+
+/**
+ * Return assets agent template assets map for package policies operation
+ */
+export async function getAgentTemplateAssetsMap({
+  savedObjectsClient,
+  packageInfo,
+  logger,
+  ignoreUnverified,
+}: {
+  savedObjectsClient: SavedObjectsClientContract;
+  packageInfo: PackageInfo;
+  logger: Logger;
+  ignoreUnverified?: boolean;
+}): Promise<PackagePolicyAssetsMap> {
+  const cache = getAgentTemplateAssetsMapCache(packageInfo.name, packageInfo.version);
+  if (cache) {
+    return cache;
+  }
+  const assetsFilter = (path: string) =>
+    filterAssetPathForParseAndVerifyArchive(path) || !!path.match(/\/agent\/.*\.hbs/);
+  const installedPackageWithAssets = await getInstalledPackageWithAssets({
+    savedObjectsClient,
+    pkgName: packageInfo.name,
+    logger,
+    assetsFilter,
+  });
+
+  try {
+    let assetsMap: AssetsMap | undefined;
+    if (installedPackageWithAssets?.installation.version !== packageInfo.version) {
+      // Try to get from registry
+      const pkg = await Registry.getPackage(packageInfo.name, packageInfo.version, {
+        ignoreUnverified,
+        useStreaming: true,
+      });
+      assetsMap = new Map();
+      await pkg.archiveIterator.traverseEntries(async (entry) => {
+        if (entry.buffer) {
+          assetsMap!.set(entry.path, entry.buffer);
+        }
+      }, assetsFilter);
+    } else {
+      assetsMap = installedPackageWithAssets.assetsMap;
+    }
+    setAgentTemplateAssetsMapCache(packageInfo.name, packageInfo.version, assetsMap);
+
+    return assetsMap;
+  } catch (error) {
+    logger.warn(`getAgentTemplateAssetsMap error: ${error}`);
+    throw error;
+  }
+}

x-pack/platform/plugins/shared/fleet/server/services/epm/packages/install_state_machine/steps/step_save_archive_entries.test.ts

@@ -210,6 +210,7 @@ describe('stepSaveArchiveEntries', () => {
       packageAssetRefs: [
         {
           id: 'test',
+          path: 'some/path',
           type: 'epm-packages-assets',
         },
       ],

x-pack/platform/plugins/shared/fleet/server/services/epm/packages/install_state_machine/steps/step_save_archive_entries.ts

@@ -40,6 +40,7 @@ export async function stepSaveArchiveEntries(context: InstallContext) {
       ...packageAssetRefs,
       ...packageAssetResults.saved_objects.map((result) => ({
         id: result.id,
+        path: result.attributes?.asset_path,
         type: ASSETS_SAVED_OBJECT_TYPE as typeof ASSETS_SAVED_OBJECT_TYPE,
       })),
     ];

x-pack/platform/plugins/shared/fleet/server/services/package_policy.test.ts

@@ -188,6 +188,7 @@ jest.mock('./epm/registry', () => ({
 
 jest.mock('./epm/packages/get', () => ({
   getPackageAssetsMap: jest.fn().mockResolvedValue(new Map()),
+  getAgentTemplateAssetsMap: jest.fn().mockResolvedValue(new Map()),
 }));
 
 jest.mock('./agent_policy');

x-pack/platform/plugins/shared/fleet/server/services/package_policy.ts

@@ -70,8 +70,8 @@ import type {
   ExperimentalDataStreamFeature,
   DeletePackagePoliciesResponse,
   PolicySecretReference,
-  AssetsMap,
   AgentPolicy,
+  PackagePolicyAssetsMap,
 } from '../../common/types';
 import {
   FleetError,
@@ -148,7 +148,7 @@ import {
   deleteSecretsIfNotReferenced as deleteSecrets,
   isSecretStorageEnabled,
 } from './secrets';
-import { getPackageAssetsMap } from './epm/packages/get';
+import { getAgentTemplateAssetsMap } from './epm/packages/get';
 import { validateAgentPolicyOutputForIntegration } from './agent_policies/outputs_helpers';
 import type { PackagePolicyClientFetchAllItemIdsOptions } from './package_policy_service';
 import {
@@ -174,12 +174,12 @@ async function getPkgInfoAssetsMap({
 }) {
   const packageInfosandAssetsMap = new Map<
     string,
-    { assetsMap: AssetsMap; pkgInfo: PackageInfo }
+    { assetsMap: PackagePolicyAssetsMap; pkgInfo: PackageInfo }
   >();
   await pMap(
     packageInfos,
     async (pkgInfo) => {
-      const assetsMap = await getPackageAssetsMap({
+      const assetsMap = await getAgentTemplateAssetsMap({
         logger,
         packageInfo: pkgInfo,
         savedObjectsClient,
@@ -385,7 +385,7 @@ class PackagePolicyClientImpl implements PackagePolicyClient {
 
         inputs = enrichedPackagePolicy.inputs as PackagePolicyInput[];
       }
-      const assetsMap = await getPackageAssetsMap({
+      const assetsMap = await getAgentTemplateAssetsMap({
         logger,
         packageInfo: pkgInfo,
         savedObjectsClient: soClient,
@@ -717,7 +717,7 @@ class PackagePolicyClientImpl implements PackagePolicyClient {
           `Package info and assets not found: ${packagePolicy.package.name}-${packagePolicy.package.version}`
         );
       }
-      const assetsMap = await getPackageAssetsMap({
+      const assetsMap = await getAgentTemplateAssetsMap({
         logger: appContextService.getLogger(),
         packageInfo: pkgInfo,
         savedObjectsClient: soClient,
@@ -1036,7 +1036,7 @@ class PackagePolicyClientImpl implements PackagePolicyClient {
         secretsToDelete = secretsRes.secretsToDelete;
         inputs = restOfPackagePolicy.inputs as PackagePolicyInput[];
       }
-      const assetsMap = await getPackageAssetsMap({
+      const assetsMap = await getAgentTemplateAssetsMap({
         logger,
         packageInfo: pkgInfo,
         savedObjectsClient: soClient,
@@ -1770,7 +1770,7 @@ class PackagePolicyClientImpl implements PackagePolicyClient {
       packageInfo,
       packageToPackagePolicyInputs(packageInfo) as InputsOverride[]
     );
-    const assetsMap = await getPackageAssetsMap({
+    const assetsMap = await getAgentTemplateAssetsMap({
       logger: appContextService.getLogger(),
       packageInfo,
       savedObjectsClient: soClient,
@@ -1809,11 +1809,19 @@ class PackagePolicyClientImpl implements PackagePolicyClient {
 
       ({ packagePolicy, packageInfo, experimentalDataStreamFeatures } =
         await this.getUpgradePackagePolicyInfo(soClient, id, packagePolicy, pkgVersion));
-      const assetsMap = await getPackageAssetsMap({
+
+      // const getAgentTemplateAssets =
+      // getAgentTemplateAssetsMap;
+      const assetsMap = await getAgentTemplateAssetsMap({
         logger: appContextService.getLogger(),
         packageInfo,
         savedObjectsClient: soClient,
       });
+      // const assetsMap = await getPackageAssetsMap({
+      //   logger: appContextService.getLogger(),
+      //   packageInfo,
+      //   savedObjectsClient: soClient,
+      // });
 
       // Ensure the experimental features from the Installation saved object come through on the package policy
       // during an upgrade dry run
@@ -1834,7 +1842,7 @@ class PackagePolicyClientImpl implements PackagePolicyClient {
     soClient: SavedObjectsClientContract,
     packagePolicy: PackagePolicy,
     packageInfo: PackageInfo,
-    assetsMap: AssetsMap
+    assetsMap: PackagePolicyAssetsMap
   ): Promise<UpgradePackagePolicyDryRunResponseItem> {
     const updatedPackagePolicy = updatePackageInputs(
       {
@@ -2507,7 +2515,7 @@ export async function _compilePackagePolicyInputs(
   pkgInfo: PackageInfo,
   vars: PackagePolicy['vars'],
   inputs: PackagePolicyInput[],
-  assetsMap: AssetsMap
+  assetsMap: PackagePolicyAssetsMap
 ): Promise<PackagePolicyInput[]> {
   const inputsPromises = inputs.map(async (input) => {
     const compiledInput = await _compilePackagePolicyInput(pkgInfo, vars, input, assetsMap);
@@ -2526,7 +2534,7 @@ async function _compilePackagePolicyInput(
   pkgInfo: PackageInfo,
   vars: PackagePolicy['vars'],
   input: PackagePolicyInput,
-  assetsMap: AssetsMap
+  assetsMap: PackagePolicyAssetsMap
 ) {
   const packagePolicyTemplate = input.policy_template
     ? pkgInfo.policy_templates?.find(
@@ -2575,7 +2583,7 @@ async function _compilePackageStreams(
   pkgInfo: PackageInfo,
   vars: PackagePolicy['vars'],
   input: PackagePolicyInput,
-  assetsMap: AssetsMap
+  assetsMap: PackagePolicyAssetsMap
 ) {
   const streamsPromises = input.streams.map((stream) =>
     _compilePackageStream(pkgInfo, vars, input, stream, assetsMap)
@@ -2646,7 +2654,7 @@ async function _compilePackageStream(
   vars: PackagePolicy['vars'],
   input: PackagePolicyInput,
   streamIn: PackagePolicyInputStream,
-  assetsMap: AssetsMap
+  assetsMap: PackagePolicyAssetsMap
 ) {
   let stream = streamIn;