Small simplifications

Author: rix0rrr

packages/@aws-cdk/aws-ec2/lib/user-data.ts

@@ -296,7 +296,7 @@ export interface MultipartBodyOptions {
   /**
    * `Content-Transfer-Encoding` header specifying part encoding.
    *
-   * @default undefined - don't add this header
+   * @default undefined - body is not encoded
    */
   readonly transferEncoding?: string;
 
@@ -312,6 +312,15 @@ export interface MultipartBodyOptions {
  * The base class for all classes which can be used as {@link MultipartUserData}.
  */
 export abstract class MultipartBody {
+  /**
+   * Content type for shell scripts
+   */
+  public static readonly SHELL_SCRIPT = 'text/x-shellscript; charset="utf-8"';
+
+  /**
+   * Content type for boot hooks
+   */
+  public static readonly CLOUD_BOOTHOOK = 'text/cloud-boothook; charset="utf-8"';
 
   /**
    * Constructs the new `MultipartBody` wrapping existing `UserData`. Modification to `UserData` are reflected
@@ -336,83 +345,72 @@ export abstract class MultipartBody {
     return new MultipartBodyRaw(opts);
   }
 
-  protected static readonly DEFAULT_CONTENT_TYPE = 'text/x-shellscript; charset="utf-8"';
-
-  /** The body of this MIME part. */
-  public abstract get body(): string | undefined;
-
-  /** `Content-Type` header of this part */
-  public abstract get contentType(): string;
+  public constructor() {
+  }
 
   /**
-   * `Content-Transfer-Encoding` header specifying part encoding.
+   * Render body part as the string.
    *
-   * @default undefined - don't add this header
+   * Subclasses should not add leading nor trailing new line characters (\r \n)
    */
-  public abstract get transferEncoding(): string | undefined;
+  public abstract renderBodyPart(): string[];
+}
 
-  public constructor() {
+/**
+ * The raw part of multi-part user data, which can be added to {@link MultipartUserData}.
+ */
+class MultipartBodyRaw extends MultipartBody {
+  public constructor(private readonly props: MultipartBodyOptions) {
+    super();
   }
 
   /**
    * Render body part as the string.
-   *
-   * Subclasses should not add leading nor trailing new line characters (\r \n)
    */
-  public renderBodyPart(): string {
+  public renderBodyPart(): string[] {
     const result: string[] = [];
 
-    result.push(`Content-Type: ${this.contentType}`);
+    result.push(`Content-Type: ${this.props.contentType}`);
 
-    if (this.transferEncoding != null) {
-      result.push(`Content-Transfer-Encoding: ${this.transferEncoding}`);
+    if (this.props.transferEncoding != null) {
+      result.push(`Content-Transfer-Encoding: ${this.props.transferEncoding}`);
     }
     // One line free after separator
     result.push('');
 
-    if (this.body != null) {
-      result.push(this.body);
+    if (this.props.body != null) {
+      result.push(this.props.body);
       // The new line added after join will be consumed by encapsulating or closing boundary
     }
 
-    return result.join('\n');
-  }
-}
-
-/**
- * The raw part of multi-part user data, which can be added to {@link MultipartUserData}.
- */
-class MultipartBodyRaw extends MultipartBody {
-  public readonly body: string | undefined;
-  public readonly contentType: string;
-  public readonly transferEncoding: string | undefined;
-
-  public constructor(props: MultipartBodyOptions) {
-    super();
-
-    this.body = props.body;
-    this.contentType = props.contentType;
+    return result;
   }
 }
 
 /**
  * Wrapper for `UserData`.
  */
 class MultipartBodyUserDataWrapper extends MultipartBody {
+  private readonly contentType: string;
 
-  public readonly contentType: string;
-  public readonly transferEncoding: string | undefined;
-
-  public constructor(public readonly userData: UserData, contentType?: string) {
+  public constructor(private readonly userData: UserData, contentType?: string) {
     super();
 
-    this.contentType = contentType || MultipartBody.DEFAULT_CONTENT_TYPE;
-    this.transferEncoding = 'base64';
+    this.contentType = contentType || MultipartBody.SHELL_SCRIPT;
   }
 
-  public get body(): string {
-    // Wrap rendered user data with Base64 function, in case data contains non ASCII characters
-    return Fn.base64(this.userData.render());
+  /**
+   * Render body part as the string.
+   */
+  public renderBodyPart(): string[] {
+    const result: string[] = [];
+
+    result.push(`Content-Type: ${this.contentType}`);
+    result.push('Content-Transfer-Encoding: base64');
+    result.push('');
+    result.push(Fn.base64(this.userData.render()));
+
+    return result;
   }
 }
 
@@ -469,10 +467,21 @@ export class MultipartUserData extends UserData {
   /**
    * Adds a part to the list of parts.
    */
-  public addPart(part: MultipartBody): this {
+  public addPart(part: MultipartBody) {
     this.parts.push(part);
+  }
 
-    return this;
+  /**
+   * Adds a multipart part based on a UserData object
+   *
+   * This is the same as calling:
+   *
+   * ```ts
+   * multiPart.addPart(MultipartBody.fromUserData(userData, contentType));
+   * ```
+   */
+  public addUserDataPart(userData: UserData, contentType?: string) {
+    this.addPart(MultipartBody.fromUserData(userData, contentType));
   }
 
   public render(): string {
@@ -491,7 +500,7 @@ export class MultipartUserData extends UserData {
     // Add parts - each part starts with boundary
     this.parts.forEach(part => {
       resultArchive.push(`--${boundary}`);
-      resultArchive.push(part.renderBodyPart());
+      resultArchive.push(...part.renderBodyPart());
     });
 
     // Add closing boundary

packages/@aws-cdk/aws-ec2/test/userdata.test.ts

@@ -275,22 +275,27 @@ nodeunitShim({
 
   'Linux user rendering multipart headers'(test: Test) {
     // GIVEN
+    const stack = new Stack();
     const linuxUserData = ec2.UserData.forLinux();
     linuxUserData.addCommands('echo "Hello world"');
 
     // WHEN
     const defaultRender1 = ec2.MultipartBody.fromUserData(linuxUserData);
     const defaultRender2 = ec2.MultipartBody.fromUserData(linuxUserData, 'text/cloud-boothook; charset=\"utf-8\"');
 
-    const defaultRender3 = ec2.MultipartBody.fromUserData(linuxUserData);
-
     // THEN
-    test.equals(defaultRender1.contentType, 'text/x-shellscript; charset=\"utf-8\"');
-    test.equals(defaultRender2.contentType, 'text/cloud-boothook; charset=\"utf-8\"');
-    test.equals(defaultRender3.contentType, 'text/x-shellscript; charset=\"utf-8\"');
-
-    test.equals(defaultRender1.transferEncoding, 'base64');
-    test.equals(defaultRender2.transferEncoding, 'base64');
+    expect(stack.resolve(defaultRender1.renderBodyPart())).toEqual([
+      'Content-Type: text/x-shellscript; charset=\"utf-8\"',
+      'Content-Transfer-Encoding: base64',
+      '',
+      { 'Fn::Base64': '#!/bin/bash\necho \"Hello world\"' },
+    ]);
+    expect(stack.resolve(defaultRender2.renderBodyPart())).toEqual([
+      'Content-Type: text/cloud-boothook; charset=\"utf-8\"',
+      'Content-Transfer-Encoding: base64',
+      '',
+      { 'Fn::Base64': '#!/bin/bash\necho \"Hello world\"' },
+    ]);
 
     test.done();
   },