Store default account in disk cache (#220)

Given an access key ID, the account ID is fixed.
This means that we don't need to look up the default
account every time the toolkit is used.

This change will store the mapping between access key
and account ID in `~/.cdk/cache/accounts.json`
and will fetch the account ID from there if it exists.

* Nuke cache if exceeds 1000 entries

packages/aws-cdk/lib/api/util/account-cache.ts

@@ -0,0 +1,93 @@
+import * as fs from 'fs-extra';
+import * as os from 'os';
+import * as path from 'path';
+import { debug } from '../../logging';
+
+
+/**
+ * Disk cache which maps access key IDs to account IDs.
+ * Usage:
+ *   cache.get(accessKey) => accountId | undefined
+ *   cache.put(accessKey, accountId)
+ */
+export class AccountAccessKeyCache {
+    /**
+     * Max number of entries in the cache, after which the cache will be reset.
+     */
+    public static readonly MAX_ENTRIES = 1000;
+
+    private readonly cacheFile: string;
+
+    /**
+     * @param filePath Path to the cache file
+     */
+    constructor(filePath?: string) {
+        this.cacheFile = filePath || path.join(os.homedir(), '.cdk', 'cache', 'accounts.json');
+    }
+
+    /**
+     * Tries to fetch the account ID from cache. If it's not in the cache, invokes
+     * the resolver function which should retrieve the account ID and return it.
+     * Then, it will be stored into disk cache returned.
+     *
+     * Example:
+     *
+     *      const accountId = cache.fetch(accessKey, async () => {
+     *          return await fetchAccountIdFromSomewhere(accessKey);
+     *      });
+     *
+     * @param accessKeyId
+     * @param resolver
+     */
+    public async fetch(accessKeyId: string, resolver: () => Promise<string | undefined>) {
+        // try to get account ID based on this access key ID from disk.
+        const cached = await this.get(accessKeyId);
+        if (cached) {
+            debug(`Retrieved account ID ${cached} from disk cache`);
+            return cached;
+        }
+
+        // if it's not in the cache, resolve and put in cache.
+        const accountId = await resolver();
+        if (accountId) {
+            await this.put(accessKeyId, accountId);
+        }
+
+        return accountId;
+    }
+
+    /** Get the account ID from an access key or undefined if not in cache */
+    public async get(accessKeyId: string): Promise<string | undefined> {
+        const map = await this.loadMap();
+        return map[accessKeyId];
+    }
+
+    /** Put a mapping betweenn access key and account ID */
+    public async put(accessKeyId: string, accountId: string) {
+        let map = await this.loadMap();
+
+        // nuke cache if it's too big.
+        if (Object.keys(map).length >= AccountAccessKeyCache.MAX_ENTRIES) {
+            map = { };
+        }
+
+        map[accessKeyId] = accountId;
+        await this.saveMap(map);
+    }
+
+    private async loadMap(): Promise<{ [accessKeyId: string]: string }> {
+        if (!(await fs.pathExists(this.cacheFile))) {
+            return { };
+        }
+
+        return await fs.readJson(this.cacheFile);
+    }
+
+    private async saveMap(map: { [accessKeyId: string]: string }) {
+        if (!(await fs.pathExists(this.cacheFile))) {
+            await fs.mkdirs(path.dirname(this.cacheFile));
+        }
+
+        await fs.writeJson(this.cacheFile, map, { spaces: 2 });
+    }
+}

packages/aws-cdk/lib/api/util/sdk.ts

@@ -1,8 +1,9 @@
 import { Environment} from '@aws-cdk/cx-api';
-import { CloudFormation, config, CredentialProviderChain , EC2, S3, SSM, STS } from 'aws-sdk';
+import { CloudFormation, config, CredentialProviderChain, EC2, S3, SSM, STS } from 'aws-sdk';
 import { debug } from '../../logging';
 import { PluginHost } from '../../plugin';
 import { CredentialProviderSource, Mode } from '../aws-auth/credentials';
+import { AccountAccessKeyCache } from './account-cache';
 
 /**
  * Source for SDK client objects
@@ -17,6 +18,7 @@ export class SDK {
     private defaultAccountFetched = false;
     private defaultAccountId?: string = undefined;
     private readonly userAgent: string;
+    private readonly accountCache = new AccountAccessKeyCache();
 
     constructor() {
         // Find the package.json from the main toolkit
@@ -70,11 +72,30 @@ export class SDK {
 
     private async lookupDefaultAccount() {
         try {
-            debug('Looking up default account ID from STS');
-            const result = await new STS().getCallerIdentity().promise();
-            return result.Account;
+            debug('Resolving default credentials');
+            const chain = new CredentialProviderChain();
+            const creds = await chain.resolvePromise();
+            const accessKeyId = creds.accessKeyId;
+            if (!accessKeyId) {
+                throw new Error('Unable to resolve AWS credentials (setup with "aws configure")');
+            }
+
+            const accountId = await this.accountCache.fetch(creds.accessKeyId, async () => {
+                // if we don't have one, resolve from STS and store in cache.
+                debug('Looking up default account ID from STS');
+                const result = await new STS().getCallerIdentity().promise();
+                const aid = result.Account;
+                if (!aid) {
+                    debug('STS didn\'t return an account ID');
+                    return undefined;
+                }
+                debug('Default account ID:', aid);
+                return aid;
+            });
+
+            return accountId;
         } catch (e) {
-            debug('Unable to retrieve default account from STS:', e);
+            debug('Unable to determine the default AWS account (did you configure "aws configure"?):', e);
             return undefined;
         }
     }

packages/aws-cdk/test/test.account-cache.ts

@@ -0,0 +1,85 @@
+import * as fs from 'fs-extra';
+import { ICallbackFunction, Test } from 'nodeunit';
+import * as path from 'path';
+import { AccountAccessKeyCache } from '../lib/api/util/account-cache';
+
+export = {
+    'setUp'(cb: ICallbackFunction) {
+        const self = this as any;
+        fs.mkdtemp('/tmp/account-cache-test').then(dir => {
+            self.file = path.join(dir, 'cache.json');
+            self.cache = new AccountAccessKeyCache(self.file);
+            return cb();
+        });
+    },
+
+    'tearDown'(cb: ICallbackFunction) {
+        const self = this as any;
+        fs.remove(path.dirname(self.file)).then(cb);
+    },
+
+    async 'get(k) when cache is empty'(test: Test) {
+        const self = this as any;
+        const cache: AccountAccessKeyCache = self.cache;
+        test.equal(await cache.get('foo'), undefined, 'get returns undefined');
+        test.equal(await fs.pathExists(self.file), false, 'cache file is not created');
+        test.done();
+    },
+
+    async 'put(k,v) and then get(k)'(test: Test) {
+        const self = this as any;
+        const cache: AccountAccessKeyCache = self.cache;
+
+        await cache.put('key', 'value');
+        await cache.put('boo', 'bar');
+        test.deepEqual(await cache.get('key'), 'value', '"key" is mapped to "value"');
+
+        // create another cache instance on the same file, should still work
+        const cache2 = new AccountAccessKeyCache(self.file);
+        test.deepEqual(await cache2.get('boo'), 'bar', '"boo" is mapped to "bar"');
+
+        // whitebox: read the file
+        test.deepEqual(await fs.readJson(self.file), {
+            key: 'value',
+            boo: 'bar'
+        });
+
+        test.done();
+    },
+
+    async 'fetch(k, resolver) can be used to "atomically" get + resolve + put'(test: Test) {
+        const self = this as any;
+        const cache: AccountAccessKeyCache = self.cache;
+
+        test.deepEqual(await cache.get('foo'), undefined, 'no value for "foo" yet');
+        test.deepEqual(await cache.fetch('foo', async () => 'bar'), 'bar', 'resolved by fetch and returned');
+        test.deepEqual(await cache.get('foo'), 'bar', 'stored in cache by fetch()');
+        test.done();
+    },
+
+    async 'cache is nuked if it exceeds 1000 entries'(test: Test) {
+        const self = this as any;
+        const cache: AccountAccessKeyCache = self.cache;
+
+        for (let i = 0; i < AccountAccessKeyCache.MAX_ENTRIES; ++i) {
+            await cache.put(`key${i}`, `value${i}`);
+        }
+
+        // verify all 1000 values are on disk
+        const otherCache = new AccountAccessKeyCache(self.file);
+        for (let i = 0; i < AccountAccessKeyCache.MAX_ENTRIES; ++i) {
+            test.equal(await otherCache.get(`key${i}`), `value${i}`);
+        }
+
+        // add another value
+        await cache.put('nuke-me', 'genesis');
+
+        // now, we expect only `nuke-me` to exist on disk
+        test.equal(await otherCache.get('nuke-me'), 'genesis');
+        for (let i = 0; i < AccountAccessKeyCache.MAX_ENTRIES; ++i) {
+            test.equal(await otherCache.get(`key${i}`), undefined);
+        }
+
+        test.done();
+    }
+};