[Fleet] Modify Agent policy and agent upgrades to handle custom source_uri

x-pack/plugins/fleet/common/constants/download_source.ts

@@ -5,7 +5,8 @@
  * 2.0.
  */
 
-export const DEFAULT_DOWNLOAD_SOURCE = 'artifactory.elastic.co';
+// Default URL used to download Elastic Agent
+export const DEFAULT_DOWNLOAD_SOURCE = 'https://artifacts.elastic.co';
 
 export const DOWNLOAD_SOURCE_SAVED_OBJECT_TYPE = 'ingest-download-sources';
 

x-pack/plugins/fleet/common/services/full_agent_policy_to_yaml.ts

@@ -23,6 +23,7 @@ const POLICY_KEYS_ORDER = [
   'use_output',
   'meta',
   'input',
+  'download',
 ];
 
 export const fullAgentPolicyToYaml = (policy: FullAgentPolicy, toYaml: typeof safeDump): string => {

x-pack/plugins/fleet/common/types/index.ts

@@ -38,7 +38,6 @@ export interface FleetConfigType {
   };
   developer?: {
     disableRegistryVersionCheck?: boolean;
-    allowAgentUpgradeSourceUri?: boolean;
     bundledPackageLocation?: string;
   };
 }

x-pack/plugins/fleet/common/types/models/agent.ts

@@ -48,6 +48,7 @@ export interface NewAgentAction {
   expiration?: string;
   start_time?: string;
   minimum_execution_duration?: number;
+  source_uri?: string;
 }
 
 export interface AgentAction extends NewAgentAction {

x-pack/plugins/fleet/common/types/models/agent_policy.ts

@@ -106,6 +106,7 @@ export interface FullAgentPolicy {
       metrics: boolean;
       logs: boolean;
     };
+    download: { source_uri: string };
   };
 }
 

x-pack/plugins/fleet/server/index.ts

