MONGOSH-523 - AutoEncryption only run on enterprise

packages/service-provider-server/src/cli-service-provider.spec.ts

@@ -3,7 +3,7 @@ import chai, { expect } from 'chai';
 import { Collection, Db, MongoClient } from 'mongodb';
 import sinonChai from 'sinon-chai';
 import sinon, { StubbedInstance, stubInterface } from 'ts-sinon';
-import CliServiceProvider from './cli-service-provider';
+import CliServiceProvider, { connectMongoClient } from './cli-service-provider';
 
 chai.use(sinonChai);
 
@@ -34,6 +34,98 @@ describe('CliServiceProvider', () => {
   let serviceProvider: CliServiceProvider;
   let collectionStub: StubbedInstance<Collection>;
 
+  describe('connectMongoClient', () => {
+    it('connects once when no AutoEncryption set', async() => {
+      const uri = 'localhost:27017';
+      const mClientType = stubInterface<typeof MongoClient>();
+      const mClient = stubInterface<MongoClient>();
+      mClientType.connect.onFirstCall().resolves(mClient);
+      const result = await connectMongoClient(uri, {}, mClientType);
+      const calls = mClientType.connect.getCalls();
+      expect(calls.length).to.equal(1);
+      expect(calls[0].args).to.deep.equal([
+        uri, {}
+      ]);
+      expect(result).to.equal(mClient);
+    });
+    it('connects once when bypassAutoEncryption is true', async() => {
+      const uri = 'localhost:27017';
+      const opts = { autoEncryption: { bypassAutoEncryption: true } };
+      const mClientType = stubInterface<typeof MongoClient>();
+      const mClient = stubInterface<MongoClient>();
+      mClientType.connect.onFirstCall().resolves(mClient);
+      const result = await connectMongoClient(uri, opts, mClientType);
+      const calls = mClientType.connect.getCalls();
+      expect(calls.length).to.equal(1);
+      expect(calls[0].args).to.deep.equal([
+        uri, opts
+      ]);
+      expect(result).to.equal(mClient);
+    });
+    it('connects twice when bypassAutoEncryption is false and enterprise via modules', async() => {
+      const uri = 'localhost:27017';
+      const opts = { autoEncryption: { bypassAutoEncryption: false } };
+      const mClientType = stubInterface<typeof MongoClient>();
+      const mClientFirst = stubInterface<MongoClient>();
+      const commandSpy = sinon.spy();
+      mClientFirst.db.returns({ admin: () => ({ command: (...args) => {
+        commandSpy(...args);
+        return { modules: [ 'enterprise' ] };
+      } } as any) } as any);
+      const mClientSecond = stubInterface<MongoClient>();
+      mClientType.connect.onFirstCall().resolves(mClientFirst);
+      mClientType.connect.onSecondCall().resolves(mClientSecond);
+      const result = await connectMongoClient(uri, opts, mClientType);
+      const calls = mClientType.connect.getCalls();
+      expect(calls.length).to.equal(2);
+      expect(calls[0].args).to.deep.equal([
+        uri, {}
+      ]);
+      expect(commandSpy).to.have.been.calledOnceWithExactly({ buildInfo: 1 });
+      expect(result).to.equal(mClientSecond);
+    });
+    it('errors when bypassAutoEncryption is falsy and not enterprise', async() => {
+      const uri = 'localhost:27017';
+      const opts = { autoEncryption: {} };
+      const mClientType = stubInterface<typeof MongoClient>();
+      const mClientFirst = stubInterface<MongoClient>();
+      const commandSpy = sinon.spy();
+      mClientFirst.db.returns({ admin: () => ({ command: (...args) => {
+        commandSpy(...args);
+        return { modules: [] };
+      } } as any) } as any);
+      const mClientSecond = stubInterface<MongoClient>();
+      mClientType.connect.onFirstCall().resolves(mClientFirst);
+      mClientType.connect.onSecondCall().resolves(mClientSecond);
+      try {
+        await connectMongoClient(uri, opts, mClientType);
+      } catch (e) {
+        return expect(e.message.toLowerCase()).to.include('automatic encryption');
+      }
+      expect.fail('Failed to throw expected error');
+    });
+    it('errors when bypassAutoEncryption is falsy, missing modules', async() => {
+      const uri = 'localhost:27017';
+      const opts = { autoEncryption: {} };
+      const mClientType = stubInterface<typeof MongoClient>();
+      const mClientFirst = stubInterface<MongoClient>();
+      const commandSpy = sinon.spy();
+      mClientFirst.db.returns({ admin: () => ({ command: (...args) => {
+        commandSpy(...args);
+        return {};
+      } } as any) } as any);
+      const mClientSecond = stubInterface<MongoClient>();
+      mClientType.connect.onFirstCall().resolves(mClientFirst);
+      mClientType.connect.onSecondCall().resolves(mClientSecond);
+      try {
+        await connectMongoClient(uri, opts, mClientType);
+      } catch (e) {
+        return expect(e.message.toLowerCase()).to.include('automatic encryption');
+      }
+      expect.fail('Failed to throw expected error');
+    });
+  });
+
   describe('#constructor', () => {
     const mongoClient: any = sinon.spy();
     serviceProvider = new CliServiceProvider(mongoClient);

packages/service-provider-server/src/cli-service-provider.ts

@@ -74,7 +74,7 @@ import {
   ConnectionString
 } from '@mongosh/service-provider-core';
 
-import { MongoshCommandFailed, MongoshInternalError } from '@mongosh/errors';
+import { MongoshCommandFailed, MongoshInternalError, MongoshRuntimeError } from '@mongosh/errors';
 
 const bsonlib = {
   Binary,
@@ -127,7 +127,35 @@ const DEFAULT_BASE_OPTIONS = Object.freeze({
 });
 
 /**
- * Encapsulates logic for the service provider for the mongosh CLI.
+ * Connect a MongoClient. If AutoEncryption is requested, first connect without the encryption options and verify that
+ * the connection is to an enterprise cluster. If not, then error, otherwise close the connection and reconnect with the
+ * options the user initially specified. Provide the client class as an additional argument in order to test.
+ * @param uri {String}
+ * @param clientOptions {MongoClientOptions}
+ * @param mClient {MongoClient}
+ */
+export async function connectMongoClient(uri: string, clientOptions: MongoClientOptions, mClient = MongoClient): Promise<MongoClient> {
+  if (clientOptions.autoEncryption !== undefined &&
+    !clientOptions.autoEncryption.bypassAutoEncryption) {
+    // connect first without autoEncryptionOptions
+    const optionsWithoutFLE = { ...clientOptions };
+    delete optionsWithoutFLE.autoEncryption;
+    const client = await mClient.connect(uri, optionsWithoutFLE);
+    const buildInfo = await client.db('admin').admin().command({ buildInfo: 1 });
+    if (
+      !(buildInfo.modules?.includes('enterprise')) &&
+      !(buildInfo.gitVersion?.match(/enterprise/))
+    ) {
+      await client.close();
+      throw new MongoshRuntimeError('Automatic encryption is only available with Atlas and MongoDB Enterprise');
+    }
+    await client.close();
+  }
+  return mClient.connect(uri, clientOptions);
+}
+
+/**
+   * Encapsulates logic for the service provider for the mongosh CLI.
  */
 class CliServiceProvider extends ServiceProviderCore implements ServiceProvider {
   /**
@@ -148,7 +176,7 @@ class CliServiceProvider extends ServiceProviderCore implements ServiceProvider
     const clientOptions = processDriverOptions(driverOptions);
 
     const mongoClient = !cliOptions.nodb ?
-      await MongoClient.connect(
+      await connectMongoClient(
         connectionString.toString(),
         clientOptions
       ) :
@@ -196,7 +224,7 @@ class CliServiceProvider extends ServiceProviderCore implements ServiceProvider
     const connectionString = new ConnectionString(uri);
     const clientOptions = processDriverOptions(options);
 
-    const mongoClient = await MongoClient.connect(
+    const mongoClient = await connectMongoClient(
       connectionString.toString(),
       clientOptions
     );
@@ -1026,7 +1054,7 @@ class CliServiceProvider extends ServiceProviderCore implements ServiceProvider
     });
     if (authDoc.mechanism) clientOptions.authMechanism = authDoc.mechanism as AuthMechanismId;
     if (authDoc.authDb) clientOptions.authSource = authDoc.authDb;
-    const mc = await MongoClient.connect(
+    const mc = await connectMongoClient(
       Object.assign((this.uri as ConnectionString).clone(), {
         username: '', password: ''
       }).toString(),
@@ -1096,7 +1124,7 @@ class CliServiceProvider extends ServiceProviderCore implements ServiceProvider
       ...this.initialOptions,
       ...options
     });
-    const mc = await MongoClient.connect(
+    const mc = await connectMongoClient(
       // TODO This seems to potentially undo a previous db.auth(), MONGOSH-529
       (this.uri as ConnectionString).toString(),
       clientOptions