[9.0] [Obs AI Assistant] Fix bug with kibana tool when using a proxy (#236653)

x-pack/platform/plugins/shared/observability_ai_assistant/server/functions/kibana.test.ts

@@ -0,0 +1,133 @@
+/*
+ * 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 { AxiosRequestConfig } from 'axios';
+import axios from 'axios';
+import { registerKibanaFunction } from './kibana';
+import type { FunctionRegistrationParameters } from '.';
+
+jest.mock('axios');
+const mockedAxios = jest.mocked(axios);
+
+function registerFunction(overrides: {
+  publicBaseUrl?: string;
+  requestUrl?: URL;
+  rewrittenUrl?: URL;
+  headers?: Record<string, string>;
+}) {
+  const logger = { info: jest.fn(), error: jest.fn() };
+  const coreStart = {
+    http: {
+      basePath: {
+        publicBaseUrl: overrides.publicBaseUrl ?? 'https://kibana.example.com:5601',
+        serverBasePath: '',
+        get: jest.fn(),
+      },
+    },
+  };
+
+  const resources = {
+    request: {
+      url:
+        overrides.requestUrl ??
+        new URL('https://source.example/internal/observability_ai_assistant/chat/complete'),
+      rewrittenUrl: overrides.rewrittenUrl,
+      headers: {
+        'content-type': 'application/json',
+        host: 'attacker.example',
+        origin: 'https://attacker.example',
+        ...(overrides.headers ?? {}),
+      },
+    },
+    logger,
+    plugins: {
+      core: {
+        start: jest.fn().mockResolvedValue(coreStart),
+      },
+    },
+  };
+
+  const functions = { registerFunction: jest.fn() };
+  registerKibanaFunction({ functions, resources } as unknown as FunctionRegistrationParameters);
+
+  return {
+    handler: functions.registerFunction.mock.calls[0][1],
+    coreStart,
+    resources,
+  };
+}
+
+describe('kibana tool', () => {
+  beforeEach(() => {
+    jest.clearAllMocks();
+    mockedAxios.mockResolvedValue({ data: { ok: true } });
+  });
+
+  it('forwards requests to the configured publicBaseUrl host only', async () => {
+    const { handler } = registerFunction({
+      headers: {
+        host: 'malicious-host:9200',
+        origin: 'https://malicious-host:9200',
+        'x-forwarded-host': 'another-host',
+      },
+    });
+
+    await handler({
+      arguments: {
+        method: 'GET',
+        pathname: '/api/saved_objects/_find',
+        query: { type: 'dashboard' },
+      },
+    });
+
+    const forwardedRequest = mockedAxios.mock.calls[0][0] as AxiosRequestConfig;
+    expect(forwardedRequest.url).toBe(
+      'https://kibana.example.com:5601/api/saved_objects/_find?type=dashboard'
+    );
+    expect(forwardedRequest.url).not.toContain('malicious-host');
+  });
+
+  it('builds the forwarded url using the space from the incoming request path', async () => {
+    const { handler } = registerFunction({
+      requestUrl: new URL(
+        'https://source.example/s/my-space/internal/observability_ai_assistant/chat/complete'
+      ),
+    });
+
+    await handler({
+      arguments: {
+        method: 'POST',
+        pathname: '/api/apm/agent_keys',
+        body: { foo: 'bar' },
+      },
+    });
+
+    expect(mockedAxios).toHaveBeenCalledWith(
+      expect.objectContaining({
+        url: 'https://kibana.example.com:5601/s/my-space/api/apm/agent_keys',
+        data: JSON.stringify({ foo: 'bar' }),
+      })
+    );
+  });
+
+  it('throws when server.publicBaseUrl is not configured', async () => {
+    const { handler, coreStart } = registerFunction({});
+    coreStart.http.basePath.publicBaseUrl = undefined as any;
+
+    await expect(
+      handler({
+        arguments: {
+          method: 'GET',
+          pathname: '/api/saved_objects/_find',
+          query: { type: 'dashboard' },
+        },
+      })
+    ).rejects.toThrow(
+      'Cannot invoke Kibana tool: "server.publicBaseUrl" must be configured in kibana.yml'
+    );
+  });
+});

x-pack/platform/plugins/shared/observability_ai_assistant/server/functions/kibana.ts

@@ -6,9 +6,10 @@
  */
 
 import axios from 'axios';
-import { format, parse } from 'url';
-import { castArray, first, pick, pickBy } from 'lodash';
+import { format } from 'url';
+import { pickBy } from 'lodash';
 import type { KibanaRequest } from '@kbn/core/server';
+import { addSpaceIdToPath, getSpaceIdFromPath } from '@kbn/spaces-plugin/common';
 import type { FunctionRegistrationParameters } from '.';
 
 export function registerKibanaFunction({
@@ -47,24 +48,43 @@ export function registerKibanaFunction({
         required: ['method', 'pathname'] as const,
       },
     },
-    ({ arguments: { method, pathname, body, query } }, signal) => {
-      const { request } = resources;
+    async ({ arguments: { method, pathname, body, query } }, signal) => {
+      const { request, logger } = resources;
+      const requestUrl = request.rewrittenUrl || request.url;
+      const core = await resources.plugins.core.start();
 
-      const { protocol, host, pathname: pathnameFromRequest } = request.rewrittenUrl || request.url;
+      function getParsedPublicBaseUrl() {
+        const { publicBaseUrl } = core.http.basePath;
+        if (!publicBaseUrl) {
+          const errorMessage = `Cannot invoke Kibana tool: "server.publicBaseUrl" must be configured in kibana.yml`;
+          logger.error(errorMessage);
+          throw new Error(errorMessage);
+        }
+        const parsedBaseUrl = new URL(publicBaseUrl);
+        return parsedBaseUrl;
+      }
 
-      const origin = first(castArray(request.headers.origin));
+      function getPathnameWithSpaceId() {
+        const { serverBasePath } = core.http.basePath;
+        const { spaceId } = getSpaceIdFromPath(requestUrl.pathname, serverBasePath);
+        const pathnameWithSpaceId = addSpaceIdToPath(serverBasePath, spaceId, pathname);
+        return pathnameWithSpaceId;
+      }
 
+      const parsedPublicBaseUrl = getParsedPublicBaseUrl();
       const nextUrl = {
-        host,
-        protocol,
-        ...(origin ? pick(parse(origin), 'host', 'protocol') : {}),
-        pathname: pathnameFromRequest.replace(
-          '/internal/observability_ai_assistant/chat/complete',
-          pathname
-        ),
+        host: parsedPublicBaseUrl.host,
+        protocol: parsedPublicBaseUrl.protocol,
+        pathname: getPathnameWithSpaceId(),
         query: query ? (query as Record<string, string>) : undefined,
       };
 
+      logger.info(
+        `Calling Kibana API by forwarding request from "${requestUrl}" to: "${method} ${format(
+          nextUrl
+        )}"`
+      );
+
       const copiedHeaderNames = [
         'accept-encoding',
         'accept-language',
@@ -87,15 +107,19 @@ export function registerKibanaFunction({
         );
       });
 
-      return axios({
-        method,
-        headers,
-        url: format(nextUrl),
-        data: body ? JSON.stringify(body) : undefined,
-        signal,
-      }).then((response) => {
+      try {
+        const response = await axios({
+          method,
+          headers,
+          url: format(nextUrl),
+          data: body ? JSON.stringify(body) : undefined,
+          signal,
+        });
         return { content: response.data };
-      });
+      } catch (e) {
+        logger.error(`Error calling Kibana API: ${method} ${format(nextUrl)}. Failed with ${e}`);
+        throw e;
+      }
     }
   );
 }