@@ -141,7 +141,6 @@ export const config: PluginConfigDescriptor = {
     agentIdVerificationEnabled: schema.boolean({ defaultValue: true }),
     developer: schema.object({
       disableRegistryVersionCheck: schema.boolean({ defaultValue: false }),
-      allowAgentUpgradeSourceUri: schema.boolean({ defaultValue: false }),
       bundledPackageLocation: schema.string({ defaultValue: DEFAULT_BUNDLED_PACKAGE_LOCATION }),
     }),
     packageVerification: schema.object({

x-pack/plugins/fleet/server/routes/agent/helpers.test.ts

@@ -0,0 +1,196 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { savedObjectsClientMock } from '@kbn/core/server/mocks';
+
+import { DOWNLOAD_SOURCE_SAVED_OBJECT_TYPE } from '../../constants';
+
+import type { Agent, AgentPolicy, DownloadSource } from '../../types';
+
+import { agentPolicyService } from '../../services/agent_policy';
+
+import { getSourceUriForAgentPolicy, getSourceUriForAgent, getDefaultSourceUri } from './helpers';
+
+const soClientMock = savedObjectsClientMock.create();
+
+const mockedAgentPolicyService = agentPolicyService as jest.Mocked<typeof agentPolicyService>;
+jest.mock('../../services/agent_policy');
+
+jest.mock('../download_source', () => {
+  return {
+    downloadSourceService: {
+      getDefaultDownloadSourceId: async () => 'default-download-source-id',
+      get: async (soClient: any, id: string): Promise<DownloadSource> => {
+        if (id === 'test-ds-1') {
+          return {
+            id: 'test-ds-1',
+            is_default: false,
+            name: 'Test',
+            host: 'http://custom-registry-test',
+          };
+        }
+        return {
+          id: 'default-download-source-id',
+          is_default: true,
+          name: 'Default host',
+          host: 'http://default-registry.co',
+        };
+      },
+    },
+  };
+});
+
+function mockDownloadSourceSO(id: string, attributes: any = {}) {
+  return {
+    id,
+    type: DOWNLOAD_SOURCE_SAVED_OBJECT_TYPE,
+    references: [],
+    attributes: {
+      source_id: id,
+      ...attributes,
+    },
+  };
+}
+
+function mockAgentPolicy(data?: Partial<AgentPolicy>) {
+  mockedAgentPolicyService.get.mockImplementation((soClient, id, withPackagePolicies) =>
+    Promise.resolve({
+      id: 'agent-policy-id',
+      status: 'active',
+      is_managed: false,
+      namespace: 'default',
+      revision: 1,
+      name: 'Policy',
+      updated_at: '2022-01-01',
+      updated_by: 'qwerty',
+      ...data,
+    } as AgentPolicy)
+  );
+}
+describe('helpers', () => {
+  beforeEach(() => {
+    soClientMock.get.mockImplementation(async (type: string, id: string) => {
+      switch (id) {
+        case 'test-ds-1': {
+          return mockDownloadSourceSO('test-ds-1', {
+            is_default: false,
+            name: 'Test',
+            host: 'http://custom-registry-test',
+          });
+        }
+        case 'default-download-source-id': {
+          return mockDownloadSourceSO('default-download-source-id', {
+            is_default: true,
+            name: 'Default host',
+            host: 'http://default-registry.co',
+          });
+        }
+        default:
+          throw new Error('not found: ' + id);
+      }
+    });
+    soClientMock.find.mockResolvedValue({
+      saved_objects: [
+        {
+          id: 'default-download-source-id',
+          is_default: true,
+          attributes: {
+            download_source_id: 'test-source-id',
+          },
+        },
+        {
+          id: 'test-ds-1',
+          attributes: {
+            download_source_id: 'test-ds-1',
+          },
+        },
+      ],
+    } as any);
+  });
+  describe('getSourceUriForAgentPolicy', () => {
+    it('should return the source_uri set on an agent policy ', async () => {
+      const agentPolicy: AgentPolicy = {
+        id: 'agent-policy-id',
+        status: 'active',
+        package_policies: [],
+        is_managed: false,
+        namespace: 'default',
+        revision: 1,
+        name: 'Policy',
+        updated_at: '2022-01-01',
+        updated_by: 'qwerty',
+        download_source_id: 'test-ds-1',
+      };
+
+      expect(await getSourceUriForAgentPolicy(soClientMock, agentPolicy)).toEqual(
+        'http://custom-registry-test'
+      );
+    });
+    it('should return the default source_uri if there is none set on the agent policy ', async () => {
+      const agentPolicy: AgentPolicy = {
+        id: 'agent-policy-id',
+        status: 'active',
+        package_policies: [],
+        is_managed: false,
+        namespace: 'default',
+        revision: 1,
+        name: 'Policy',
+        updated_at: '2022-01-01',
+        updated_by: 'qwerty',
+      };
+
+      expect(await getSourceUriForAgentPolicy(soClientMock, agentPolicy)).toEqual(
+        'http://default-registry.co'
+      );
+    });
+  });
+
+  describe('getSourceUriForAgent', () => {
+    beforeEach(() => {
+      mockedAgentPolicyService.get.mockReset();
+    });
+    it('should return the source_uri set on an agent policy for that agent', async () => {
+      mockAgentPolicy({ download_source_id: 'test-ds-1' });
+      const agent: Agent = {
+        id: 'de9006e1-54a7-4320-b24e-927e6fe518a8',
+        active: true,
+        policy_id: 'agent-policy-id',
+        type: 'PERMANENT',
+        enrolled_at: '2022-09-30T20:24:08.347Z',
+        user_provided_metadata: {},
+        local_metadata: {},
+        packages: ['system'],
+        status: 'online',
+      };
+
+      expect(await getSourceUriForAgent(soClientMock, agent)).toEqual(
+        'http://custom-registry-test'
+      );
+    });
+  });
+  it('should return the default source_uri set if there is none set on that policy', async () => {
+    mockAgentPolicy();
+    const agent: Agent = {
+      id: 'de9006e1-54a7-4320-b24e-927e6fe518a8',
+      active: true,
+      policy_id: 'agent-policy-id',
+      type: 'PERMANENT',
+      enrolled_at: '2022-09-30T20:24:08.347Z',
+      user_provided_metadata: {},
+      local_metadata: {},
+      packages: ['system'],
+      status: 'online',
+    };
+    expect(await getSourceUriForAgent(soClientMock, agent)).toEqual('http://default-registry.co');
+  });
+
+  describe('getDefaultSourceUri', () => {
+    it('should return the global default source_uri', async () => {
+      expect(await getDefaultSourceUri(soClientMock)).toEqual('http://default-registry.co');
+    });
+  });
+});

x-pack/plugins/fleet/server/routes/agent/helpers.ts

@@ -0,0 +1,52 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import type { SavedObjectsClientContract } from '@kbn/core/server';
+
+import { downloadSourceService, agentPolicyService } from '../../services';
+import type { Agent, AgentPolicy } from '../../types';
+
+export const getSourceUriForAgent = async (soClient: SavedObjectsClientContract, agent: Agent) => {
+  if (!agent.policy_id) {
+    throw new Error('Agent.policy_id not found');
+  }
+  const agentPolicy = await agentPolicyService.get(soClient, agent.policy_id, false);
+  if (!agentPolicy) {
+    throw new Error('Agent Policy not found');
+  }
+  return getSourceUriForAgentPolicy(soClient, agentPolicy);
+};
+
+export const getSourceUriForAgentPolicy = async (
+  soClient: SavedObjectsClientContract,
+  agentPolicy: AgentPolicy
+) => {
+  const defaultDownloadSourceId = await downloadSourceService.getDefaultDownloadSourceId(soClient);
+
+  if (!defaultDownloadSourceId) {
+    throw new Error('Default download source host is not setup');
+  }
+  const downloadSourceId: string = agentPolicy.download_source_id || defaultDownloadSourceId;
+  const downloadSource = await downloadSourceService.get(soClient, downloadSourceId);
+  if (!downloadSource) {
+    throw new Error(`Download source host not found ${downloadSourceId}`);
+  }
+  return downloadSource.host;
+};
+
+export const getDefaultSourceUri = async (soClient: SavedObjectsClientContract) => {
+  const defaultDownloadSourceId = await downloadSourceService.getDefaultDownloadSourceId(soClient);
+
+  if (!defaultDownloadSourceId) {
+    throw new Error('Default download source host is not setup');
+  }
+  const defaultDownloadSource = await downloadSourceService.get(soClient, defaultDownloadSourceId);
+  if (!defaultDownloadSource) {
+    throw new Error(`Download source host not found ${defaultDownloadSourceId}`);
+  }
+  return defaultDownloadSource.host;
+};

x-pack/plugins/fleet/server/routes/agent/upgrade_handler.ts

@@ -27,6 +27,8 @@ import type { Agent } from '../../types';
 
 import { getAllFleetServerAgents } from '../../collectors/get_all_fleet_server_agents';
 
+import { getSourceUriForAgent } from './helpers';
+
 export const postAgentUpgradeHandler: RequestHandler<
   TypeOf<typeof PostAgentUpgradeRequestSchema.params>,
   undefined,
@@ -35,11 +37,10 @@ export const postAgentUpgradeHandler: RequestHandler<
   const coreContext = await context.core;
   const soClient = coreContext.savedObjects.client;
   const esClient = coreContext.elasticsearch.client.asInternalUser;
-  const { version, source_uri: sourceUri, force } = request.body;
+  const { version, source_uri: requestedSourceUri, force } = request.body;
   const kibanaVersion = appContextService.getKibanaVersion();
   try {
     checkKibanaVersion(version, kibanaVersion);
-    checkSourceUriAllowed(sourceUri);
   } catch (err) {
     return response.customError({
       statusCode: 400,
@@ -48,7 +49,25 @@ export const postAgentUpgradeHandler: RequestHandler<
       },
     });
   }
+
   const agent = await getAgentById(esClient, request.params.agentId);
+
+  let sourceUri;
+  if (requestedSourceUri) {
+    sourceUri = requestedSourceUri;
+  } else {
+    try {
+      sourceUri = await getSourceUriForAgent(soClient, agent);
+    } catch (err) {
+      return response.customError({
+        statusCode: 400,
+        body: {
+          message: err.message,
+        },
+      });
+    }
+  }
+
   if (agent.unenrollment_started_at || agent.unenrolled_at) {
     return response.customError({
       statusCode: 400,
@@ -102,7 +121,6 @@ export const postBulkAgentsUpgradeHandler: RequestHandler<
   const kibanaVersion = appContextService.getKibanaVersion();
   try {
     checkKibanaVersion(version, kibanaVersion);
-    checkSourceUriAllowed(sourceUri);
     const fleetServerAgents = await getAllFleetServerAgents(soClient, esClient);
     checkFleetServerVersion(version, fleetServerAgents);
   } catch (err) {
@@ -167,14 +185,6 @@ export const checkKibanaVersion = (version: string, kibanaVersion: string) => {
     );
 };
 
-const checkSourceUriAllowed = (sourceUri?: string) => {
-  if (sourceUri && !appContextService.getConfig()?.developer?.allowAgentUpgradeSourceUri) {
-    throw new Error(
-      `source_uri is not allowed or recommended in production. Set xpack.fleet.developer.allowAgentUpgradeSourceUri in kibana.yml to true.`
-    );
-  }
-};
-
 // Check the installed fleet server version
 const checkFleetServerVersion = (versionToUpgradeNumber: string, fleetServerAgents: Agent[]) => {
   const fleetServerVersions = fleetServerAgents.map(