From 748a6e288a002078ff5b6ec225c92c276c91fa1e Mon Sep 17 00:00:00 2001 From: Nicolas Chaulet Date: Wed, 5 Mar 2025 11:31:59 -0500 Subject: [PATCH] [Fleet] Fetch only relevant assets for package policies operation --- .../shared/fleet/common/types/models/epm.ts | 3 + .../server/routes/package_policy/handlers.ts | 13 ++-- .../fleet/server/services/epm/agent/agent.ts | 2 +- .../server/services/epm/archive/storage.ts | 6 +- .../server/services/epm/packages/cache.ts | 29 +++++++- .../fleet/server/services/epm/packages/get.ts | 68 ++++++++++++++++++- .../epm/packages/get_template_inputs.ts | 4 +- .../epm/packages/get_templates_inputs.test.ts | 14 ++-- .../steps/step_save_archive_entries.test.ts | 1 + .../steps/step_save_archive_entries.ts | 1 + .../server/services/package_policy.test.ts | 4 +- .../fleet/server/services/package_policy.ts | 36 ++++++---- 12 files changed, 146 insertions(+), 35 deletions(-) diff --git a/x-pack/platform/plugins/shared/fleet/common/types/models/epm.ts b/x-pack/platform/plugins/shared/fleet/common/types/models/epm.ts index 84c19295adcaf..ae5d3b46e6f3c 100644 --- a/x-pack/platform/plugins/shared/fleet/common/types/models/epm.ts +++ b/x-pack/platform/plugins/shared/fleet/common/types/models/epm.ts @@ -126,6 +126,8 @@ export type InstallablePackage = RegistryPackage | ArchivePackage; export type AssetsMap = Map; +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; } diff --git a/x-pack/platform/plugins/shared/fleet/server/routes/package_policy/handlers.ts b/x-pack/platform/plugins/shared/fleet/server/routes/package_policy/handlers.ts index 0e537c90d06c0..e488888db4ae1 100644 --- a/x-pack/platform/plugins/shared/fleet/server/routes/package_policy/handlers.ts +++ b/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); diff --git a/x-pack/platform/plugins/shared/fleet/server/services/epm/agent/agent.ts b/x-pack/platform/plugins/shared/fleet/server/services/epm/agent/agent.ts index 20bb9ed557572..57e5cf3c32819 100644 --- a/x-pack/platform/plugins/shared/fleet/server/services/epm/agent/agent.ts +++ b/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( diff --git a/x-pack/platform/plugins/shared/fleet/server/services/epm/archive/storage.ts b/x-pack/platform/plugins/shared/fleet/server/services/epm/archive/storage.ts index 8f6f151383d5a..9de6c87286398 100644 --- a/x-pack/platform/plugins/shared/fleet/server/services/epm/archive/storage.ts +++ b/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(bulkBody, { refresh: false }); + const results = await savedObjectsClient.bulkCreate(bulkBody, { + refresh: false, + overwrite: true, + }); + return results; } diff --git a/x-pack/platform/plugins/shared/fleet/server/services/epm/packages/cache.ts b/x-pack/platform/plugins/shared/fleet/server/services/epm/packages/cache.ts index 63eb825365fbe..2ca5fb6249b5a 100644 --- a/x-pack/platform/plugins/shared/fleet/server/services/epm/packages/cache.ts +++ b/x-pack/platform/plugins/shared/fleet/server/services/epm/packages/cache.ts @@ -9,7 +9,7 @@ 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'; @@ -17,12 +17,15 @@ const cacheStore = new AsyncLocalStorage(); 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; private _packageAssetsMap?: LRUCache; + private _agentTemplateAssetsMap?: LRUCache; + getPackageInfoCache() { if (!this._packageInfoCache) { this._packageInfoCache = new LRUCache({ @@ -40,6 +43,15 @@ class CacheSession { } return this._packageAssetsMap; } + + getAgentTemplateAssetsMapCache() { + if (!this._agentTemplateAssetsMap) { + this._agentTemplateAssetsMap = new LRUCache({ + 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(cb: () => Promise): Promise { const cache = new CacheSession(); diff --git a/x-pack/platform/plugins/shared/fleet/server/services/epm/packages/get.ts b/x-pack/platform/plugins/shared/fleet/server/services/epm/packages/get.ts index 2753f0514f6ab..03accf80f8d95 100644 --- a/x-pack/platform/plugins/shared/fleet/server/services/epm/packages/get.ts +++ b/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 { + 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; + } +} diff --git a/x-pack/platform/plugins/shared/fleet/server/services/epm/packages/get_template_inputs.ts b/x-pack/platform/plugins/shared/fleet/server/services/epm/packages/get_template_inputs.ts index f76b739904f3f..7da752aa987f6 100644 --- a/x-pack/platform/plugins/shared/fleet/server/services/epm/packages/get_template_inputs.ts +++ b/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, diff --git a/x-pack/platform/plugins/shared/fleet/server/services/epm/packages/get_templates_inputs.test.ts b/x-pack/platform/plugins/shared/fleet/server/services/epm/packages/get_templates_inputs.test.ts index 1a3738d8eaa82..836c536019558 100644 --- a/x-pack/platform/plugins/shared/fleet/server/services/epm/packages/get_templates_inputs.test.ts +++ b/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}`); diff --git a/x-pack/platform/plugins/shared/fleet/server/services/epm/packages/install_state_machine/steps/step_save_archive_entries.test.ts b/x-pack/platform/plugins/shared/fleet/server/services/epm/packages/install_state_machine/steps/step_save_archive_entries.test.ts index 1a4f2998d49dd..26a4f17132357 100644 --- a/x-pack/platform/plugins/shared/fleet/server/services/epm/packages/install_state_machine/steps/step_save_archive_entries.test.ts +++ b/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', }, ], diff --git a/x-pack/platform/plugins/shared/fleet/server/services/epm/packages/install_state_machine/steps/step_save_archive_entries.ts b/x-pack/platform/plugins/shared/fleet/server/services/epm/packages/install_state_machine/steps/step_save_archive_entries.ts index 54221723cb28b..059920afd0aaa 100644 --- a/x-pack/platform/plugins/shared/fleet/server/services/epm/packages/install_state_machine/steps/step_save_archive_entries.ts +++ b/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, })), ]; diff --git a/x-pack/platform/plugins/shared/fleet/server/services/package_policy.test.ts b/x-pack/platform/plugins/shared/fleet/server/services/package_policy.test.ts index 215559640995e..c024561e2ad40 100644 --- a/x-pack/platform/plugins/shared/fleet/server/services/package_policy.test.ts +++ b/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'); diff --git a/x-pack/platform/plugins/shared/fleet/server/services/package_policy.ts b/x-pack/platform/plugins/shared/fleet/server/services/package_policy.ts index ab4b8f83ceb56..e1d3d856c929d 100644 --- a/x-pack/platform/plugins/shared/fleet/server/services/package_policy.ts +++ b/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 { const updatedPackagePolicy = updatePackageInputs( { @@ -2507,7 +2515,7 @@ export async function _compilePackagePolicyInputs( pkgInfo: PackageInfo, vars: PackagePolicy['vars'], inputs: PackagePolicyInput[], - assetsMap: AssetsMap + assetsMap: PackagePolicyAssetsMap ): Promise { 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;