chore(atlas-service): add validation for backend response (#4745)

* chore(atlas-service): add validation for backend response

* chore(atlas-service): use some for boolean array check; fix test description

Co-authored-by: Anna Henningsen <[email protected]>

* chore(atlas-service): use inspect for better error messages

Co-authored-by: Anna Henningsen <[email protected]>

---------

Co-authored-by: Anna Henningsen <[email protected]>

Author: gribnoysup

packages/atlas-service/src/main.spec.ts

@@ -14,20 +14,6 @@ function getListenerCount(emitter: EventEmitter) {
   }, 0);
 }
 
-const atlasAIServiceTests: {
-  functionName: 'getQueryFromUserInput' | 'getAggregationFromUserInput';
-  aiEndpoint: string;
-}[] = [
-  {
-    functionName: 'getQueryFromUserInput',
-    aiEndpoint: 'mql-query',
-  },
-  {
-    functionName: 'getAggregationFromUserInput',
-    aiEndpoint: 'mql-aggregation',
-  },
-];
-
 describe('AtlasServiceMain', function () {
   const sandbox = Sinon.createSandbox();
 
@@ -216,15 +202,46 @@ describe('AtlasServiceMain', function () {
     });
   });
 
-  for (const { functionName, aiEndpoint } of atlasAIServiceTests) {
+  const atlasAIServiceTests = [
+    {
+      functionName: 'getQueryFromUserInput',
+      aiEndpoint: 'mql-query',
+      responses: {
+        success: {
+          content: { query: { filter: "{ test: 'pineapple' }" } },
+        },
+        invalid: [
+          {},
+          { countent: {} },
+          { content: { qooery: {} } },
+          { content: { query: { filter: { foo: 1 } } } },
+        ],
+      },
+    },
+    {
+      functionName: 'getAggregationFromUserInput',
+      aiEndpoint: 'mql-aggregation',
+      responses: {
+        success: {
+          content: { aggregation: { pipeline: "[{ test: 'pineapple' }]" } },
+        },
+        invalid: [
+          {},
+          { content: { aggregation: {} } },
+          { content: { aggrogation: {} } },
+          { content: { aggregation: { pipeline: true } } },
+        ],
+      },
+    },
+  ] as const;
+
+  for (const { functionName, aiEndpoint, responses } of atlasAIServiceTests) {
     describe(functionName, function () {
       it('makes a post request with the user input to the endpoint in the environment', async function () {
         AtlasService['fetch'] = sandbox.stub().resolves({
           ok: true,
           json() {
-            return Promise.resolve({
-              content: { query: { filter: "{ test: 'pineapple' }" } },
-            });
+            return Promise.resolve(responses.success);
           },
         }) as any;
 
@@ -246,10 +263,28 @@ describe('AtlasServiceMain', function () {
         expect(args[1].body).to.eq(
           '{"userInput":"test","collectionName":"jam","databaseName":"peanut","schema":{"_id":{"types":[{"bsonType":"ObjectId"}]}},"sampleDocuments":[{"_id":1234}]}'
         );
-        expect(res).to.have.nested.property(
-          'content.query.filter',
-          "{ test: 'pineapple' }"
-        );
+        expect(res).to.deep.eq(responses.success);
+      });
+
+      it('should fail when response is not matching expected schema', async function () {
+        for (const res of responses.invalid) {
+          AtlasService['fetch'] = sandbox.stub().resolves({
+            ok: true,
+            json() {
+              return Promise.resolve(res);
+            },
+          }) as any;
+          try {
+            await AtlasService[functionName]({
+              userInput: 'test',
+              collectionName: 'test',
+              databaseName: 'peanut',
+            });
+            expect.fail(`Expected ${functionName} to throw`);
+          } catch (err) {
+            expect((err as Error).message).to.match(/Unexpected.+?response/);
+          }
+        }
       });
 
       it('uses the abort signal in the fetch request', async function () {
@@ -289,7 +324,7 @@ describe('AtlasServiceMain', function () {
         AtlasService['fetch'] = sandbox.stub().resolves({
           ok: true,
           json() {
-            return Promise.resolve({});
+            return Promise.resolve(responses.success);
           },
         }) as any;
 
@@ -356,13 +391,13 @@ describe('AtlasServiceMain', function () {
         AtlasService['fetch'] = sandbox.stub().resolves({
           ok: true,
           json() {
-            return Promise.resolve({ test: 1 });
+            return Promise.resolve(responses.success);
           },
         }) as any;
         AtlasService['oidcPluginLogger'].emit(
           'mongodb-oidc-plugin:refresh-started'
         );
-        const [query] = await Promise.all([
+        const [res] = await Promise.all([
           AtlasService[functionName]({
             userInput: 'test',
             collectionName: 'test',
@@ -378,7 +413,7 @@ describe('AtlasServiceMain', function () {
             );
           })(),
         ]);
-        expect(query).to.deep.eq({ test: 1 });
+        expect(res).to.deep.eq(responses.success);
       });
     });
   }

packages/atlas-service/src/main.ts

@@ -19,12 +19,14 @@ import type { Response } from 'node-fetch';
 import fetch from 'node-fetch';
 import type { SimplifiedSchema } from 'mongodb-schema';
 import type { Document } from 'mongodb';
-import type {
-  AIAggregation,
-  AIQuery,
-  IntrospectInfo,
-  Token,
-  UserInfo,
+import {
+  validateAIQueryResponse,
+  type AIAggregation,
+  type AIQuery,
+  type IntrospectInfo,
+  type Token,
+  type UserInfo,
+  validateAIAggregationResponse,
 } from './util';
 import {
   broadcast,
@@ -733,7 +735,7 @@ export class AtlasService {
     schema?: SimplifiedSchema;
     sampleDocuments?: Document[];
     signal?: AbortSignal;
-  }) {
+  }): Promise<AIAggregation> {
     throwIfAborted(signal);
     throwIfNetworkTrafficDisabled();
 
@@ -780,7 +782,11 @@ export class AtlasService {
 
     await throwIfNotOk(res);
 
-    return res.json() as Promise<AIAggregation>;
+    const body = await res.json();
+
+    validateAIAggregationResponse(body);
+
+    return body;
   }
 
   static async getQueryFromUserInput({
@@ -797,7 +803,7 @@ export class AtlasService {
     schema?: SimplifiedSchema;
     sampleDocuments?: Document[];
     signal?: AbortSignal;
-  }) {
+  }): Promise<AIQuery> {
     throwIfAborted(signal);
     throwIfNetworkTrafficDisabled();
 
@@ -841,7 +847,11 @@ export class AtlasService {
 
     await throwIfNotOk(res);
 
-    return res.json() as Promise<AIQuery>;
+    const body = await res.json();
+
+    validateAIQueryResponse(body);
+
+    return body;
   }
 
   static async onExit() {

packages/atlas-service/src/renderer.ts

@@ -110,4 +110,11 @@ export class AtlasService implements AtlasServiceEmitter {
 
 export { AtlasSignIn } from './components/atlas-signin';
 
-export type { UserInfo, IntrospectInfo, Token } from './util';
+export type {
+  UserInfo,
+  IntrospectInfo,
+  Token,
+  AtlasServiceNetworkError,
+  AIQuery,
+  AIAggregation,
+} from './util';

packages/atlas-service/src/util.ts

@@ -1,4 +1,5 @@
 import type * as plugin from '@mongodb-js/oidc-plugin';
+import util from 'util';
 
 export type UserInfo = {
   firstName: string;
@@ -11,16 +12,111 @@ export type IntrospectInfo = { active: boolean };
 
 export type Token = plugin.IdPServerResponse;
 
+function hasExtraneousKeys(obj: any, expectedKeys: string[]) {
+  return Object.keys(obj).some((key) => !expectedKeys.includes(key));
+}
+
 export type AIAggregation = {
-  content?: {
+  content: {
     aggregation?: {
-      pipeline?: unknown;
+      pipeline?: string;
     };
   };
 };
 
+export function validateAIAggregationResponse(
+  response: any
+): asserts response is AIAggregation {
+  const { content } = response;
+
+  if (typeof content !== 'object' || content === null) {
+    throw new Error('Unexpected response: expected content to be an object');
+  }
+
+  if (hasExtraneousKeys(content, ['aggregation'])) {
+    throw new Error('Unexpected keys in response: expected aggregation');
+  }
+
+  if (content.aggregation && typeof content.aggregation.pipeline !== 'string') {
+    // Compared to queries where we will always get the `query` field, for
+    // aggregations backend deletes the whole `aggregation` key if pipeline is
+    // empty, so we only validate `pipeline` key if `aggregation` key is present
+    throw new Error(
+      `Unexpected response: expected aggregation to be a string, got ${String(
+        content.aggregation.pipeline
+      )}`
+    );
+  }
+}
+
 export type AIQuery = {
-  content?: {
-    query?: unknown;
+  content: {
+    query: Record<
+      'filter' | 'project' | 'collation' | 'sort' | 'skip' | 'limit',
+      string
+    > & { aggregation?: { pipeline: string } };
   };
 };
+
+export function validateAIQueryResponse(
+  response: any
+): asserts response is AIQuery {
+  const { content } = response ?? {};
+
+  if (typeof content !== 'object' || content === null) {
+    throw new Error('Unexpected response: expected content to be an object');
+  }
+
+  if (hasExtraneousKeys(content, ['query'])) {
+    throw new Error('Unexpected keys in response: expected query');
+  }
+
+  const { query } = content;
+
+  if (typeof query !== 'object' || query === null) {
+    throw new Error('Unexpected response: expected query to be an object');
+  }
+
+  if (
+    hasExtraneousKeys(query, [
+      'filter',
+      'project',
+      'collation',
+      'sort',
+      'skip',
+      'limit',
+      'aggregation',
+    ])
+  ) {
+    throw new Error(
+      'Unexpected keys in response: expected filter, project, collation, sort, skip, limit, aggregation'
+    );
+  }
+
+  for (const field of [
+    'filter',
+    'project',
+    'collation',
+    'sort',
+    'skip',
+    'limit',
+  ]) {
+    if (query[field] && typeof query[field] !== 'string') {
+      throw new Error(
+        `Unexpected response: expected field ${field} to be a string, got ${util.inspect(
+          query[field]
+        )}`
+      );
+    }
+  }
+
+  if (query.aggregation && typeof query.aggregation.pipeline !== 'string') {
+    throw new Error(
+      `Unexpected response: expected aggregation pipeline to be a string, got ${util.inspect(
+        query.aggregation
+      )}`
+    );
+  }
+}
+
+export type AtlasServiceNetworkError = Error & { statusCode: number };

packages/compass-aggregations/src/modules/pipeline-builder/pipeline-ai.spec.ts

@@ -28,7 +28,7 @@ describe('AIPipelineReducer', function () {
       it('should succeed', async function () {
         const mockAtlasService = {
           getAggregationFromUserInput: sandbox.stub().resolves({
-            content: { aggregation: { pipeline: [{ $match: { _id: 1 } }] } },
+            content: { aggregation: { pipeline: '[{ $match: { _id: 1 } }]' } },
           }),
         };