Store default account in disk cache

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/package.json

@@ -11,7 +11,8 @@
     "prepare": "/bin/bash generate.sh && tslint -p . && tsc && chmod +x bin/cdk && pkglint",
     "watch": "tsc -w",
     "lint": "tsc && tslint -p . --force",
-    "pkglint": "pkglint -f"
+    "pkglint": "pkglint -f",
+    "test": "nodeunit test/test.*.js"
   },
   "author": {
     "name": "Amazon Web Services",

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();
+    }
+};