diff --git a/package-lock.json b/package-lock.json index d5addf2749b..cf6ca208dfb 100644 --- a/package-lock.json +++ b/package-lock.json @@ -11,9 +11,6 @@ "configs/*", "scripts" ], - "dependencies": { - "electron": "^23.3.12" - }, "devDependencies": { "@babel/core": "7.16.0", "@babel/parser": "7.16.0", @@ -40087,6 +40084,7 @@ "@mongodb-js/devtools-connect": "^2.4.0", "@mongodb-js/oidc-plugin": "^0.3.0", "electron": "^23.3.12", + "keytar": "^7.9.0", "node-fetch": "^2.6.7", "react": "^17.0.2", "react-redux": "^8.0.5", @@ -54867,6 +54865,7 @@ "depcheck": "^1.4.1", "electron": "^23.3.12", "eslint": "^7.25.0", + "keytar": "^7.9.0", "mocha": "^10.2.0", "mongodb": "^5.7.0", "mongodb-schema": "^11.2.1", diff --git a/packages/atlas-service/package.json b/packages/atlas-service/package.json index 011fe4d5a59..c3bd3ac6de4 100644 --- a/packages/atlas-service/package.json +++ b/packages/atlas-service/package.json @@ -73,6 +73,7 @@ "@mongodb-js/devtools-connect": "^2.4.0", "@mongodb-js/oidc-plugin": "^0.3.0", "electron": "^23.3.12", + "keytar": "^7.9.0", "node-fetch": "^2.6.7", "react": "^17.0.2", "react-redux": "^8.0.5", diff --git a/packages/atlas-service/src/main.spec.ts b/packages/atlas-service/src/main.spec.ts index 8b45a9e03b0..6834d050ae4 100644 --- a/packages/atlas-service/src/main.spec.ts +++ b/packages/atlas-service/src/main.spec.ts @@ -37,11 +37,14 @@ describe('AtlasServiceMain', function () { AtlasService['attachOidcPluginLoggerEvents'](); const fetch = AtlasService['fetch']; + const ipcMain = AtlasService['ipcMain']; const apiBaseUrl = process.env.COMPASS_ATLAS_SERVICE_BASE_URL; const issuer = process.env.COMPASS_OIDC_ISSUER; const clientId = process.env.COMPASS_CLIENT_ID; beforeEach(function () { + AtlasService['ipcMain'] = { handle: sandbox.stub() }; + process.env.COMPASS_ATLAS_SERVICE_BASE_URL = 'http://example.com'; process.env.COMPASS_OIDC_ISSUER = 'http://example.com'; process.env.COMPASS_CLIENT_ID = '1234abcd'; @@ -53,7 +56,9 @@ describe('AtlasServiceMain', function () { process.env.COMPASS_CLIENT_ID = clientId; AtlasService['fetch'] = fetch; + AtlasService['ipcMain'] = ipcMain; AtlasService['token'] = null; + AtlasService['initPromise'] = null; AtlasService['oidcPluginSyncedFromLoggerState'] = 'initial'; sandbox.resetHistory(); @@ -182,10 +187,11 @@ describe('AtlasServiceMain', function () { (async () => { await wait(20); AtlasService['oidcPluginLogger'].emit( - 'mongodb-oidc-plugin:refresh-succeeded' + 'mongodb-oidc-plugin:refresh-started' ); AtlasService['oidcPluginLogger'].emit( - 'mongodb-oidc-plugin:state-updated' + 'mongodb-oidc-plugin:auth-succeeded', + {} as any ); })(), ]); @@ -346,10 +352,11 @@ describe('AtlasServiceMain', function () { (async () => { await wait(20); AtlasService['oidcPluginLogger'].emit( - 'mongodb-oidc-plugin:refresh-succeeded' + 'mongodb-oidc-plugin:refresh-started' ); AtlasService['oidcPluginLogger'].emit( - 'mongodb-oidc-plugin:state-updated' + 'mongodb-oidc-plugin:auth-succeeded', + {} as any ); })(), ]); @@ -480,40 +487,24 @@ describe('AtlasServiceMain', function () { ); }); - it('should set AtlasService state to error on `mongodb-oidc-plugin:refresh-failed` event', function () { - AtlasService['oidcPluginLogger'].emit( - 'mongodb-oidc-plugin:refresh-failed' as any, - { error: 'Stringified logger error' } - ); - expect(AtlasService).to.have.property( - 'oidcPluginSyncedFromLoggerState', - 'error' - ); - }); - - it('should refresh token in atlas service state on `mongodb-oidc-plugin:refresh-succeeded` event', async function () { - const initialListenerCount = getListenerCount( - mockOidcPlugin.logger as EventEmitter - ); + it('should refresh token in atlas service state on `mongodb-oidc-plugin:refresh-started` event', async function () { // Checking that multiple events while we are refreshing don't cause - // multiple calls to REFRESH_TOKEN_CALLBACK - mockOidcPlugin.logger.emit('mongodb-oidc-plugin:refresh-succeeded'); - mockOidcPlugin.logger.emit('mongodb-oidc-plugin:refresh-succeeded'); - mockOidcPlugin.logger.emit('mongodb-oidc-plugin:refresh-succeeded'); - // Checking that refresh-succeeded doesn't update the service state as we - // are just starting the refresh actually - expect(AtlasService).to.have.property( - 'oidcPluginSyncedFromLoggerState', - 'initial' + // multiple calls to REQUEST_TOKEN_CALLBACK + mockOidcPlugin.logger.emit('mongodb-oidc-plugin:refresh-started'); + mockOidcPlugin.logger.emit('mongodb-oidc-plugin:refresh-started'); + mockOidcPlugin.logger.emit('mongodb-oidc-plugin:refresh-started'); + // Make it look like oidc-plugin successfully updated + mockOidcPlugin.logger.emit( + 'mongodb-oidc-plugin:auth-succeeded', + {} as any ); - mockOidcPlugin.logger.emit('mongodb-oidc-plugin:state-updated'); await once( AtlasService['oidcPluginLogger'], 'atlas-service-token-refreshed' ); expect( mockOidcPlugin.mongoClientOptions.authMechanismProperties - .REFRESH_TOKEN_CALLBACK + .REQUEST_TOKEN_CALLBACK ).to.have.been.calledOnce; expect(AtlasService).to.have.property( 'oidcPluginSyncedFromLoggerState', @@ -521,10 +512,42 @@ describe('AtlasServiceMain', function () { ); expect(AtlasService) .to.have.property('token') - .deep.eq({ accessToken: '4321' }); - // Checking that we cleaned up all listeners - expect(getListenerCount(mockOidcPlugin.logger as EventEmitter)).to.eq( - initialListenerCount + .deep.eq({ accessToken: '1234' }); + }); + }); + + describe('init', function () { + it('should set service to unauthenticated state if requesting token throws', async function () { + mockOidcPlugin.mongoClientOptions.authMechanismProperties.REQUEST_TOKEN_CALLBACK = + sandbox + .stub() + .rejects(new Error('Could not retrieve valid access token')); + const createPlugin = () => mockOidcPlugin; + const initPromise = AtlasService.init(createPlugin); + expect(AtlasService).to.have.property( + 'oidcPluginSyncedFromLoggerState', + 'restoring' + ); + await initPromise; + expect(AtlasService).to.have.property( + 'oidcPluginSyncedFromLoggerState', + 'unauthenticated' + ); + }); + + it('should set service to autenticated state if token was returned', async function () { + mockOidcPlugin.mongoClientOptions.authMechanismProperties.REQUEST_TOKEN_CALLBACK = + sandbox.stub().resolves({ accessToken: 'token' }); + const createPlugin = () => mockOidcPlugin; + const initPromise = AtlasService.init(createPlugin); + expect(AtlasService).to.have.property( + 'oidcPluginSyncedFromLoggerState', + 'restoring' + ); + await initPromise; + expect(AtlasService).to.have.property( + 'oidcPluginSyncedFromLoggerState', + 'authenticated' ); }); }); diff --git a/packages/atlas-service/src/main.ts b/packages/atlas-service/src/main.ts index 213ba15cae9..d16db5701d8 100644 --- a/packages/atlas-service/src/main.ts +++ b/packages/atlas-service/src/main.ts @@ -1,7 +1,12 @@ -import { shell } from 'electron'; +import { ipcMain, shell } from 'electron'; import { URL, URLSearchParams } from 'url'; import { EventEmitter, once } from 'events'; -import type { MongoDBOIDCPluginOptions } from '@mongodb-js/oidc-plugin'; +import keytar from 'keytar'; +import type { + AuthFlowType, + MongoDBOIDCPlugin, + MongoDBOIDCPluginOptions, +} from '@mongodb-js/oidc-plugin'; import { createMongoDBOIDCPlugin } from '@mongodb-js/oidc-plugin'; import { oidcServerRequestHandler } from '@mongodb-js/devtools-connect'; // TODO(https://github.com/node-fetch/node-fetch/issues/1652): Remove this when @@ -14,6 +19,7 @@ import type { Document } from 'mongodb'; import type { AIQuery, IntrospectInfo, Token, UserInfo } from './util'; import { broadcast, + getStoragePaths, ipcExpose, throwIfAborted, } from '@mongodb-js/compass-utils'; @@ -62,19 +68,110 @@ export async function throwIfNotOk( throw err; } -const MAX_REQUEST_SIZE = 10000; +const AI_MAX_REQUEST_SIZE = 10000; -const MIN_SAMPLE_DOCUMENTS = 1; +const AI_MIN_SAMPLE_DOCUMENTS = 1; type MongoDBOIDCPluginLogger = Required['logger']; +type AtlasServiceAuthState = + // Instance was just created + | 'initial' + // Trying to restore state if it was persisted before + | 'restoring' + // Successfully got token from oidc-plugin + | 'authenticated' + // Token expired (and being refreshed at the moment) + | 'expired' + // Encountered an error while requesting token (either on sign in or refresh) + | 'unauthenticated'; + +type SecretStore = { + getItem(key: string): Promise; + setItem(key: string, value: string): Promise; +}; + +const SECRET_STORE_KEY = 'AtlasLoginOIDCPluginState'; + +const defaultSecretStore: SecretStore = { + async getItem(key: string) { + try { + const { appName } = getStoragePaths() ?? {}; + if ( + process.env.COMPASS_E2E_DISABLE_KEYCHAIN_USAGE === 'true' || + !appName + ) { + throw new Error('Unsupported environment'); + } + return (await keytar.getPassword(appName, key)) ?? undefined; + } catch { + return undefined; + } + }, + async setItem(key: string, value: string) { + try { + const { appName } = getStoragePaths() ?? {}; + if ( + process.env.COMPASS_E2E_DISABLE_KEYCHAIN_USAGE === 'true' || + !appName + ) { + throw new Error('Unsupported environment'); + } + return await keytar.setPassword(appName, key, value); + } catch { + return; + } + }, +}; + export class AtlasService { private constructor() { // singleton } - private static calledOnce = false; - + private static initPromise: Promise | null = null; + + /** + * oidc-plugin logger event flow for plugin creation and token request / refresh: + * + * createPlugin + * ↓ + * serialized state provided? → no → no event + * ↓ + * yes → failed to deserialize? → no → `plugin:state-updated` + * ↓ + * yes → `plugin:deserialization-failed` + * + * + * + * requestToken + * ↓ + * token expired? + * │ + * no ←─┴─→ yes ─→ trying refresh + * │ ↓ + * │ `plugin:refresh-started` + * │ ↓ + * │ got token from issuer? → no ─┐ + * │ ↓ │ + * │ yes │ + * │ ↓ │ + * │ `plugin:refresh-succeeded` │ + * │ ↓ │ + * │ start state update │ + * │ ↓ │ ┌────────────────────────────────────────────┐ + * │ state update failed → yes ───┴→ `plugin:refresh-failed` │ `plugin:auth-attempt-started` │ + * │ ↓ ↓ │ ↓ │ + * │ no for flow in getAllowedFlows → │ is attempt successfull? → no │ + * ↓ ↓ │ ↓ ↓ │ + * `plugin:skip-auth-attempt` ← `plugin:state-updated` │ yes `plugin:auth-attempt-failed` │ + * ↓ │ ↓ │ + * `plugin:auth-succeeded` ←─────── yes ←──────── do we have a new token set in state? ← │ `plugin:state-updated` │ + * ↓ │ ↓ │ + * no │ `plugin:auth-attempt-succeeded` │ + * ↓ └────────────────────────────────────────────┘ + * `plugin:auth-failed` + */ private static oidcPluginLogger: MongoDBOIDCPluginLogger & { on(evt: 'atlas-service-token-refreshed', fn: () => void): void; on(evt: 'atlas-service-token-refresh-failed', fn: () => void): void; @@ -82,33 +179,10 @@ export class AtlasService { emit(evt: 'atlas-service-token-refresh-failed'): void; } = new EventEmitter(); - private static oidcPluginSyncedFromLoggerState: - | 'initial' - | 'authenticated' - | 'expired' - | 'error' = 'initial'; - - private static plugin = createMongoDBOIDCPlugin({ - redirectServerRequestHandler(data) { - if (data.result === 'redirecting') { - const { res, status, location } = data; - res.statusCode = status; - const redirectUrl = new URL( - 'https://account.mongodb.com/account/login' - ); - redirectUrl.searchParams.set('fromURI', location); - res.setHeader('Location', redirectUrl.toString()); - res.end(); - return; - } + private static oidcPluginSyncedFromLoggerState: AtlasServiceAuthState = + 'initial'; - redirectRequestHandler(data); - }, - async openBrowser({ url }) { - await shell.openExternal(url); - }, - logger: this.oidcPluginLogger, - }); + private static plugin: MongoDBOIDCPlugin | null = null; private static token: Token | null = null; @@ -116,48 +190,139 @@ export class AtlasService { private static fetch: typeof fetch = fetch; + private static secretStore: SecretStore = defaultSecretStore; + + private static ipcMain: Pick = ipcMain; + private static refreshing = false; private static get clientId() { - if (!process.env.COMPASS_CLIENT_ID) { + const clientId = + process.env.COMPASS_CLIENT_ID_OVERRIDE || process.env.COMPASS_CLIENT_ID; + if (!clientId) { throw new Error('COMPASS_CLIENT_ID is required'); } - return process.env.COMPASS_CLIENT_ID; + return clientId; } private static get issuer() { - if (!process.env.COMPASS_OIDC_ISSUER) { + const issuer = + process.env.COMPASS_OIDC_ISSUER_OVERRIDE || + process.env.COMPASS_OIDC_ISSUER; + if (!issuer) { throw new Error('COMPASS_OIDC_ISSUER is required'); } - return process.env.COMPASS_OIDC_ISSUER; + return issuer; } private static get apiBaseUrl() { - if (!process.env.COMPASS_ATLAS_SERVICE_BASE_URL) { + const apiBaseUrl = + process.env.COMPASS_ATLAS_SERVICE_BASE_URL_OVERRIDE || + process.env.COMPASS_ATLAS_SERVICE_BASE_URL; + if (!apiBaseUrl) { throw new Error( 'No AI Query endpoint to fetch. Please set the environment variable `COMPASS_ATLAS_SERVICE_BASE_URL`' ); } - return process.env.COMPASS_ATLAS_SERVICE_BASE_URL; + return apiBaseUrl; } - static init() { - if (this.calledOnce) { - return; + // We use `allowedFlows` plugin option to control whether or not plugin is + // allowed to start sign in flow or just try refresh and fail if refresh is + // not possible. + // - If we are authenticated, token expired, or we are restoring state, we + // only allow token refresh + // - Any other state: initial or error, will allow sign in using auth-code + // flow in addition to token refresh + private static getAllowedAuthFlows(): AuthFlowType[] { + return ['restoring', 'authenticated', 'expired'].includes( + this.oidcPluginSyncedFromLoggerState + ) + ? [] + : ['auth-code']; + } + + static init( + _createMongoDBOIDCPlugin = createMongoDBOIDCPlugin + ): Promise { + return (this.initPromise ??= (async () => { + ipcExpose( + 'AtlasService', + this, + [ + 'getUserInfo', + 'introspect', + 'isAuthenticated', + 'signIn', + 'getQueryFromUserInput', + ], + this.ipcMain + ); + this.attachOidcPluginLoggerEvents(); + log.info( + mongoLogId(1_001_000_210), + 'AtlasService', + 'Atlas service initialized' + ); + this.oidcPluginSyncedFromLoggerState = 'restoring'; + const serializedState = await this.secretStore.getItem(SECRET_STORE_KEY); + this.plugin = _createMongoDBOIDCPlugin({ + redirectServerRequestHandler(data) { + if (data.result === 'redirecting') { + const { res, status, location } = data; + res.statusCode = status; + const redirectUrl = new URL( + 'https://account.mongodb.com/account/login' + ); + redirectUrl.searchParams.set('fromURI', location); + res.setHeader('Location', redirectUrl.toString()); + res.end(); + return; + } + + redirectRequestHandler(data); + }, + async openBrowser({ url }) { + await shell.openExternal(url); + }, + allowedFlows: this.getAllowedAuthFlows.bind(this), + logger: this.oidcPluginLogger, + serializedState, + }); + // Whether or not we got the state, try refreshing the token. If there was + // no serialized state returned, this will just put the service in + // `unauthenticated` state quickly. If there was some state, we need to + // refresh the token and get the value back from the oidc-plugin to make + // sure the state between atlas-service and oidc-plugin is in sync + await this.refreshToken({ + // In case where we refresh token after plugin state deserialisation + // happened we don't expect plugin refresh events to be emitted, and + // so we skip waiting for them. Everything else in this flow is + // exactly as if we were handling refresh-success event from + // oidc-plugin + waitForOIDCPluginStateUpdateEvents: false, + }); + })()); + } + + private static async requestOAuthToken({ + signal, + }: { signal?: AbortSignal } = {}) { + if (!this.plugin) { + throw new Error( + 'Trying to use the oidc-plugin before service is initialised' + ); } - this.calledOnce = true; - ipcExpose('AtlasService', this, [ - 'getUserInfo', - 'introspect', - 'isAuthenticated', - 'signIn', - 'getQueryFromUserInput', - ]); - this.attachOidcPluginLoggerEvents(); - log.info( - mongoLogId(1_001_000_210), - 'AtlasService', - 'Atlas service initialized' + + return this.plugin.mongoClientOptions.authMechanismProperties.REQUEST_TOKEN_CALLBACK( + { clientId: this.clientId, issuer: this.issuer }, + { + // Required driver specific stuff + version: 0, + // While called timeoutContext, this is actually an abort signal that + // plugin will listen to, not a timeout + timeoutContext: signal, + } ); } @@ -173,33 +338,6 @@ export class AtlasService { 'Token refresh started by oidc-plugin' ); this.oidcPluginSyncedFromLoggerState = 'expired'; - }); - this.oidcPluginLogger.on( - 'mongodb-oidc-plugin:refresh-failed', - ({ error }) => { - log.error( - mongoLogId(1_001_000_212), - 'AtlasService', - 'Oidc-plugin failed to refresh token', - { error } - ); - this.token = null; - this.oidcPluginSyncedFromLoggerState = 'error'; - this.oidcPluginLogger.emit('atlas-service-token-refresh-failed'); - } - ); - // NB: oidc-plugin only has a logger interface to listen to state updates, - // it also doesn't expose renewed tokens or any other state in any usable - // way from those events or otherwise, so to get the renewed token we listen - // to the refresh-succeeded event and then kick off the "refresh" - // programmatically to be able to get the actual tokens back and sync them - // to the service state - this.oidcPluginLogger.on('mongodb-oidc-plugin:refresh-succeeded', () => { - log.info( - mongoLogId(1_001_000_213), - 'AtlasService', - 'Oidc-plugin refreshed token successfully' - ); void this.refreshToken(); }); this.oidcPluginLogger.on('atlas-service-token-refresh-failed', () => { @@ -210,47 +348,54 @@ export class AtlasService { }); } - private static async refreshToken() { - // In case our call to REFRESH_TOKEN_CALLBACK somehow started another - // token refresh instead of just returning token from the plugin state, we - // short circuit if plugin logged a refresh-succeeded event and this - // listener got triggered + private static async refreshToken({ + waitForOIDCPluginStateUpdateEvents = true, + } = {}) { + // In case our call to requestToken started another token refresh instead of + // just returning token from the plugin state, we short circuit if plugin + // logged a refresh-succeeded event and this listener got triggered if (this.refreshing) { return; } this.refreshing = true; + const onRefreshFailed = () => { + this.token = null; + this.oidcPluginSyncedFromLoggerState = 'unauthenticated'; + this.oidcPluginLogger.emit('atlas-service-token-refresh-failed'); + }; log.info( mongoLogId(1_001_000_214), 'AtlasService', 'Start atlas service token refresh' ); try { - // We expect only one promise below to resolve, to clean up listeners that - // never fired we are setting up an abort controller - const listenerController = new AbortController(); - try { - await Promise.race([ - // When oidc-plugin logged that token was refreshed, the token is not - // actually refreshed yet in the plugin state and so calling `REFRESH_TOKEN_CALLBACK` - // causes weird behavior that actually opens the browser again, to work - // around that we wait for the state update event in addition. We can't - // guarantee that this event will be emitted for our particular state as - // this is not something oidc-plugin exposes, but we can ignore this for - // now as only one auth state is created in this instance of oidc-plugin - once(this.oidcPluginLogger, 'mongodb-oidc-plugin:state-updated', { - signal: listenerController.signal, - }), - // At the same time refresh can still fail at this stage, so to avoid - // refresh being stuck, we also wait for refresh-failed event and throw - // if it happens to avoid calling `REFRESH_TOKEN_CALLBACK` - once(this.oidcPluginLogger, 'mongodb-oidc-plugin:refresh-failed', { - signal: listenerController.signal, - }).then(() => { - throw new Error('Refresh failed'); - }), - ]); - } finally { - listenerController.abort(); + if (waitForOIDCPluginStateUpdateEvents) { + // We expect only one promise below to resolve, to clean up listeners + // that never fired we are setting up an abort controller + const listenerController = new AbortController(); + try { + // When oidc-plugin starts internal token refresh, we will wait + // first for the auth flow to either succeed to fail and then + // proceed by requesting token once again from the plugin to sync + // token back to the atlas service state + await Promise.race([ + // We are not using `refresh-succeeded` / `refresh-failed` events + // here because those happen BEFORE `allowedFlows` method is called + // (see oidc-plugin flow diagram) meaning that we need to wait until + // the whole request token flow went through to make sure that sign + // in flow is not allowed while we are refreshing the token + once(this.oidcPluginLogger, 'mongodb-oidc-plugin:auth-succeeded', { + signal: listenerController.signal, + }), + once(this.oidcPluginLogger, 'mongodb-oidc-plugin:auth-failed', { + signal: listenerController.signal, + }).then(() => { + throw new Error('Token refresh failed'); + }), + ]); + } finally { + listenerController.abort(); + } } try { log.info( @@ -258,19 +403,10 @@ export class AtlasService { 'AtlasService', 'Request token refresh from oidc-plugin' ); - const token = - await this.plugin.mongoClientOptions.authMechanismProperties - // WARN: in the oidc plugin refresh callback is actually the same - // method as sign in, so calling it here means that potentially - // this can start an actual sign in flow for the user instead of - // just trying to refresh the token - .REFRESH_TOKEN_CALLBACK( - { clientId: this.clientId, issuer: this.issuer }, - { - // Required driver specific stuff - version: 0, - } - ); + // Request and refresh token are the same methods in oidc-plugin, what + // actually controls whether or not sign in flow is allowed is + // `allowedFlows` property passed to the plugin + const token = await this.requestOAuthToken(); log.info( mongoLogId(1_001_000_216), 'AtlasService', @@ -286,14 +422,16 @@ export class AtlasService { 'Oidc-plugin failed to return refreshed token', { error: (err as Error).stack } ); - // REFRESH_TOKEN_CALLBACK call failed for some reason - this.token = null; - this.oidcPluginSyncedFromLoggerState = 'error'; - this.oidcPluginLogger.emit('atlas-service-token-refresh-failed'); + onRefreshFailed(); } - } catch { - // encountered 'mongodb-oidc-plugin:refresh-failed' event, do nothing, we - // already have a listener for this event + } catch (err) { + log.error( + mongoLogId(1_001_000_222), + 'AtlasService', + 'Oidc-plugin failed to refresh token', + { error: (err as Error).stack } + ); + onRefreshFailed(); } finally { this.refreshing = false; } @@ -305,18 +443,23 @@ export class AtlasService { if (signal?.aborted) { return; } - // In cases where we ended up in expired state, we know that oidc-plugin - // is trying to refresh the token automatically, we can wait for this process - // to finish before proceeding with a request - if (this.oidcPluginSyncedFromLoggerState === 'expired') { + // There are two cases where it makes sense to wait for token refresh before + // proceeding with running atlas service commands: + // - In case where we ended up in `expired` state, we know that oidc-plugin + // is trying to refresh the token automatically at the moment + // - In case of `restoring` saved auth state, we will try to refresh the + // token right after restore happened + if ( + ['expired', 'restoring'].includes(this.oidcPluginSyncedFromLoggerState) + ) { // We expect only one promise below to resolve, to clean up listeners that // never fired we are setting up an abort controller const listenerController = new AbortController(); try { await Promise.race([ - // We are using our own events here and not oidc plugin ones because - // after plugin logged that token was refreshed, we still need to run - // REFRESH_TOKEN_CALLBACK to get the actual token value in the state + // We are using our own events here and not oidc-plugin ones because + // after plugin logged that token was refreshed, we still need to do + // some additional work to get token from the oidc-plugin once(this.oidcPluginLogger, 'atlas-service-token-refreshed', { signal: listenerController.signal, }), @@ -361,18 +504,12 @@ export class AtlasService { this.signInPromise = (async () => { log.info(mongoLogId(1_001_000_218), 'AtlasService', 'Starting sign in'); + + // We might be refreshing token or restoring state at the moment + await this.maybeWaitForToken({ signal }); + try { - this.token = - await this.plugin.mongoClientOptions.authMechanismProperties.REQUEST_TOKEN_CALLBACK( - { clientId: this.clientId, issuer: this.issuer }, - { - // Required driver specific stuff - version: 0, - // This seems to be just an abort signal? We probably can make it - // explicit when adding a proper interface for this - timeoutContext: signal, - } - ); + this.token = await this.requestOAuthToken({ signal }); this.oidcPluginSyncedFromLoggerState = 'authenticated'; log.info( mongoLogId(1_001_000_219), @@ -387,7 +524,7 @@ export class AtlasService { 'Failed to sign in', { error: (err as Error).stack } ); - this.oidcPluginSyncedFromLoggerState = 'error'; + this.oidcPluginSyncedFromLoggerState = 'unauthenticated'; throw err; } })(); @@ -401,7 +538,9 @@ export class AtlasService { signal, }: { signal?: AbortSignal } = {}): Promise { throwIfAborted(signal); + await this.maybeWaitForToken({ signal }); + const res = await this.fetch(`${this.issuer}/v1/userinfo`, { headers: { Authorization: `Bearer ${this.token?.accessToken ?? ''}`, @@ -464,7 +603,7 @@ export class AtlasService { schema, sampleDocuments, }); - if (msgBody.length > MAX_REQUEST_SIZE) { + if (msgBody.length > AI_MAX_REQUEST_SIZE) { // When the message body is over the max size, we try // to see if with fewer sample documents we can still perform the request. // If that fails we throw an error indicating this collection's @@ -474,10 +613,10 @@ export class AtlasService { collectionName, databaseName, schema, - sampleDocuments: sampleDocuments?.slice(0, MIN_SAMPLE_DOCUMENTS), + sampleDocuments: sampleDocuments?.slice(0, AI_MIN_SAMPLE_DOCUMENTS), }); // Why this is not happening on the backend? - if (msgBody.length > MAX_REQUEST_SIZE) { + if (msgBody.length > AI_MAX_REQUEST_SIZE) { throw new Error( 'Error: too large of a request to send to the ai. Please use a smaller prompt or collection with smaller documents.' ); @@ -500,4 +639,24 @@ export class AtlasService { return res.json() as Promise; } + + static async onExit() { + // Haven't even created the oidc-plugin yet, nothing to store + if (!this.plugin) { + return; + } + try { + await this.secretStore.setItem( + SECRET_STORE_KEY, + await this.plugin.serialize() + ); + } catch (err) { + log.warn( + mongoLogId(1_001_000_221), + 'AtlasService', + 'Failed to save auth state', + { erroe: (err as Error).stack } + ); + } + } } diff --git a/packages/compass/src/main/application.ts b/packages/compass/src/main/application.ts index 171cd55b574..f55c9278caf 100644 --- a/packages/compass/src/main/application.ts +++ b/packages/compass/src/main/application.ts @@ -90,8 +90,7 @@ class CompassApplication { return; } - AtlasService.init(); - + void this.setupAtlasService(); this.setupAutoUpdate(); await setupCSFLELibrary(); setupTheme(); @@ -109,6 +108,13 @@ class CompassApplication { return (this.initPromise ??= this._init(mode, globalPreferences)); } + private static async setupAtlasService() { + await AtlasService.init(); + this.addExitHandler(() => { + return AtlasService.onExit(); + }); + } + private static setupJavaScriptArguments(): void { // For Linux users with drivers that are avoided by Chromium we disable the // GPU check to attempt to bypass the disabled WebGL settings.