test: add validation tests for Elasticsearch tool operations

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

@@ -9,6 +9,24 @@ import type { FunctionRegistrationParameters } from '.';
 
 export const ELASTICSEARCH_FUNCTION_NAME = 'elasticsearch';
 
+/**
+ * Determines if the given HTTP method and path are allowed for the Elasticsearch tool.
+ * Only GET requests and POST requests to the "_search" endpoint are permitted.
+ *
+ * @param method - The HTTP method (GET, POST, PUT, DELETE, PATCH)
+ * @param path - The Elasticsearch API path
+ * @returns true if the operation is allowed, false otherwise
+ */
+function isOperationAllowed(method: string, path: string): boolean {
+  // Allowlist: (1) all GET requests, (2) POST requests whose *final* path segment is exactly "_search".
+  const [pathWithoutQuery] = path.split('?');
+  const pathSegments = pathWithoutQuery.replace(/^\//, '').split('/');
+  const lastPathSegment = pathSegments[pathSegments.length - 1];
+  const isSearchEndpoint = lastPathSegment === '_search';
+
+  return method === 'GET' || (method === 'POST' && isSearchEndpoint);
+}
+
 export function registerElasticsearchFunction({
   functions,
   resources,
@@ -52,13 +70,7 @@ export function registerElasticsearchFunction({
       },
     },
     async ({ arguments: { method, path, body } }) => {
-      // Allowlist: (1) all GET requests, (2) POST requests whose *final* path segment is exactly "_search".
-      const [pathWithoutQuery] = path.split('?');
-      const pathSegments = pathWithoutQuery.replace(/^\//, '').split('/');
-      const lastPathSegment = pathSegments[pathSegments.length - 1];
-      const isSearchEndpoint = lastPathSegment === '_search';
-
-      if (method !== 'GET' && !(method === 'POST' && isSearchEndpoint)) {
+      if (!isOperationAllowed(method, path)) {
         throw new Error(
           'Only GET requests or POST requests to the "_search" endpoint are permitted through this assistant function.'
         );

x-pack/solutions/observability/test/api_integration_deployment_agnostic/apis/ai_assistant/elasticsearch_tool/elasticsearch_tool.spec.ts

@@ -0,0 +1,328 @@
+import expect from '@kbn/expect';
+import { MessageRole, type Message } from '@kbn/observability-ai-assistant-plugin/common';
+import { ELASTICSEARCH_FUNCTION_NAME } from '@kbn/observability-ai-assistant-plugin/server/functions/elasticsearch';
+import { createLlmProxy, LlmProxy } from '../utils/create_llm_proxy';
+import {
+  getMessageAddedEvents,
+  invokeChatCompleteWithFunctionRequest,
+} from '../utils/conversation';
+import type { DeploymentAgnosticFtrProviderContext } from '../../../ftr_provider_context';
+
+export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderContext) {
+  const log = getService('log');
+  const observabilityAIAssistantAPIClient = getService('observabilityAIAssistantApi');
+
+  describe('tool: elasticsearch tool restrictions', function () {
+    // Fails on MKI: https://github.com/elastic/kibana/issues/205581
+    this.tags(['skipCloud']);
+    let proxy: LlmProxy;
+    let connectorId: string;
+
+    before(async () => {
+      proxy = await createLlmProxy(log);
+      connectorId = await observabilityAIAssistantAPIClient.createProxyActionConnector({
+        port: proxy.getPort(),
+      });
+    });
+
+    after(async () => {
+      proxy.close();
+      await observabilityAIAssistantAPIClient.deleteActionConnector({
+        actionId: connectorId,
+      });
+    });
+
+    describe('Allowed Operations', () => {
+      it('should allow GET requests to any endpoint', async () => {
+        void proxy.interceptWithResponse('Hello from LLM Proxy');
+
+        const responseBody = await invokeChatCompleteWithFunctionRequest({
+          connectorId,
+          observabilityAIAssistantAPIClient,
+          functionCall: {
+            name: ELASTICSEARCH_FUNCTION_NAME,
+            trigger: MessageRole.User,
+            arguments: JSON.stringify({
+              method: 'GET',
+              path: '/_cluster/health',
+            }),
+          },
+        });
+
+        await proxy.waitForAllInterceptorsToHaveBeenCalled();
+
+        const events = getMessageAddedEvents(responseBody);
+        const esFunctionResponse = events[0];
+
+        expect(esFunctionResponse.message.message.name).to.be(ELASTICSEARCH_FUNCTION_NAME);
+        // Function should succeed - response should contain actual data, not an error
+        const content = JSON.parse(esFunctionResponse.message.message.content!);
+        expect(content).to.have.property('response');
+        expect(content.response).to.not.have.property('error');
+      });
+
+      it('should allow POST requests to _search endpoint', async () => {
+        void proxy.interceptWithResponse('Hello from LLM Proxy');
+
+        const responseBody = await invokeChatCompleteWithFunctionRequest({
+          connectorId,
+          observabilityAIAssistantAPIClient,
+          functionCall: {
+            name: ELASTICSEARCH_FUNCTION_NAME,
+            trigger: MessageRole.User,
+            arguments: JSON.stringify({
+              method: 'POST',
+              path: '/_search',
+              body: { query: { match_all: {} } },
+            }),
+          },
+        });
+
+        await proxy.waitForAllInterceptorsToHaveBeenCalled();
+
+        const events = getMessageAddedEvents(responseBody);
+        const esFunctionResponse = events[0];
+
+        expect(esFunctionResponse.message.message.name).to.be(ELASTICSEARCH_FUNCTION_NAME);
+        // Function should succeed - response should contain actual data, not an error
+        const content = JSON.parse(esFunctionResponse.message.message.content!);
+        expect(content).to.have.property('response');
+        expect(content.response).to.not.have.property('error');
+      });
+
+      it('should allow POST requests to nested _search endpoints', async () => {
+        void proxy.interceptWithResponse('Hello from LLM Proxy');
+
+        const responseBody = await invokeChatCompleteWithFunctionRequest({
+          connectorId,
+          observabilityAIAssistantAPIClient,
+          functionCall: {
+            name: ELASTICSEARCH_FUNCTION_NAME,
+            trigger: MessageRole.User,
+            arguments: JSON.stringify({
+              method: 'POST',
+              path: '/my-index/_search',
+              body: { query: { match_all: {} } },
+            }),
+          },
+        });
+
+        await proxy.waitForAllInterceptorsToHaveBeenCalled();
+
+        const events = getMessageAddedEvents(responseBody);
+        const esFunctionResponse = events[0];
+
+        expect(esFunctionResponse.message.message.name).to.be(ELASTICSEARCH_FUNCTION_NAME);
+        // Function should succeed - response should contain actual data, not an error
+        const content = JSON.parse(esFunctionResponse.message.message.content!);
+        expect(content).to.have.property('response');
+        expect(content.response).to.not.have.property('error');
+      });
+    });
+
+    describe('Disallowed Operations', () => {
+      it('should reject PUT requests', async () => {
+        void proxy.interceptWithResponse('Hello from LLM Proxy');
+
+        const responseBody = await invokeChatCompleteWithFunctionRequest({
+          connectorId,
+          observabilityAIAssistantAPIClient,
+          functionCall: {
+            name: ELASTICSEARCH_FUNCTION_NAME,
+            trigger: MessageRole.User,
+            arguments: JSON.stringify({
+              method: 'PUT',
+              path: '/my-index',
+              body: { mappings: { properties: {} } },
+            }),
+          },
+        });
+
+        await proxy.waitForAllInterceptorsToHaveBeenCalled();
+
+        const events = getMessageAddedEvents(responseBody);
+        const esFunctionResponse = events[0];
+
+        expect(esFunctionResponse.message.message.name).to.be(ELASTICSEARCH_FUNCTION_NAME);
+        // Function should fail - response should contain an error
+        const content = JSON.parse(esFunctionResponse.message.message.content!);
+        expect(content).to.have.property('error');
+        expect(content.error).to.contain('Only GET requests or POST requests to the "_search" endpoint are permitted');
+      });
+
+      it('should reject DELETE requests', async () => {
+        void proxy.interceptWithResponse('Hello from LLM Proxy');
+
+        const responseBody = await invokeChatCompleteWithFunctionRequest({
+          connectorId,
+          observabilityAIAssistantAPIClient,
+          functionCall: {
+            name: ELASTICSEARCH_FUNCTION_NAME,
+            trigger: MessageRole.User,
+            arguments: JSON.stringify({
+              method: 'DELETE',
+              path: '/my-index',
+            }),
+          },
+        });
+
+        await proxy.waitForAllInterceptorsToHaveBeenCalled();
+
+        const events = getMessageAddedEvents(responseBody);
+        const esFunctionResponse = events[0];
+
+        expect(esFunctionResponse.message.message.name).to.be(ELASTICSEARCH_FUNCTION_NAME);
+        // Function should fail - response should contain an error
+        const content = JSON.parse(esFunctionResponse.message.message.content!);
+        expect(content).to.have.property('error');
+        expect(content.error).to.contain('Only GET requests or POST requests to the "_search" endpoint are permitted');
+      });
+
+      it('should reject PATCH requests', async () => {
+        void proxy.interceptWithResponse('Hello from LLM Proxy');
+
+        const responseBody = await invokeChatCompleteWithFunctionRequest({
+          connectorId,
+          observabilityAIAssistantAPIClient,
+          functionCall: {
+            name: ELASTICSEARCH_FUNCTION_NAME,
+            trigger: MessageRole.User,
+            arguments: JSON.stringify({
+              method: 'PATCH',
+              path: '/my-index/_doc/1',
+              body: { doc: { field: 'value' } },
+            }),
+          },
+        });
+
+        await proxy.waitForAllInterceptorsToHaveBeenCalled();
+
+        const events = getMessageAddedEvents(responseBody);
+        const esFunctionResponse = events[0];
+
+        expect(esFunctionResponse.message.message.name).to.be(ELASTICSEARCH_FUNCTION_NAME);
+        // Function should fail - response should contain an error
+        const content = JSON.parse(esFunctionResponse.message.message.content!);
+        expect(content).to.have.property('error');
+        expect(content.error).to.contain('Only GET requests or POST requests to the "_search" endpoint are permitted');
+      });
+
+      it('should reject POST requests to non-search endpoints', async () => {
+        void proxy.interceptWithResponse('Hello from LLM Proxy');
+
+        const responseBody = await invokeChatCompleteWithFunctionRequest({
+          connectorId,
+          observabilityAIAssistantAPIClient,
+          functionCall: {
+            name: ELASTICSEARCH_FUNCTION_NAME,
+            trigger: MessageRole.User,
+            arguments: JSON.stringify({
+              method: 'POST',
+              path: '/my-index/_doc',
+              body: { field: 'value' },
+            }),
+          },
+        });
+
+        await proxy.waitForAllInterceptorsToHaveBeenCalled();
+
+        const events = getMessageAddedEvents(responseBody);
+        const esFunctionResponse = events[0];
+
+        expect(esFunctionResponse.message.message.name).to.be(ELASTICSEARCH_FUNCTION_NAME);
+        // Function should fail - response should contain an error
+        const content = JSON.parse(esFunctionResponse.message.message.content!);
+        expect(content).to.have.property('error');
+        expect(content.error).to.contain('Only GET requests or POST requests to the "_search" endpoint are permitted');
+      });
+
+      it('should reject POST requests to endpoints ending with _search but not being _search', async () => {
+        void proxy.interceptWithResponse('Hello from LLM Proxy');
+
+        const responseBody = await invokeChatCompleteWithFunctionRequest({
+          connectorId,
+          observabilityAIAssistantAPIClient,
+          functionCall: {
+            name: ELASTICSEARCH_FUNCTION_NAME,
+            trigger: MessageRole.User,
+            arguments: JSON.stringify({
+              method: 'POST',
+              path: '/my-index/custom_search',
+              body: { query: { match_all: {} } },
+            }),
+          },
+        });
+
+        await proxy.waitForAllInterceptorsToHaveBeenCalled();
+
+        const events = getMessageAddedEvents(responseBody);
+        const esFunctionResponse = events[0];
+
+        expect(esFunctionResponse.message.message.name).to.be(ELASTICSEARCH_FUNCTION_NAME);
+        // Function should fail - response should contain an error
+        const content = JSON.parse(esFunctionResponse.message.message.content!);
+        expect(content).to.have.property('error');
+        expect(content.error).to.contain('Only GET requests or POST requests to the "_search" endpoint are permitted');
+      });
+    });
+
+    describe('Edge Cases', () => {
+      it('should handle paths with query parameters correctly', async () => {
+        void proxy.interceptWithResponse('Hello from LLM Proxy');
+
+        const responseBody = await invokeChatCompleteWithFunctionRequest({
+          connectorId,
+          observabilityAIAssistantAPIClient,
+          functionCall: {
+            name: ELASTICSEARCH_FUNCTION_NAME,
+            trigger: MessageRole.User,
+            arguments: JSON.stringify({
+              method: 'GET',
+              path: '/_cluster/health?pretty=true',
+            }),
+          },
+        });
+
+        await proxy.waitForAllInterceptorsToHaveBeenCalled();
+
+        const events = getMessageAddedEvents(responseBody);
+        const esFunctionResponse = events[0];
+
+        expect(esFunctionResponse.message.message.name).to.be(ELASTICSEARCH_FUNCTION_NAME);
+        // Function should succeed - response should contain actual data, not an error
+        const content = JSON.parse(esFunctionResponse.message.message.content!);
+        expect(content).to.have.property('response');
+        expect(content.response).to.not.have.property('error');
+      });
+
+      it('should handle POST _search with query parameters', async () => {
+        void proxy.interceptWithResponse('Hello from LLM Proxy');
+
+        const responseBody = await invokeChatCompleteWithFunctionRequest({
+          connectorId,
+          observabilityAIAssistantAPIClient,
+          functionCall: {
+            name: ELASTICSEARCH_FUNCTION_NAME,
+            trigger: MessageRole.User,
+            arguments: JSON.stringify({
+              method: 'POST',
+              path: '/_search?scroll=1m',
+              body: { query: { match_all: {} } },
+            }),
+          },
+        });
+
+        await proxy.waitForAllInterceptorsToHaveBeenCalled();
+
+        const events = getMessageAddedEvents(responseBody);
+        const esFunctionResponse = events[0];
+
+        expect(esFunctionResponse.message.message.name).to.be(ELASTICSEARCH_FUNCTION_NAME);
+        // Function should succeed - response should contain actual data, not an error
+        const content = JSON.parse(esFunctionResponse.message.message.content!);
+        expect(content).to.have.property('response');
+        expect(content.response).to.not.have.property('error');
+      });
+    });
+  });
+}

x-pack/solutions/observability/test/api_integration_deployment_agnostic/apis/ai_assistant/index.ts

@@ -38,6 +38,7 @@ export default function aiAssistantApiIntegrationTests({
 
     // Misc.
     loadTestFile(require.resolve('./chat/chat.spec.ts'));
+    loadTestFile(require.resolve('./elasticsearch_tool/elasticsearch_tool.spec.ts'));
     loadTestFile(require.resolve('./complete/complete.spec.ts'));
     loadTestFile(require.resolve('./index_assets/index_assets.spec.ts'));
     loadTestFile(require.resolve('./connectors/connectors.spec.ts'));