From 8d5cd4f2d752a8d47c37225c91561af3763464d3 Mon Sep 17 00:00:00 2001 From: Nicolas Chaulet Date: Thu, 12 Dec 2024 08:23:50 -0500 Subject: [PATCH] [Fleet] Fix changing space for agent actions (#203683) --- .../services/spaces/agent_policy.test.ts | 6 ++ .../server/services/spaces/agent_policy.ts | 60 ++++++++++++++++++- .../change_space_agent_policies.ts | 52 ++++++++++++++++ .../apis/space_awareness/helpers.ts | 6 ++ 4 files changed, 121 insertions(+), 3 deletions(-) diff --git a/x-pack/plugins/fleet/server/services/spaces/agent_policy.test.ts b/x-pack/plugins/fleet/server/services/spaces/agent_policy.test.ts index 079148a6ae041..6b2c63aae9126 100644 --- a/x-pack/plugins/fleet/server/services/spaces/agent_policy.test.ts +++ b/x-pack/plugins/fleet/server/services/spaces/agent_policy.test.ts @@ -7,6 +7,7 @@ import { createAppContextStartContractMock } from '../../mocks'; import { agentPolicyService } from '../agent_policy'; +import { getAgentsByKuery } from '../agents'; import { appContextService } from '../app_context'; import { packagePolicyService } from '../package_policy'; @@ -16,6 +17,7 @@ import { isSpaceAwarenessEnabled } from './helpers'; jest.mock('./helpers'); jest.mock('../agent_policy'); jest.mock('../package_policy'); +jest.mock('../agents'); describe('updateAgentPolicySpaces', () => { beforeEach(() => { @@ -55,6 +57,10 @@ describe('updateAgentPolicySpaces', () => { } as any, ], }); + + jest.mocked(getAgentsByKuery).mockResolvedValue({ + agents: [], + } as any); }); it('does nothings if agent policy already in correct space', async () => { diff --git a/x-pack/plugins/fleet/server/services/spaces/agent_policy.ts b/x-pack/plugins/fleet/server/services/spaces/agent_policy.ts index 50f20443c3262..dce40f4cc3ff4 100644 --- a/x-pack/plugins/fleet/server/services/spaces/agent_policy.ts +++ b/x-pack/plugins/fleet/server/services/spaces/agent_policy.ts @@ -6,27 +6,29 @@ */ import deepEqual from 'fast-deep-equal'; - import { DEFAULT_SPACE_ID } from '@kbn/spaces-plugin/common'; +import type { SortResults } from '@elastic/elasticsearch/lib/api/types'; import { AGENTS_INDEX, + AGENT_ACTIONS_INDEX, AGENT_POLICY_SAVED_OBJECT_TYPE, PACKAGE_POLICY_SAVED_OBJECT_TYPE, SO_SEARCH_LIMIT, UNINSTALL_TOKENS_SAVED_OBJECT_TYPE, } from '../../../common/constants'; - import { appContextService } from '../app_context'; import { agentPolicyService } from '../agent_policy'; import { ENROLLMENT_API_KEYS_INDEX } from '../../constants'; import { packagePolicyService } from '../package_policy'; import { FleetError, HostedAgentPolicyRestrictionRelatedError } from '../../errors'; - import type { UninstallTokenSOAttributes } from '../security/uninstall_token_service'; +import { closePointInTime, getAgentsByKuery, openPointInTime } from '../agents'; import { isSpaceAwarenessEnabled } from './helpers'; +const UPDATE_AGENT_BATCH_SIZE = 1000; + export async function updateAgentPolicySpaces({ agentPolicyId, currentSpaceId, @@ -49,6 +51,7 @@ export async function updateAgentPolicySpaces({ const soClient = appContextService.getInternalUserSOClientWithoutSpaceExtension(); const currentSpaceSoClient = appContextService.getInternalUserSOClientForSpaceId(currentSpaceId); + const newSpaceSoClient = appContextService.getInternalUserSOClientForSpaceId(newSpaceIds[0]); const existingPolicy = await agentPolicyService.get(currentSpaceSoClient, agentPolicyId); const existingPackagePolicies = await packagePolicyService.findAllForAgentPolicy( @@ -166,4 +169,55 @@ export async function updateAgentPolicySpaces({ ignore_unavailable: true, refresh: true, }); + + const agentIndexExists = await esClient.indices.exists({ + index: AGENTS_INDEX, + }); + + // Update agent actions + if (agentIndexExists) { + const pitId = await openPointInTime(esClient); + + try { + let hasMore = true; + let searchAfter: SortResults | undefined; + while (hasMore) { + const { agents } = await getAgentsByKuery(esClient, newSpaceSoClient, { + kuery: `policy_id:"${agentPolicyId}"`, + showInactive: true, + perPage: UPDATE_AGENT_BATCH_SIZE, + pitId, + searchAfter, + }); + + if (agents.length === 0) { + hasMore = false; + break; + } + + const lastAgent = agents[agents.length - 1]; + searchAfter = lastAgent.sort; + + await esClient.updateByQuery({ + index: AGENT_ACTIONS_INDEX, + query: { + bool: { + must: { + terms: { + agents: agents.map(({ id }) => id), + }, + }, + }, + }, + script: `ctx._source.namespaces = [${newSpaceIds + .map((spaceId) => `"${spaceId}"`) + .join(',')}]`, + ignore_unavailable: true, + refresh: true, + }); + } + } finally { + await closePointInTime(esClient, pitId); + } + } } diff --git a/x-pack/test/fleet_api_integration/apis/space_awareness/change_space_agent_policies.ts b/x-pack/test/fleet_api_integration/apis/space_awareness/change_space_agent_policies.ts index df3aa6f8f257b..519ae1109e511 100644 --- a/x-pack/test/fleet_api_integration/apis/space_awareness/change_space_agent_policies.ts +++ b/x-pack/test/fleet_api_integration/apis/space_awareness/change_space_agent_policies.ts @@ -6,7 +6,10 @@ */ import expect from '@kbn/expect'; +import { v4 as uuidV4 } from 'uuid'; +import { Client } from '@elastic/elasticsearch'; import { CreateAgentPolicyResponse, GetOnePackagePolicyResponse } from '@kbn/fleet-plugin/common'; +import { FleetServerAgentAction } from '@kbn/fleet-plugin/common/types'; import { FtrProviderContext } from '../../../api_integration/ftr_provider_context'; import { skipIfNoDockerRegistry } from '../../helpers'; import { SpaceTestApiClient } from './api_helper'; @@ -19,6 +22,33 @@ import { } from './helpers'; import { testUsers, setupTestUsers } from '../test_users'; +export async function createFleetAction(esClient: Client, agentId: string, spaceId?: string) { + const actionResponse = await esClient.index({ + index: '.fleet-actions', + refresh: 'wait_for', + body: { + '@timestamp': new Date().toISOString(), + expiration: new Date().toISOString(), + agents: [agentId], + action_id: uuidV4(), + data: {}, + type: 'UPGRADE', + namespaces: spaceId ? [spaceId] : undefined, + }, + }); + + return actionResponse._id; +} + +async function getFleetActionDoc(esClient: Client, actionId: string) { + const actionResponse = await esClient.get({ + index: '.fleet-actions', + id: actionId, + }); + + return actionResponse; +} + export default function (providerContext: FtrProviderContext) { const { getService } = providerContext; const supertest = getService('supertest'); @@ -39,6 +69,9 @@ export default function (providerContext: FtrProviderContext) { let policy1AgentId: string; let policy2AgentId: string; + let agent1ActionId: string; + let agent2ActionId: string; + before(async () => { TEST_SPACE_1 = spaces.getDefaultTestSpace(); await setupTestUsers(getService('security'), true); @@ -63,6 +96,9 @@ export default function (providerContext: FtrProviderContext) { policy1AgentId = await createFleetAgent(esClient, defaultSpacePolicy1.item.id); policy2AgentId = await createFleetAgent(esClient, defaultSpacePolicy2.item.id); + agent1ActionId = await createFleetAction(esClient, policy1AgentId, 'default'); + agent2ActionId = await createFleetAction(esClient, policy2AgentId, 'default'); + const packagePolicyRes = await apiClient.createPackagePolicy(undefined, { policy_ids: [defaultSpacePolicy1.item.id], name: `test-nginx-${Date.now()}`, @@ -172,6 +208,15 @@ export default function (providerContext: FtrProviderContext) { } } + async function assertActionSpaces(actionId: string, expectedSpaces: string[]) { + const actionDoc = await getFleetActionDoc(esClient, actionId); + if (expectedSpaces.length === 1 && expectedSpaces[0] === 'default') { + expect(actionDoc._source?.namespaces ?? ['default']).to.eql(expectedSpaces); + } else { + expect(actionDoc._source?.namespaces).to.eql(expectedSpaces); + } + } + it('should allow set policy in multiple space', async () => { await apiClient.putAgentPolicy(defaultSpacePolicy1.item.id, { name: 'tata', @@ -189,6 +234,9 @@ export default function (providerContext: FtrProviderContext) { await assertAgentSpaces(policy1AgentId, ['default', TEST_SPACE_1]); await assertAgentSpaces(policy2AgentId, ['default']); + await assertActionSpaces(agent1ActionId, ['default', TEST_SPACE_1]); + await assertActionSpaces(agent2ActionId, ['default']); + await assertEnrollemntApiKeysForSpace('default', [ defaultSpacePolicy1.item.id, defaultSpacePolicy2.item.id, @@ -213,6 +261,10 @@ export default function (providerContext: FtrProviderContext) { await assertPackagePolicyNotAvailableInSpace(); await assertAgentSpaces(policy1AgentId, [TEST_SPACE_1]); await assertAgentSpaces(policy2AgentId, ['default']); + + await assertActionSpaces(agent1ActionId, [TEST_SPACE_1]); + await assertActionSpaces(agent2ActionId, ['default']); + await assertEnrollemntApiKeysForSpace('default', [defaultSpacePolicy2.item.id]); await assertEnrollemntApiKeysForSpace(TEST_SPACE_1, [defaultSpacePolicy1.item.id]); // Ensure no side effect on other policies diff --git a/x-pack/test/fleet_api_integration/apis/space_awareness/helpers.ts b/x-pack/test/fleet_api_integration/apis/space_awareness/helpers.ts index 92f4e3a387678..e4e5ff2563cc1 100644 --- a/x-pack/test/fleet_api_integration/apis/space_awareness/helpers.ts +++ b/x-pack/test/fleet_api_integration/apis/space_awareness/helpers.ts @@ -49,6 +49,12 @@ export async function cleanFleetIndices(esClient: Client) { ignore_unavailable: true, refresh: true, }), + esClient.deleteByQuery({ + index: AGENT_ACTIONS_INDEX, + q: '*', + ignore_unavailable: true, + refresh: true, + }), ]); }