[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

@@ -9,20 +9,23 @@ import { AsyncLocalStorage } from 'async_hooks';
 
 import LRUCache from 'lru-cache';
 
-import type { AssetsMap } from '../../../../common/types';
+import type { AssetsMap, PackagePolicyAssetsMap } from '../../../../common/types';
 
 import type { PackageInfo } from '../../../../common';
 
 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, PackagePolicyAssetsMap>;
+
   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, PackagePolicyAssetsMap>({
+        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: PackagePolicyAssetsMap
+) {
+  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: PackagePolicyAssetsMap | 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() as PackagePolicyAssetsMap;
+      await pkg.archiveIterator.traverseEntries(async (entry) => {
+        if (entry.buffer) {
+          assetsMap!.set(entry.path, entry.buffer);
+        }
+      }, assetsFilter);
+    } else {
+      assetsMap = installedPackageWithAssets.assetsMap as PackagePolicyAssetsMap;
+    }
+    setAgentTemplateAssetsMapCache(packageInfo.name, packageInfo.version, assetsMap);
+
+    return assetsMap as PackagePolicyAssetsMap;
+  } catch (error) {
+    logger.warn(`getAgentTemplateAssetsMap error: ${error}`);
+    throw error;
+  }
+}

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

@@ -33,7 +33,7 @@ import { _sortYamlKeys } from '../../../../common/services/full_agent_policy_to_
 import { getFullInputStreams } from '../../agent_policies/package_policies_to_agent_inputs';
 
 import { getPackageInfo } from '.';
-import { getPackageAssetsMap } from './get';
+import { getAgentTemplateAssetsMap } from './get';
 
 type Format = 'yml' | 'json';
 
@@ -161,7 +161,7 @@ export async function getTemplateInputs(
     }
   }
 
-  const assetsMap = await getPackageAssetsMap({
+  const assetsMap = await getAgentTemplateAssetsMap({
     logger: appContextService.getLogger(),
     packageInfo,
     savedObjectsClient: soClient,

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

@@ -8,12 +8,12 @@
 import { savedObjectsClientMock } from '@kbn/core-saved-objects-api-server-mocks';
 
 import { createAppContextStartContractMock } from '../../../mocks';
-import type { PackagePolicyInput } from '../../../../common/types';
+import type { PackagePolicyAssetsMap, PackagePolicyInput } from '../../../../common/types';
 import { appContextService } from '../..';
 
 import { getTemplateInputs, templatePackagePolicyToFullInputStreams } from './get_template_inputs';
 import REDIS_1_18_0_PACKAGE_INFO from './__fixtures__/redis_1_18_0_package_info.json';
-import { getPackageAssetsMap, getPackageInfo } from './get';
+import { getAgentTemplateAssetsMap, getPackageInfo } from './get';
 import { REDIS_ASSETS_MAP } from './__fixtures__/redis_1_18_0_streams_template';
 import { LOGS_2_3_0_ASSETS_MAP, LOGS_2_3_0_PACKAGE_INFO } from './__fixtures__/logs_2_3_0';
 import { DOCKER_2_11_0_PACKAGE_INFO, DOCKER_2_11_0_ASSETS_MAP } from './__fixtures__/docker_2_11_0';
@@ -324,19 +324,19 @@ describe('Fleet - templatePackagePolicyToFullInputStreams', () => {
 describe('Fleet - getTemplateInputs', () => {
   beforeEach(() => {
     appContextService.start(createAppContextStartContractMock());
-    jest.mocked(getPackageAssetsMap).mockImplementation(async ({ packageInfo }) => {
+    jest.mocked(getAgentTemplateAssetsMap).mockImplementation(async ({ packageInfo }) => {
       if (packageInfo.name === 'redis' && packageInfo.version === '1.18.0') {
-        return REDIS_ASSETS_MAP;
+        return REDIS_ASSETS_MAP as PackagePolicyAssetsMap;
       }
 
       if (packageInfo.name === 'log') {
-        return LOGS_2_3_0_ASSETS_MAP;
+        return LOGS_2_3_0_ASSETS_MAP as PackagePolicyAssetsMap;
       }
       if (packageInfo.name === 'docker') {
-        return DOCKER_2_11_0_ASSETS_MAP;
+        return DOCKER_2_11_0_ASSETS_MAP as PackagePolicyAssetsMap;
       }
 
-      return new Map();
+      return new Map() as PackagePolicyAssetsMap;
     });
     jest.mocked(getPackageInfo).mockImplementation(async ({ pkgName, pkgVersion }) => {
       const pkgInfo = packageInfoCache.get(`${pkgName}-${pkgVersion}`);

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

@@ -50,6 +50,7 @@ import type {
   NewPackagePolicyInput,
   PackagePolicyPackage,
   DeletePackagePoliciesResponse,
+  PackagePolicyAssetsMap,
 } from '../../common/types';
 import { packageToPackagePolicy } from '../../common/services';
 
@@ -120,7 +121,7 @@ const ASSETS_MAP_FIXTURES = new Map([
   {{/each}}
   `),
   ],
-]);
+]) as PackagePolicyAssetsMap;
 
 async function mockedGetInstallation(params: any) {
   let pkg;
@@ -188,6 +189,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');