opusdev/vector-similarity-api

1import * as fs from 'fs/promises';2import { type MongoCryptContext, type MongoCryptKMSRequest } from 'mongodb-client-encryption';3import * as net from 'net';4import * as tls from 'tls';5 6import {7  type BSONSerializeOptions,8  deserialize,9  type Document,10  pluckBSONSerializeOptions,11  serialize12} from '../bson';13import { type ProxyOptions } from '../cmap/connection';14import { CursorTimeoutContext } from '../cursor/abstract_cursor';15import { getSocks, type SocksLib } from '../deps';16import { MongoOperationTimeoutError } from '../error';17import { type MongoClient, type MongoClientOptions } from '../mongo_client';18import { type Abortable } from '../mongo_types';19import { type CollectionInfo } from '../operations/list_collections';20import { Timeout, type TimeoutContext, TimeoutError } from '../timeout';21import {22  addAbortListener,23  BufferPool,24  kDispose,25  MongoDBCollectionNamespace,26  promiseWithResolvers27} from '../utils';28import { autoSelectSocketOptions, type DataKey } from './client_encryption';29import { MongoCryptError } from './errors';30import { type MongocryptdManager } from './mongocryptd_manager';31import { type KMSProviders } from './providers';32 33let socks: SocksLib | null = null;34function loadSocks(): SocksLib {35  if (socks == null) {36    const socksImport = getSocks();37    if ('kModuleError' in socksImport) {38      throw socksImport.kModuleError;39    }40    socks = socksImport;41  }42  return socks;43}44 45// libmongocrypt states46const MONGOCRYPT_CTX_ERROR = 0;47const MONGOCRYPT_CTX_NEED_MONGO_COLLINFO = 1;48const MONGOCRYPT_CTX_NEED_MONGO_MARKINGS = 2;49const MONGOCRYPT_CTX_NEED_MONGO_KEYS = 3;50const MONGOCRYPT_CTX_NEED_KMS_CREDENTIALS = 7;51const MONGOCRYPT_CTX_NEED_KMS = 4;52const MONGOCRYPT_CTX_READY = 5;53const MONGOCRYPT_CTX_DONE = 6;54 55const HTTPS_PORT = 443;56 57const stateToString = new Map([58  [MONGOCRYPT_CTX_ERROR, 'MONGOCRYPT_CTX_ERROR'],59  [MONGOCRYPT_CTX_NEED_MONGO_COLLINFO, 'MONGOCRYPT_CTX_NEED_MONGO_COLLINFO'],60  [MONGOCRYPT_CTX_NEED_MONGO_MARKINGS, 'MONGOCRYPT_CTX_NEED_MONGO_MARKINGS'],61  [MONGOCRYPT_CTX_NEED_MONGO_KEYS, 'MONGOCRYPT_CTX_NEED_MONGO_KEYS'],62  [MONGOCRYPT_CTX_NEED_KMS_CREDENTIALS, 'MONGOCRYPT_CTX_NEED_KMS_CREDENTIALS'],63  [MONGOCRYPT_CTX_NEED_KMS, 'MONGOCRYPT_CTX_NEED_KMS'],64  [MONGOCRYPT_CTX_READY, 'MONGOCRYPT_CTX_READY'],65  [MONGOCRYPT_CTX_DONE, 'MONGOCRYPT_CTX_DONE']66]);67 68const INSECURE_TLS_OPTIONS = [69  'tlsInsecure',70  'tlsAllowInvalidCertificates',71  'tlsAllowInvalidHostnames'72];73 74/**75 * Helper function for logging. Enabled by setting the environment flag MONGODB_CRYPT_DEBUG.76 * @param msg - Anything you want to be logged.77 */78function debug(msg: unknown) {79  if (process.env.MONGODB_CRYPT_DEBUG) {80    // eslint-disable-next-line no-console81    console.error(msg);82  }83}84 85declare module 'mongodb-client-encryption' {86  // the properties added to `MongoCryptContext` here are only used for the `StateMachine`'s87  // execute method and are not part of the C++ bindings.88  interface MongoCryptContext {89    id: number;90    document: Document;91    ns: string;92  }93}94 95/**96 * @public97 *98 * TLS options to use when connecting. The spec specifically calls out which insecure99 * tls options are not allowed:100 *101 *  - tlsAllowInvalidCertificates102 *  - tlsAllowInvalidHostnames103 *  - tlsInsecure104 *105 * These options are not included in the type, and are ignored if provided.106 */107export type ClientEncryptionTlsOptions = Pick<108  MongoClientOptions,109  'tlsCAFile' | 'tlsCertificateKeyFile' | 'tlsCertificateKeyFilePassword' | 'secureContext'110>;111 112/** @public */113export type CSFLEKMSTlsOptions = {114  aws?: ClientEncryptionTlsOptions;115  gcp?: ClientEncryptionTlsOptions;116  kmip?: ClientEncryptionTlsOptions;117  local?: ClientEncryptionTlsOptions;118  azure?: ClientEncryptionTlsOptions;119 120  [key: string]: ClientEncryptionTlsOptions | undefined;121};122 123/**124 * @public125 *126 * Socket options to use for KMS requests.127 */128export type ClientEncryptionSocketOptions = Pick<129  MongoClientOptions,130  'autoSelectFamily' | 'autoSelectFamilyAttemptTimeout'131>;132 133/**134 * This is kind of a hack.  For `rewrapManyDataKey`, we have tests that135 * guarantee that when there are no matching keys, `rewrapManyDataKey` returns136 * nothing.  We also have tests for auto encryption that guarantee for `encrypt`137 * we return an error when there are no matching keys.  This error is generated in138 * subsequent iterations of the state machine.139 * Some apis (`encrypt`) throw if there are no filter matches and others (`rewrapManyDataKey`)140 * do not.  We set the result manually here, and let the state machine continue.  `libmongocrypt`141 * will inform us if we need to error by setting the state to `MONGOCRYPT_CTX_ERROR` but142 * otherwise we'll return `{ v: [] }`.143 */144let EMPTY_V;145 146/**147 * @internal148 *149 * An interface representing an object that can be passed to the `StateMachine.execute` method.150 *151 * Not all properties are required for all operations.152 */153export interface StateMachineExecutable {154  _keyVaultNamespace: string;155  _keyVaultClient: MongoClient;156  askForKMSCredentials: () => Promise<KMSProviders>;157 158  /** only used for auto encryption */159  _metaDataClient?: MongoClient;160  /** only used for auto encryption */161  _mongocryptdClient?: MongoClient;162  /** only used for auto encryption */163  _mongocryptdManager?: MongocryptdManager;164}165 166export type StateMachineOptions = {167  /** socks5 proxy options, if set. */168  proxyOptions: ProxyOptions;169 170  /** TLS options for KMS requests, if set. */171  tlsOptions: CSFLEKMSTlsOptions;172 173  /** Socket specific options we support. */174  socketOptions: ClientEncryptionSocketOptions;175} & Pick<BSONSerializeOptions, 'promoteLongs' | 'promoteValues'>;176 177/**178 * @internal179 * An internal class that executes across a MongoCryptContext until either180 * a finishing state or an error is reached. Do not instantiate directly.181 */182// TODO(DRIVERS-2671): clarify CSOT behavior for FLE APIs183export class StateMachine {184  private options: StateMachineOptions;185  private bsonOptions: BSONSerializeOptions;186 187  constructor(options: StateMachineOptions, bsonOptions = pluckBSONSerializeOptions(options)) {188    this.options = options;189    this.bsonOptions = bsonOptions;190  }191 192  /**193   * Executes the state machine according to the specification194   */195  async execute(196    executor: StateMachineExecutable,197    context: MongoCryptContext,198    options: { timeoutContext?: TimeoutContext } & Abortable199  ): Promise<Uint8Array> {200    const keyVaultNamespace = executor._keyVaultNamespace;201    const keyVaultClient = executor._keyVaultClient;202    const metaDataClient = executor._metaDataClient;203    const mongocryptdClient = executor._mongocryptdClient;204    const mongocryptdManager = executor._mongocryptdManager;205    let result: Uint8Array | null = null;206 207    // Typescript treats getters just like properties: Once you've tested it for equality208    // it cannot change. Which is exactly the opposite of what we use state and status for.209    // Every call to at least `addMongoOperationResponse` and `finalize` can change the state.210    // These wrappers let us write code more naturally and not add compiler exceptions211    // to conditions checks inside the state machine.212    const getStatus = () => context.status;213    const getState = () => context.state;214 215    while (getState() !== MONGOCRYPT_CTX_DONE && getState() !== MONGOCRYPT_CTX_ERROR) {216      options.signal?.throwIfAborted();217      debug(`[context#${context.id}] ${stateToString.get(getState()) || getState()}`);218 219      switch (getState()) {220        case MONGOCRYPT_CTX_NEED_MONGO_COLLINFO: {221          const filter = deserialize(context.nextMongoOperation());222          if (!metaDataClient) {223            throw new MongoCryptError(224              'unreachable state machine state: entered MONGOCRYPT_CTX_NEED_MONGO_COLLINFO but metadata client is undefined'225            );226          }227 228          const collInfoCursor = this.fetchCollectionInfo(229            metaDataClient,230            context.ns,231            filter,232            options233          );234 235          for await (const collInfo of collInfoCursor) {236            context.addMongoOperationResponse(serialize(collInfo));237            if (getState() === MONGOCRYPT_CTX_ERROR) break;238          }239 240          if (getState() === MONGOCRYPT_CTX_ERROR) break;241 242          context.finishMongoOperation();243          break;244        }245 246        case MONGOCRYPT_CTX_NEED_MONGO_MARKINGS: {247          const command = context.nextMongoOperation();248          if (getState() === MONGOCRYPT_CTX_ERROR) break;249 250          if (!mongocryptdClient) {251            throw new MongoCryptError(252              'unreachable state machine state: entered MONGOCRYPT_CTX_NEED_MONGO_MARKINGS but mongocryptdClient is undefined'253            );254          }255 256          // When we are using the shared library, we don't have a mongocryptd manager.257          const markedCommand: Uint8Array = mongocryptdManager258            ? await mongocryptdManager.withRespawn(259                this.markCommand.bind(this, mongocryptdClient, context.ns, command, options)260              )261            : await this.markCommand(mongocryptdClient, context.ns, command, options);262 263          context.addMongoOperationResponse(markedCommand);264          context.finishMongoOperation();265          break;266        }267 268        case MONGOCRYPT_CTX_NEED_MONGO_KEYS: {269          const filter = context.nextMongoOperation();270          const keys = await this.fetchKeys(keyVaultClient, keyVaultNamespace, filter, options);271 272          if (keys.length === 0) {273            // See docs on EMPTY_V274            result = EMPTY_V ??= serialize({ v: [] });275          }276          for (const key of keys) {277            context.addMongoOperationResponse(serialize(key));278          }279 280          context.finishMongoOperation();281 282          break;283        }284 285        case MONGOCRYPT_CTX_NEED_KMS_CREDENTIALS: {286          const kmsProviders = await executor.askForKMSCredentials();287          context.provideKMSProviders(serialize(kmsProviders));288          break;289        }290 291        case MONGOCRYPT_CTX_NEED_KMS: {292          await Promise.all(this.requests(context, options));293          context.finishKMSRequests();294          break;295        }296 297        case MONGOCRYPT_CTX_READY: {298          const finalizedContext = context.finalize();299          if (getState() === MONGOCRYPT_CTX_ERROR) {300            const message = getStatus().message || 'Finalization error';301            throw new MongoCryptError(message);302          }303          result = finalizedContext;304          break;305        }306 307        default:308          throw new MongoCryptError(`Unknown state: ${getState()}`);309      }310    }311 312    if (getState() === MONGOCRYPT_CTX_ERROR || result == null) {313      const message = getStatus().message;314      if (!message) {315        debug(316          `unidentifiable error in MongoCrypt - received an error status from \`libmongocrypt\` but received no error message.`317        );318      }319      throw new MongoCryptError(320        message ??321          'unidentifiable error in MongoCrypt - received an error status from `libmongocrypt` but received no error message.'322      );323    }324 325    return result;326  }327 328  /**329   * Handles the request to the KMS service. Exposed for testing purposes. Do not directly invoke.330   * @param kmsContext - A C++ KMS context returned from the bindings331   * @returns A promise that resolves when the KMS reply has be fully parsed332   */333  async kmsRequest(334    request: MongoCryptKMSRequest,335    options?: { timeoutContext?: TimeoutContext } & Abortable336  ): Promise<void> {337    const parsedUrl = request.endpoint.split(':');338    const port = parsedUrl[1] != null ? Number.parseInt(parsedUrl[1], 10) : HTTPS_PORT;339    const socketOptions: tls.ConnectionOptions & {340      host: string;341      port: number;342      autoSelectFamily?: boolean;343      autoSelectFamilyAttemptTimeout?: number;344    } = {345      host: parsedUrl[0],346      servername: parsedUrl[0],347      port,348      ...autoSelectSocketOptions(this.options.socketOptions || {})349    };350    const message = request.message;351    const buffer = new BufferPool();352 353    let netSocket: net.Socket;354    let socket: tls.TLSSocket;355 356    function destroySockets() {357      for (const sock of [socket, netSocket]) {358        if (sock) {359          sock.destroy();360        }361      }362    }363 364    function onerror(cause: Error) {365      return new MongoCryptError('KMS request failed', { cause });366    }367 368    function onclose() {369      return new MongoCryptError('KMS request closed');370    }371 372    const tlsOptions = this.options.tlsOptions;373    if (tlsOptions) {374      const kmsProvider = request.kmsProvider;375      const providerTlsOptions = tlsOptions[kmsProvider];376      if (providerTlsOptions) {377        const error = this.validateTlsOptions(kmsProvider, providerTlsOptions);378        if (error) {379          throw error;380        }381        try {382          await this.setTlsOptions(providerTlsOptions, socketOptions);383        } catch (err) {384          throw onerror(err);385        }386      }387    }388 389    let abortListener;390 391    try {392      if (this.options.proxyOptions && this.options.proxyOptions.proxyHost) {393        netSocket = new net.Socket();394 395        const {396          promise: willConnect,397          reject: rejectOnNetSocketError,398          resolve: resolveOnNetSocketConnect399        } = promiseWithResolvers<void>();400 401        netSocket402          .once('error', err => rejectOnNetSocketError(onerror(err)))403          .once('close', () => rejectOnNetSocketError(onclose()))404          .once('connect', () => resolveOnNetSocketConnect());405 406        const netSocketOptions = {407          ...socketOptions,408          host: this.options.proxyOptions.proxyHost,409          port: this.options.proxyOptions.proxyPort || 1080410        };411 412        netSocket.connect(netSocketOptions);413 414        await willConnect;415 416        try {417          socks ??= loadSocks();418          socketOptions.socket = (419            await socks.SocksClient.createConnection({420              existing_socket: netSocket,421              command: 'connect',422              destination: { host: socketOptions.host, port: socketOptions.port },423              proxy: {424                // host and port are ignored because we pass existing_socket425                host: 'iLoveJavaScript',426                port: 0,427                type: 5,428                userId: this.options.proxyOptions.proxyUsername,429                password: this.options.proxyOptions.proxyPassword430              }431            })432          ).socket;433        } catch (err) {434          throw onerror(err);435        }436      }437 438      socket = tls.connect(socketOptions, () => {439        socket.write(message);440      });441 442      const {443        promise: willResolveKmsRequest,444        reject: rejectOnTlsSocketError,445        resolve446      } = promiseWithResolvers<void>();447 448      abortListener = addAbortListener(options?.signal, function () {449        destroySockets();450        rejectOnTlsSocketError(this.reason);451      });452 453      socket454        .once('error', err => rejectOnTlsSocketError(onerror(err)))455        .once('close', () => rejectOnTlsSocketError(onclose()))456        .on('data', data => {457          buffer.append(data);458          while (request.bytesNeeded > 0 && buffer.length) {459            const bytesNeeded = Math.min(request.bytesNeeded, buffer.length);460            request.addResponse(buffer.read(bytesNeeded));461          }462 463          if (request.bytesNeeded <= 0) {464            resolve();465          }466        });467      await (options?.timeoutContext?.csotEnabled()468        ? Promise.all([469            willResolveKmsRequest,470            Timeout.expires(options.timeoutContext?.remainingTimeMS)471          ])472        : willResolveKmsRequest);473    } catch (error) {474      if (error instanceof TimeoutError)475        throw new MongoOperationTimeoutError('KMS request timed out');476      throw error;477    } finally {478      // There's no need for any more activity on this socket at this point.479      destroySockets();480      abortListener?.[kDispose]();481    }482  }483 484  *requests(context: MongoCryptContext, options?: { timeoutContext?: TimeoutContext } & Abortable) {485    for (486      let request = context.nextKMSRequest();487      request != null;488      request = context.nextKMSRequest()489    ) {490      yield this.kmsRequest(request, options);491    }492  }493 494  /**495   * Validates the provided TLS options are secure.496   *497   * @param kmsProvider - The KMS provider name.498   * @param tlsOptions - The client TLS options for the provider.499   *500   * @returns An error if any option is invalid.501   */502  validateTlsOptions(503    kmsProvider: string,504    tlsOptions: ClientEncryptionTlsOptions505  ): MongoCryptError | void {506    const tlsOptionNames = Object.keys(tlsOptions);507    for (const option of INSECURE_TLS_OPTIONS) {508      if (tlsOptionNames.includes(option)) {509        return new MongoCryptError(`Insecure TLS options prohibited for ${kmsProvider}: ${option}`);510      }511    }512  }513 514  /**515   * Sets only the valid secure TLS options.516   *517   * @param tlsOptions - The client TLS options for the provider.518   * @param options - The existing connection options.519   */520  async setTlsOptions(521    tlsOptions: ClientEncryptionTlsOptions,522    options: tls.ConnectionOptions523  ): Promise<void> {524    // If a secureContext is provided, ensure it is set.525    if (tlsOptions.secureContext) {526      options.secureContext = tlsOptions.secureContext;527    }528    if (tlsOptions.tlsCertificateKeyFile) {529      const cert = await fs.readFile(tlsOptions.tlsCertificateKeyFile);530      options.cert = options.key = cert;531    }532    if (tlsOptions.tlsCAFile) {533      options.ca = await fs.readFile(tlsOptions.tlsCAFile);534    }535    if (tlsOptions.tlsCertificateKeyFilePassword) {536      options.passphrase = tlsOptions.tlsCertificateKeyFilePassword;537    }538  }539 540  /**541   * Fetches collection info for a provided namespace, when libmongocrypt542   * enters the `MONGOCRYPT_CTX_NEED_MONGO_COLLINFO` state. The result is543   * used to inform libmongocrypt of the schema associated with this544   * namespace. Exposed for testing purposes. Do not directly invoke.545   *546   * @param client - A MongoClient connected to the topology547   * @param ns - The namespace to list collections from548   * @param filter - A filter for the listCollections command549   * @param callback - Invoked with the info of the requested collection, or with an error550   */551  fetchCollectionInfo(552    client: MongoClient,553    ns: string,554    filter: Document,555    options?: { timeoutContext?: TimeoutContext } & Abortable556  ): AsyncIterable<CollectionInfo> {557    const { db } = MongoDBCollectionNamespace.fromString(ns);558 559    const cursor = client.db(db).listCollections(filter, {560      promoteLongs: false,561      promoteValues: false,562      timeoutContext:563        options?.timeoutContext && new CursorTimeoutContext(options?.timeoutContext, Symbol()),564      signal: options?.signal,565      nameOnly: false566    });567 568    return cursor;569  }570 571  /**572   * Calls to the mongocryptd to provide markings for a command.573   * Exposed for testing purposes. Do not directly invoke.574   * @param client - A MongoClient connected to a mongocryptd575   * @param ns - The namespace (database.collection) the command is being executed on576   * @param command - The command to execute.577   * @param callback - Invoked with the serialized and marked bson command, or with an error578   */579  async markCommand(580    client: MongoClient,581    ns: string,582    command: Uint8Array,583    options?: { timeoutContext?: TimeoutContext } & Abortable584  ): Promise<Uint8Array> {585    const { db } = MongoDBCollectionNamespace.fromString(ns);586    const bsonOptions = { promoteLongs: false, promoteValues: false };587    const rawCommand = deserialize(command, bsonOptions);588 589    const commandOptions: {590      timeoutMS?: number;591      signal?: AbortSignal;592    } = {593      timeoutMS: undefined,594      signal: undefined595    };596 597    if (options?.timeoutContext?.csotEnabled()) {598      commandOptions.timeoutMS = options.timeoutContext.remainingTimeMS;599    }600    if (options?.signal) {601      commandOptions.signal = options.signal;602    }603 604    const response = await client.db(db).command(rawCommand, {605      ...bsonOptions,606      ...commandOptions607    });608 609    return serialize(response, this.bsonOptions);610  }611 612  /**613   * Requests keys from the keyVault collection on the topology.614   * Exposed for testing purposes. Do not directly invoke.615   * @param client - A MongoClient connected to the topology616   * @param keyVaultNamespace - The namespace (database.collection) of the keyVault Collection617   * @param filter - The filter for the find query against the keyVault Collection618   * @param callback - Invoked with the found keys, or with an error619   */620  fetchKeys(621    client: MongoClient,622    keyVaultNamespace: string,623    filter: Uint8Array,624    options?: { timeoutContext?: TimeoutContext } & Abortable625  ): Promise<Array<DataKey>> {626    const { db: dbName, collection: collectionName } =627      MongoDBCollectionNamespace.fromString(keyVaultNamespace);628 629    const commandOptions: {630      timeoutContext?: CursorTimeoutContext;631      signal?: AbortSignal;632    } = {633      timeoutContext: undefined,634      signal: undefined635    };636 637    if (options?.timeoutContext != null) {638      commandOptions.timeoutContext = new CursorTimeoutContext(options.timeoutContext, Symbol());639    }640    if (options?.signal != null) {641      commandOptions.signal = options.signal;642    }643 644    return client645      .db(dbName)646      .collection<DataKey>(collectionName, { readConcern: { level: 'majority' } })647      .find(deserialize(filter), commandOptions)648      .toArray();649  }650}651