[Fleet] Fix update query when change agent policy spaces (#198175)

elastic · Oct 30, 2024 · fdd5e0b · fdd5e0b
1 parent 74cf0e4
commit fdd5e0b
Show file tree

Hide file tree

Showing 3 changed files with 84 additions and 5 deletions.
diff --git a/x-pack/plugins/fleet/server/services/spaces/agent_policy.ts b/x-pack/plugins/fleet/server/services/spaces/agent_policy.ts
@@ -115,12 +115,30 @@ export async function updateAgentPolicySpaces({
   // Update fleet server index agents, enrollment api keys
   await esClient.updateByQuery({
     index: ENROLLMENT_API_KEYS_INDEX,
+    query: {
+      bool: {
+        must: {
+          terms: {
+            policy_id: [agentPolicyId],
+          },
+        },
+      },
+    },
     script: `ctx._source.namespaces = [${newSpaceIds.map((spaceId) => `"${spaceId}"`).join(',')}]`,
     ignore_unavailable: true,
     refresh: true,
   });
   await esClient.updateByQuery({
     index: AGENTS_INDEX,
+    query: {
+      bool: {
+        must: {
+          terms: {
+            policy_id: [agentPolicyId],
+          },
+        },
+      },
+    },
     script: `ctx._source.namespaces = [${newSpaceIds.map((spaceId) => `"${spaceId}"`).join(',')}]`,
     ignore_unavailable: true,
     refresh: true,

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
@@ -15,6 +15,7 @@ import {
   createFleetAgent,
   expectToRejectWithError,
   expectToRejectWithNotFound,
+  getFleetAgentDoc,
 } from './helpers';
 import { testUsers, setupTestUsers } from '../test_users';
 
@@ -32,8 +33,12 @@ export default function (providerContext: FtrProviderContext) {
     const apiClient = new SpaceTestApiClient(supertest);
 
     let defaultSpacePolicy1: CreateAgentPolicyResponse;
+    let defaultSpacePolicy2: CreateAgentPolicyResponse;
     let defaultPackagePolicy1: GetOnePackagePolicyResponse;
 
+    let policy1AgentId: string;
+    let policy2AgentId: string;
+
     before(async () => {
       TEST_SPACE_1 = spaces.getDefaultTestSpace();
       await setupTestUsers(getService('security'), true);
@@ -44,14 +49,20 @@ export default function (providerContext: FtrProviderContext) {
       await cleanFleetIndices(esClient);
 
       await apiClient.postEnableSpaceAwareness();
-      const _policyRes = await apiClient.createAgentPolicy();
-      defaultSpacePolicy1 = _policyRes;
+      const [_policyRes1, _policyRes2] = await Promise.all([
+        apiClient.createAgentPolicy(),
+        apiClient.createAgentPolicy(),
+      ]);
+      defaultSpacePolicy1 = _policyRes1;
+      defaultSpacePolicy2 = _policyRes2;
       await apiClient.installPackage({
         pkgName: 'nginx',
         pkgVersion: '1.20.0',
         force: true, // To avoid package verification
       });
-      await createFleetAgent(esClient, defaultSpacePolicy1.item.id);
+      policy1AgentId = await createFleetAgent(esClient, defaultSpacePolicy1.item.id);
+      policy2AgentId = await createFleetAgent(esClient, defaultSpacePolicy2.item.id);
+
       const packagePolicyRes = await apiClient.createPackagePolicy(undefined, {
         policy_ids: [defaultSpacePolicy1.item.id],
         name: `test-nginx-${Date.now()}`,
@@ -107,7 +118,22 @@ export default function (providerContext: FtrProviderContext) {
         ).not.to.be(undefined);
 
         const agents = await apiClient.getAgents(spaceId);
-        expect(agents.total).to.be(1);
+        expect(
+          agents.items.filter((a) => a.policy_id === defaultSpacePolicy1.item.id).length
+        ).to.be(1);
+      }
+
+      async function assertEnrollemntApiKeysForSpace(spaceId?: string, policyIds?: string[]) {
+        const spaceApiKeys = await apiClient.getEnrollmentApiKeys(spaceId);
+
+        const foundPolicyIds = spaceApiKeys.items.reduce((acc, apiKey) => {
+          if (apiKey.policy_id) {
+            acc.add(apiKey.policy_id);
+          }
+          return acc;
+        }, new Set<string>());
+
+        expect([...foundPolicyIds].sort()).to.eql(policyIds?.sort());
       }
 
       async function assertPolicyNotAvailableInSpace(spaceId?: string) {
@@ -124,7 +150,19 @@ export default function (providerContext: FtrProviderContext) {
         ).to.be(undefined);
 
         const agents = await apiClient.getAgents(spaceId);
-        expect(agents.total).to.be(0);
+        expect(
+          agents.items.filter((a) => a.policy_id === defaultSpacePolicy1.item.id).length
+        ).to.be(0);
+      }
+
+      async function assertAgentSpaces(agentId: string, expectedSpaces: string[]) {
+        const agentDoc = await getFleetAgentDoc(esClient, agentId);
+
+        if (expectedSpaces.length === 1 && expectedSpaces[0] === 'default') {
+          expect(agentDoc._source?.namespaces ?? ['default']).to.eql(expectedSpaces);
+        } else {
+          expect(agentDoc._source?.namespaces).to.eql(expectedSpaces);
+        }
       }
 
       it('should allow set policy in multiple space', async () => {
@@ -137,6 +175,15 @@ export default function (providerContext: FtrProviderContext) {
 
         await assertPolicyAvailableInSpace();
         await assertPolicyAvailableInSpace(TEST_SPACE_1);
+
+        await assertAgentSpaces(policy1AgentId, ['default', TEST_SPACE_1]);
+        await assertAgentSpaces(policy2AgentId, ['default']);
+
+        await assertEnrollemntApiKeysForSpace('default', [
+          defaultSpacePolicy1.item.id,
+          defaultSpacePolicy2.item.id,
+        ]);
+        await assertEnrollemntApiKeysForSpace(TEST_SPACE_1, [defaultSpacePolicy1.item.id]);
       });
 
       it('should allow set policy in test space only', async () => {
@@ -149,6 +196,10 @@ export default function (providerContext: FtrProviderContext) {
 
         await assertPolicyNotAvailableInSpace();
         await assertPolicyAvailableInSpace(TEST_SPACE_1);
+        await assertAgentSpaces(policy1AgentId, [TEST_SPACE_1]);
+        await assertAgentSpaces(policy2AgentId, ['default']);
+        await assertEnrollemntApiKeysForSpace('default', [defaultSpacePolicy2.item.id]);
+        await assertEnrollemntApiKeysForSpace(TEST_SPACE_1, [defaultSpacePolicy1.item.id]);
       });
 
       it('should not allow add policy to a space where user do not have access', async () => {

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
@@ -13,6 +13,7 @@ import {
   AGENT_ACTIONS_RESULTS_INDEX,
   AGENT_POLICY_INDEX,
   AGENTS_INDEX,
+  type FleetServerAgent,
 } from '@kbn/fleet-plugin/common';
 import { ENROLLMENT_API_KEYS_INDEX } from '@kbn/fleet-plugin/common/constants';
 import { asyncForEach } from '@kbn/std';
@@ -117,6 +118,15 @@ export async function createFleetAgent(esClient: Client, agentPolicyId: string,
   return agentResponse._id;
 }
 
+export async function getFleetAgentDoc(esClient: Client, agentId: string) {
+  const agentResponse = await esClient.get<FleetServerAgent>({
+    index: '.fleet-agents',
+    id: agentId,
+  });
+
+  return agentResponse;
+}
+
 export async function makeAgentsUpgradeable(esClient: Client, agentIds: string[], version: string) {
   await asyncForEach(agentIds, async (agentId) => {
     await esClient.update({