From dd9f882043fe89a6829f22255d96315b27759195 Mon Sep 17 00:00:00 2001 From: Sergey Petushkov Date: Thu, 10 Aug 2023 17:57:10 +0200 Subject: [PATCH 1/6] feat(atlas-service): securely persist and restore atlas login state through keychain --- packages/atlas-service/src/main.spec.ts | 6 +- packages/atlas-service/src/main.ts | 341 +++++++++++++++-------- packages/compass/src/main/application.ts | 10 +- 3 files changed, 242 insertions(+), 115 deletions(-) diff --git a/packages/atlas-service/src/main.spec.ts b/packages/atlas-service/src/main.spec.ts index 8b45a9e03b0..2476ceabbbb 100644 --- a/packages/atlas-service/src/main.spec.ts +++ b/packages/atlas-service/src/main.spec.ts @@ -496,7 +496,7 @@ describe('AtlasServiceMain', function () { mockOidcPlugin.logger as EventEmitter ); // Checking that multiple events while we are refreshing don't cause - // multiple calls to REFRESH_TOKEN_CALLBACK + // multiple calls to REQUEST_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'); @@ -513,7 +513,7 @@ describe('AtlasServiceMain', function () { ); expect( mockOidcPlugin.mongoClientOptions.authMechanismProperties - .REFRESH_TOKEN_CALLBACK + .REQUEST_TOKEN_CALLBACK ).to.have.been.calledOnce; expect(AtlasService).to.have.property( 'oidcPluginSyncedFromLoggerState', @@ -521,7 +521,7 @@ describe('AtlasServiceMain', function () { ); expect(AtlasService) .to.have.property('token') - .deep.eq({ accessToken: '4321' }); + .deep.eq({ accessToken: '1234' }); // Checking that we cleaned up all listeners expect(getListenerCount(mockOidcPlugin.logger as EventEmitter)).to.eq( initialListenerCount diff --git a/packages/atlas-service/src/main.ts b/packages/atlas-service/src/main.ts index 213ba15cae9..4e9ab4a1b96 100644 --- a/packages/atlas-service/src/main.ts +++ b/packages/atlas-service/src/main.ts @@ -1,7 +1,12 @@ import { 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,18 +68,68 @@ 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) + | 'error'; + +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 Promise.resolve(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 Promise.resolve(); + } + }, +}; + export class AtlasService { private constructor() { // singleton } - private static calledOnce = false; + private static initPromise: Promise | null = null; private static oidcPluginLogger: MongoDBOIDCPluginLogger & { on(evt: 'atlas-service-token-refreshed', fn: () => void): void; @@ -82,33 +138,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; private static token: Token | null = null; @@ -116,24 +149,37 @@ export class AtlasService { private static fetch: typeof fetch = fetch; + private static secretStore: SecretStore = defaultSecretStore; + 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; } 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; } 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`' ); @@ -141,23 +187,94 @@ export class AtlasService { return process.env.COMPASS_ATLAS_SERVICE_BASE_URL; } - 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(): Promise { + if (this.initPromise) { + return this.initPromise; } - this.calledOnce = true; - ipcExpose('AtlasService', this, [ - 'getUserInfo', - 'introspect', - 'isAuthenticated', - 'signIn', - 'getQueryFromUserInput', - ]); - this.attachOidcPluginLoggerEvents(); - log.info( - mongoLogId(1_001_000_210), - 'AtlasService', - 'Atlas service initialized' + this.initPromise = (async () => { + ipcExpose('AtlasService', this, [ + 'getUserInfo', + 'introspect', + 'isAuthenticated', + 'signIn', + 'getQueryFromUserInput', + ]); + 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, + }); + // If we got serialisedState from store, we will try to refresh the token + // right away to sync token state from the oidc-plugin to AtlasService + if (serializedState) { + 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, + }); + } else { + // If no serialised state was found, set service back to initial state + // so that the sign in flow is allowed again + this.oidcPluginSyncedFromLoggerState = 'initial'; + } + })(); + return this.initPromise; + } + + private static requestOAuthToken({ signal }: { signal?: AbortSignal } = {}) { + 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, + } ); } @@ -210,11 +327,12 @@ 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; } @@ -225,32 +343,35 @@ export class 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 { + 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(); + } } try { log.info( @@ -258,19 +379,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', @@ -305,10 +417,12 @@ 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 + // 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') { + 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(); @@ -361,18 +475,9 @@ export class AtlasService { this.signInPromise = (async () => { log.info(mongoLogId(1_001_000_218), 'AtlasService', 'Starting sign in'); + 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), @@ -464,7 +569,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 +579,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 +605,20 @@ export class AtlasService { return res.json() as Promise; } + + static async onExit() { + 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 504cda524a6..cc51488c75d 100644 --- a/packages/compass/src/main/application.ts +++ b/packages/compass/src/main/application.ts @@ -92,8 +92,7 @@ class CompassApplication { return; } - AtlasService.init(); - + void this.setupAtlasService(); this.setupAutoUpdate(); await setupCSFLELibrary(); setupTheme(); @@ -111,6 +110,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. From d780136cad2c0a57448a06c974a5f879beabb94b Mon Sep 17 00:00:00 2001 From: Sergey Petushkov Date: Fri, 11 Aug 2023 19:02:39 +0200 Subject: [PATCH 2/6] chore(atlas-service): refactor logger listening logic; document oidc-plugin flow --- packages/atlas-service/src/main.spec.ts | 62 ++++--- packages/atlas-service/src/main.ts | 216 ++++++++++++++---------- 2 files changed, 153 insertions(+), 125 deletions(-) diff --git a/packages/atlas-service/src/main.spec.ts b/packages/atlas-service/src/main.spec.ts index 2476ceabbbb..d12669af20b 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,6 +56,7 @@ describe('AtlasServiceMain', function () { process.env.COMPASS_CLIENT_ID = clientId; AtlasService['fetch'] = fetch; + AtlasService['ipcMain'] = ipcMain; AtlasService['token'] = null; AtlasService['oidcPluginSyncedFromLoggerState'] = 'initial'; @@ -182,10 +186,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 +351,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,33 +486,17 @@ 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 REQUEST_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' + 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' @@ -522,10 +512,18 @@ describe('AtlasServiceMain', function () { expect(AtlasService) .to.have.property('token') .deep.eq({ accessToken: '1234' }); - // Checking that we cleaned up all listeners - expect(getListenerCount(mockOidcPlugin.logger as EventEmitter)).to.eq( - initialListenerCount - ); + }); + }); + + describe('init', function () { + describe('with no stored auth state', function () { + it('should set service to unauthenticated state', function () {}); + }); + + describe('with stored auth state', function () { + it('should set service to unauthenticated state if token expired', function () {}); + + it('should set service to autenticated state if token can be refreshed', function () {}); }); }); }); diff --git a/packages/atlas-service/src/main.ts b/packages/atlas-service/src/main.ts index 4e9ab4a1b96..7ffdd5bad5e 100644 --- a/packages/atlas-service/src/main.ts +++ b/packages/atlas-service/src/main.ts @@ -1,4 +1,4 @@ -import { shell } from 'electron'; +import { ipcMain, shell } from 'electron'; import { URL, URLSearchParams } from 'url'; import { EventEmitter, once } from 'events'; import keytar from 'keytar'; @@ -84,7 +84,7 @@ type AtlasServiceAuthState = // Token expired (and being refreshed at the moment) | 'expired' // Encountered an error while requesting token (either on sign in or refresh) - | 'error'; + | 'unauthenticated'; type SecretStore = { getItem(key: string): Promise; @@ -131,6 +131,47 @@ export class AtlasService { 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 → no event + * ↓ + * 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:auth-attempt-succeeded` │ + * no └────────────────────────────────────────────┘ + * ↓ + * `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; @@ -151,40 +192,39 @@ export class AtlasService { private static secretStore: SecretStore = defaultSecretStore; + private static ipcMain: Pick = ipcMain; + private static refreshing = false; private static get clientId() { 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() { 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() { 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; } // We use `allowedFlows` plugin option to control whether or not plugin is @@ -203,17 +243,19 @@ export class AtlasService { } static init(): Promise { - if (this.initPromise) { - return this.initPromise; - } - this.initPromise = (async () => { - ipcExpose('AtlasService', this, [ - 'getUserInfo', - 'introspect', - 'isAuthenticated', - 'signIn', - 'getQueryFromUserInput', - ]); + return (this.initPromise ??= (async () => { + ipcExpose( + 'AtlasService', + this, + [ + 'getUserInfo', + 'introspect', + 'isAuthenticated', + 'signIn', + 'getQueryFromUserInput', + ], + this.ipcMain + ); this.attachOidcPluginLoggerEvents(); log.info( mongoLogId(1_001_000_210), @@ -245,27 +287,25 @@ export class AtlasService { logger: this.oidcPluginLogger, serializedState, }); - // If we got serialisedState from store, we will try to refresh the token - // right away to sync token state from the oidc-plugin to AtlasService - if (serializedState) { - 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, - }); - } else { - // If no serialised state was found, set service back to initial state - // so that the sign in flow is allowed again - this.oidcPluginSyncedFromLoggerState = 'initial'; - } - })(); - return this.initPromise; + // 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 requestOAuthToken({ signal }: { signal?: AbortSignal } = {}) { + private static async requestOAuthToken({ + signal, + }: { signal?: AbortSignal } = {}) { return this.plugin.mongoClientOptions.authMechanismProperties.REQUEST_TOKEN_CALLBACK( { clientId: this.clientId, issuer: this.issuer }, { @@ -290,33 +330,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', () => { @@ -337,6 +350,11 @@ export class AtlasService { 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', @@ -348,25 +366,23 @@ export class AtlasService { // 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([ - // 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', { + // NB: While `auth-{succeeded,failed}` events can also fire when one + // of the actual sign in flows is activated (see logger event + // diagram), we make sure that only token refresh is allowed during + // this process by providing an empty array of allowedFlows to the + // plugin through the `getAllowedFlows` method + once(this.oidcPluginLogger, 'mongodb-oidc-plugin:auth-succeeded', { 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', { + once(this.oidcPluginLogger, 'mongodb-oidc-plugin:auth-failed', { signal: listenerController.signal, }).then(() => { - throw new Error('Refresh failed'); + throw new Error('Token refresh failed'); }), ]); } finally { @@ -398,14 +414,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; } @@ -417,9 +435,12 @@ 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 + // 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) ) { @@ -428,9 +449,9 @@ export class AtlasService { 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, }), @@ -475,7 +496,10 @@ 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.requestOAuthToken({ signal }); this.oidcPluginSyncedFromLoggerState = 'authenticated'; @@ -492,7 +516,7 @@ export class AtlasService { 'Failed to sign in', { error: (err as Error).stack } ); - this.oidcPluginSyncedFromLoggerState = 'error'; + this.oidcPluginSyncedFromLoggerState = 'unauthenticated'; throw err; } })(); @@ -506,7 +530,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 ?? ''}`, @@ -607,6 +633,10 @@ export class AtlasService { } static async onExit() { + // Haven't even c the oidc-plugin yet, nothing to store + if (!this.plugin) { + return; + } try { await this.secretStore.setItem( SECRET_STORE_KEY, From 2e9f3b9c8179817a1277367655ef49467f2d3e0a Mon Sep 17 00:00:00 2001 From: Sergey Petushkov Date: Fri, 11 Aug 2023 19:10:12 +0200 Subject: [PATCH 3/6] chore(atlas-service): add missing keytar dependency --- package-lock.json | 5 ++--- packages/atlas-service/package.json | 1 + 2 files changed, 3 insertions(+), 3 deletions(-) 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", From d8ef4791b28dbca3d5f2d34f4960465cb453fa8c Mon Sep 17 00:00:00 2001 From: Sergey Petushkov Date: Mon, 14 Aug 2023 10:53:19 +0200 Subject: [PATCH 4/6] chore(atlas-service): add tests; update diagram --- packages/atlas-service/src/main.spec.ts | 37 +++++++++++++++---- packages/atlas-service/src/main.ts | 48 ++++++++++++++----------- 2 files changed, 59 insertions(+), 26 deletions(-) diff --git a/packages/atlas-service/src/main.spec.ts b/packages/atlas-service/src/main.spec.ts index d12669af20b..6834d050ae4 100644 --- a/packages/atlas-service/src/main.spec.ts +++ b/packages/atlas-service/src/main.spec.ts @@ -58,6 +58,7 @@ describe('AtlasServiceMain', function () { AtlasService['fetch'] = fetch; AtlasService['ipcMain'] = ipcMain; AtlasService['token'] = null; + AtlasService['initPromise'] = null; AtlasService['oidcPluginSyncedFromLoggerState'] = 'initial'; sandbox.resetHistory(); @@ -516,14 +517,38 @@ describe('AtlasServiceMain', function () { }); describe('init', function () { - describe('with no stored auth state', function () { - it('should set service to unauthenticated state', 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' + ); }); - describe('with stored auth state', function () { - it('should set service to unauthenticated state if token expired', function () {}); - - it('should set service to autenticated state if token can be refreshed', function () {}); + 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 7ffdd5bad5e..8b32eaf2d01 100644 --- a/packages/atlas-service/src/main.ts +++ b/packages/atlas-service/src/main.ts @@ -159,17 +159,17 @@ export class AtlasService { * │ `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:auth-attempt-succeeded` │ - * no └────────────────────────────────────────────┘ - * ↓ + * │ ↓ │ ┌────────────────────────────────────────────┐ + * │ 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:auth-attempt-succeeded` │ + * ↓ │ ↓ │ + * no │ `plugin:state-updated` │ + * ↓ └────────────────────────────────────────────┘ * `plugin:auth-failed` */ private static oidcPluginLogger: MongoDBOIDCPluginLogger & { @@ -182,7 +182,7 @@ export class AtlasService { private static oidcPluginSyncedFromLoggerState: AtlasServiceAuthState = 'initial'; - private static plugin: MongoDBOIDCPlugin; + private static plugin: MongoDBOIDCPlugin | null = null; private static token: Token | null = null; @@ -242,7 +242,9 @@ export class AtlasService { : ['auth-code']; } - static init(): Promise { + static init( + _createMongoDBOIDCPlugin = createMongoDBOIDCPlugin + ): Promise { return (this.initPromise ??= (async () => { ipcExpose( 'AtlasService', @@ -264,7 +266,7 @@ export class AtlasService { ); this.oidcPluginSyncedFromLoggerState = 'restoring'; const serializedState = await this.secretStore.getItem(SECRET_STORE_KEY); - this.plugin = createMongoDBOIDCPlugin({ + this.plugin = _createMongoDBOIDCPlugin({ redirectServerRequestHandler(data) { if (data.result === 'redirecting') { const { res, status, location } = data; @@ -306,6 +308,12 @@ export class AtlasService { private static async requestOAuthToken({ signal, }: { signal?: AbortSignal } = {}) { + if (!this.plugin) { + throw new Error( + 'Trying to use the oidc-plugin before service is initialised' + ); + } + return this.plugin.mongoClientOptions.authMechanismProperties.REQUEST_TOKEN_CALLBACK( { clientId: this.clientId, issuer: this.issuer }, { @@ -371,11 +379,11 @@ export class AtlasService { // proceed by requesting token once again from the plugin to sync // token back to the atlas service state await Promise.race([ - // NB: While `auth-{succeeded,failed}` events can also fire when one - // of the actual sign in flows is activated (see logger event - // diagram), we make sure that only token refresh is allowed during - // this process by providing an empty array of allowedFlows to the - // plugin through the `getAllowedFlows` method + // 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, }), @@ -633,7 +641,7 @@ export class AtlasService { } static async onExit() { - // Haven't even c the oidc-plugin yet, nothing to store + // Haven't even created the oidc-plugin yet, nothing to store if (!this.plugin) { return; } From 6bdcbdebd80f50e1a2bf29acf1863329c1498897 Mon Sep 17 00:00:00 2001 From: Sergey Petushkov Date: Mon, 14 Aug 2023 11:20:54 +0200 Subject: [PATCH 5/6] chore(atlas-service): small updates to the plugin flow diagram --- packages/atlas-service/src/main.ts | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/packages/atlas-service/src/main.ts b/packages/atlas-service/src/main.ts index 8b32eaf2d01..5f80a41c947 100644 --- a/packages/atlas-service/src/main.ts +++ b/packages/atlas-service/src/main.ts @@ -138,7 +138,7 @@ export class AtlasService { * ↓ * serialized state provided? → no → no event * ↓ - * yes → failed to deserialize? → no → no event + * yes → failed to deserialize? → no → `plugin:state-updated` * ↓ * yes → `plugin:deserialization-failed` * @@ -166,9 +166,9 @@ export class AtlasService { * ↓ ↓ │ ↓ ↓ │ * `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:auth-attempt-succeeded` │ + * `plugin:auth-succeeded` ←─────── yes ←──────── do we have a new token set in state? ← │ `plugin:state-updated` │ * ↓ │ ↓ │ - * no │ `plugin:state-updated` │ + * no │ `plugin:auth-attempt-succeeded` │ * ↓ └────────────────────────────────────────────┘ * `plugin:auth-failed` */ From d5eafcfe7c133210d02120473ea8ee7c00778d2a Mon Sep 17 00:00:00 2001 From: Sergey Petushkov Date: Mon, 14 Aug 2023 11:41:15 +0200 Subject: [PATCH 6/6] chore(atlas-service): don't use Promise inside async functions Co-authored-by: Anna Henningsen --- packages/atlas-service/src/main.ts | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/packages/atlas-service/src/main.ts b/packages/atlas-service/src/main.ts index 5f80a41c947..d16db5701d8 100644 --- a/packages/atlas-service/src/main.ts +++ b/packages/atlas-service/src/main.ts @@ -105,7 +105,7 @@ const defaultSecretStore: SecretStore = { } return (await keytar.getPassword(appName, key)) ?? undefined; } catch { - return Promise.resolve(undefined); + return undefined; } }, async setItem(key: string, value: string) { @@ -119,7 +119,7 @@ const defaultSecretStore: SecretStore = { } return await keytar.setPassword(appName, key, value); } catch { - return Promise.resolve(); + return; } }, };