fix: clear instance cache for init with fresh credentials (#106)

adobe · Apr 25, 2023 · 85f0542 · 85f0542
1 parent df56cfa
commit 85f0542
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 5 deletions.
diff --git a/lib/impl/CosmosStateStore.js b/lib/impl/CosmosStateStore.js
@@ -73,12 +73,13 @@ class CosmosStateStore extends StateStore {
    * @override
    * @private
    */
-  constructor (container, partitionKey) {
+  constructor (container, partitionKey, /* istanbul ignore next */ options = { expiration: null }) {
     super()
     /** @private */
     this._cosmos = {}
     this._cosmos.container = container
     this._cosmos.partitionKey = partitionKey
+    this.expiration = options.expiration
   }
 
   /**
@@ -102,7 +103,7 @@ class CosmosStateStore extends StateStore {
       containerId: joi.string().required(),
       partitionKey: joi.string().required(),
 
-      expiration: joi.string() // allowed for tvm response
+      expiration: joi.string() // allowed for tvm response, in ISO format
     }).xor('masterKey', 'resourceToken').required()
       .validate(credentials)
     if (validation.error) {
@@ -112,9 +113,16 @@ class CosmosStateStore extends StateStore {
       }))
     }
 
+    const inMemoryInstance = CosmosStateStore.inMemoryInstance[credentials.partitionKey]
+    if (inMemoryInstance && inMemoryInstance.expiration !== credentials.expiration) {
+      // the TVM credentials have changed, aio-lib-core-tvm has generated new one likely because of expiration.
+      delete CosmosStateStore.inMemoryInstance[credentials.partitionKey]
+    }
+
     if (!CosmosStateStore.inMemoryInstance[credentials.partitionKey]) {
       let cosmosClient
       if (credentials.resourceToken) {
+        // Note: resourceToken doesn't necessarily mean that the TVM provided the credentials, it can have been provided by a user.
         logger.debug('using azure cosmos resource token')
         cosmosClient = new cosmos.CosmosClient({ endpoint: credentials.endpoint, consistencyLevel: 'Session', tokenProvider: /* istanbul ignore next */ async () => credentials.resourceToken })
       } else {
@@ -125,7 +133,7 @@ class CosmosStateStore extends StateStore {
         // container = (await database.containers.createIfNotExists({ id: credentials.containerId })).container
       }
       const container = cosmosClient.database(credentials.databaseId).container(credentials.containerId)
-      CosmosStateStore.inMemoryInstance[credentials.partitionKey] = new CosmosStateStore(container, credentials.partitionKey)
+      CosmosStateStore.inMemoryInstance[credentials.partitionKey] = new CosmosStateStore(container, credentials.partitionKey, { expiration: credentials.expiration })
     } else {
       logger.debug('reusing exising in-memory CosmosClient initialization')
     }

diff --git a/lib/init.js b/lib/init.js
@@ -88,8 +88,6 @@ async function wrapTVMRequest (promise, params) {
  * @returns {Promise<StateStore>} A StateStore instance
  */
 async function init (config = {}) {
-  // todo joi-validate config here or leave it to StateStore impl + TvmClient?
-
   // 0. log
   const logConfig = utils.withHiddenFields(config, ['ow.auth', 'cosmos.resourceToken', 'cosmos.masterKey'])
 

diff --git a/test/impl/CosmosStateStore.test.js b/test/impl/CosmosStateStore.test.js
@@ -181,6 +181,7 @@ describe('init', () => {
         expect(cosmos.CosmosClient).toHaveBeenCalledTimes(0)
       })
       test('successive calls should reuse the CosmosStateStore instance - with masterkey', async () => {
+        // note this test may be confusing as no reuse is made with BYO credentials, but the cache is actually deleted by the top level init file. See test below.
         await testInitOK(fakeCosmosMasterCredentials)
         expect(cosmos.CosmosClient).toHaveBeenCalledTimes(1)
         cosmos.CosmosClient.mockReset()
@@ -194,6 +195,17 @@ describe('init', () => {
         // New CosmosClient instance generated again
         expect(cosmos.CosmosClient).toHaveBeenCalledTimes(2)
       })
+      test('No reuse if TVM credential expiration changed', async () => {
+        await testInitOK(fakeCosmosTVMResponse)
+        expect(cosmos.CosmosClient).toHaveBeenCalledTimes(1)
+        await CosmosStateStore.init({ ...fakeCosmosTVMResponse, expiration: new Date(0).toISOString() })
+        expect(cosmos.CosmosClient).toHaveBeenCalledTimes(2)
+        await CosmosStateStore.init({ ...fakeCosmosTVMResponse, expiration: new Date(1).toISOString() })
+        expect(cosmos.CosmosClient).toHaveBeenCalledTimes(3)
+        // double check, same credentials no additional call
+        await CosmosStateStore.init({ ...fakeCosmosTVMResponse, expiration: new Date(1).toISOString() })
+        expect(cosmos.CosmosClient).toHaveBeenCalledTimes(3)
+      })
     })
   })
 })