fix(core): fix use of references in toJsonString()

When tokens are identified to be References, leave their type unchanged
in `toJsonString()`, so that they'll be properly identified as
references and handled in a potentially cross-stack way.

Add checking to references so that they'll be used correctly in tests.

packages/@aws-cdk/core/lib/private/cfn-reference.ts

@@ -101,7 +101,12 @@ export class CfnReference extends Reference {
   public resolve(context: IResolveContext): any {
     // If we have a special token for this consuming stack, resolve that. Otherwise resolve as if
     // we are in the same stack.
-    const token = this.replacementTokens.get(Stack.of(context.scope));
+    const consumingStack = Stack.of(context.scope);
+    const token = this.replacementTokens.get(consumingStack);
+    if (!token && this.isCrossStackReference(consumingStack) && !context.preparing) {
+      throw new Error('Should have hadd a cross-stack reference, but not found. prepare() the reference first.');
+    }
+
     if (token) {
       return token.resolve(context);
     } else {
@@ -120,9 +125,9 @@ export class CfnReference extends Reference {
     }
 
     // tslint:disable-next-line:max-line-length
-    if (this.producingStack && this.producingStack !== consumingStack && !this.replacementTokens.has(consumingStack)) {
+    if (!this.replacementTokens.has(consumingStack) && this.isCrossStackReference(consumingStack)) {
       // We're trying to resolve a cross-stack reference
-      consumingStack.addDependency(this.producingStack, `${consumingConstruct.node.path} -> ${this.target.node.path}.${this.originalDisplayName}`);
+      consumingStack.addDependency(this.producingStack!, `${consumingConstruct.node.path} -> ${this.target.node.path}.${this.originalDisplayName}`);
       this.replacementTokens.set(consumingStack, this.exportValue(consumingStack));
     }
   }
@@ -180,6 +185,10 @@ export class CfnReference extends Reference {
     const exportName = prefix + makeUniqueId(components);
     return exportName;
   }
+
+  private isCrossStackReference(consumingStack: Stack) {
+    return this.producingStack && this.producingStack !== consumingStack;
+  }
 }
 
 import { CfnElement } from "../cfn-element";

packages/@aws-cdk/core/lib/private/cloudformation-lang.ts

@@ -1,4 +1,5 @@
 import { Lazy } from "../lazy";
+import { Reference } from "../reference";
 import { DefaultTokenResolver, IFragmentConcatenator, IPostProcessor, IResolvable, IResolveContext  } from "../resolvable";
 import { TokenizedStringFragments } from "../string-fragments";
 import { Token } from "../token";
@@ -47,6 +48,11 @@ export class CloudFormationLang {
       }
 
       public resolveToken(t: IResolvable, context: IResolveContext, postProcess: IPostProcessor) {
+        // Return References directly, so their type is maintained and the references will
+        // continue to work.
+        if (Reference.isReference(t)) { return t; }
+
+        // Deep-resolve and wrap. This is necessary for Lazy tokens so we can see "inside" them.
         return wrap(super.resolveToken(t, context, postProcess));
       }
       public resolveString(fragments: TokenizedStringFragments, context: IResolveContext) {
@@ -60,6 +66,7 @@ export class CloudFormationLang {
     // We need a ResolveContext to get started so return a Token
     return Lazy.stringValue({ produce: (ctx: IResolveContext) =>
       JSON.stringify(resolve(obj, {
+        preparing: ctx.preparing,
         scope: ctx.scope,
         resolver: new IntrinsincWrapper()
       }), undefined, space)

packages/@aws-cdk/core/lib/private/resolve.ts

@@ -18,6 +18,7 @@ const tokenMap = TokenMap.instance();
  */
 export interface IResolveOptions {
   scope: IConstruct;
+  preparing: boolean;
   resolver: ITokenResolver;
   prefix?: string[];
 }
@@ -42,6 +43,7 @@ export function resolve(obj: any, options: IResolveOptions): any {
     let postProcessor: IPostProcessor | undefined;
 
     const context: IResolveContext = {
+      preparing: options.preparing,
       scope: options.scope,
       registerPostProcessor(pp) { postProcessor = pp; },
       resolve(x: any) { return resolve(x, { ...options, prefix: newPrefix }); },
@@ -168,7 +170,7 @@ export function resolve(obj: any, options: IResolveOptions): any {
 export function findTokens(scope: IConstruct, fn: () => any): IResolvable[] {
   const resolver = new RememberingTokenResolver(new StringConcat());
 
-  resolve(fn(), { scope, prefix: [], resolver });
+  resolve(fn(), { scope, prefix: [], resolver, preparing: true });
 
   return resolver.tokens;
 }

packages/@aws-cdk/core/lib/resolvable.ts

@@ -12,6 +12,11 @@ export interface IResolveContext {
    */
   readonly scope: IConstruct;
 
+  /**
+   * True when we are still preparing, false if we're rendering the final output
+   */
+  readonly preparing: boolean;
+
   /**
    * Resolve an inner object
    */

packages/@aws-cdk/core/lib/stack.ts

@@ -207,6 +207,7 @@ export class Stack extends Construct implements ITaggable {
       scope: this,
       prefix: [],
       resolver: CLOUDFORMATION_TOKEN_RESOLVER,
+      preparing: false
     });
   }
 

packages/@aws-cdk/core/lib/token.ts

@@ -111,7 +111,10 @@ export class Tokenization {
    * @param options Prefix key path components for diagnostics.
    */
   public static resolve(obj: any, options: ResolveOptions): any {
-    return resolve(obj, options);
+    return resolve(obj, {
+      ...options,
+      preparing: false
+    });
   }
 
   /**

packages/@aws-cdk/core/test/test.cloudformation-json.ts

@@ -1,5 +1,5 @@
 import { Test } from 'nodeunit';
-import { Fn, Lazy, Stack, Token } from '../lib';
+import { App, CfnOutput, Fn, Lazy, Stack, Token } from '../lib';
 import { Intrinsic } from '../lib/private/intrinsic';
 import { evaluateCFN } from './evaluate-cfn';
 
@@ -200,6 +200,39 @@ export = {
 
     test.done();
   },
+
+  'cross-stack references are also properly converted by toJsonString()'(test: Test) {
+    // GIVEN
+    const app = new App();
+    const stack1 = new Stack(app, 'Stack1');
+    const stack2 = new Stack(app, 'Stack2');
+
+    // WHEN
+    new CfnOutput(stack2, 'Stack1Id', { value: stack2.toJsonString({
+      Stack1Id: stack1.stackId,
+      Stack2Id: stack2.stackId,
+    })});
+
+    // THEN
+    const asm = app.synth();
+    test.deepEqual(asm.getStack('Stack2').template, {
+      Outputs: {
+        Stack1Id: {
+          Value: {
+            'Fn::Join': [ '', [
+              '{"Stack1Id":"',
+              { 'Fn::ImportValue': 'Stack1:ExportsOutputRefAWSStackIdB2DD5BAA' },
+              '","Stack2Id":"',
+              { Ref: 'AWS::StackId' },
+              '"}'
+            ] ]
+          }
+        }
+      }
+    });
+
+    test.done();
+  },
 };
 
 /**