From b503cb377edbf32df7e3d001ae37ab2e0dedf9d9 Mon Sep 17 00:00:00 2001
From: Daniel Neilson <ddneilson@ieee.org>
Date: Sun, 21 Aug 2022 22:08:35 -0500
Subject: [PATCH] feat(deadline): adds UserData property to WorkerInstanceFleet

This PR adds the ability to provide your own UserData object to
the deadline.WorkerInstanceFleet.
---
 .../aws-rfdk/lib/deadline/lib/worker-fleet.ts | 11 ++++++++++
 .../lib/deadline/test/worker-fleet.test.ts    | 20 +++++++++++++++++++
 .../lib/x509-certs/test/certificate.test.ts   |  4 +++-
 3 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/packages/aws-rfdk/lib/deadline/lib/worker-fleet.ts b/packages/aws-rfdk/lib/deadline/lib/worker-fleet.ts
index 2ec30b9e1..3501fc90d 100644
--- a/packages/aws-rfdk/lib/deadline/lib/worker-fleet.ts
+++ b/packages/aws-rfdk/lib/deadline/lib/worker-fleet.ts
@@ -28,6 +28,7 @@ import {
   Port,
   SubnetSelection,
   SubnetType,
+  UserData,
 } from 'aws-cdk-lib/aws-ec2';
 import {IApplicationLoadBalancerTarget} from 'aws-cdk-lib/aws-elasticloadbalancingv2';
 import {
@@ -223,6 +224,15 @@ export interface WorkerInstanceFleetProps extends WorkerSettings {
    */
   readonly blockDevices?: BlockDevice[];
 
+  /**
+   * The specific UserData to use.
+   *
+   * The UserData will be mutated by this construct and may be mutated afterwards as well.
+   *
+   * @default A UserData object appropriate for the MachineImage's Operating System is created.
+   */
+  readonly userData?: UserData;
+
   /**
    * An optional provider of user data commands to be injected at various points during the Worker configuration lifecycle.
    * You can provide a subclass of InstanceUserDataProvider with the methods overridden as desired.
@@ -455,6 +465,7 @@ export class WorkerInstanceFleet extends WorkerInstanceFleetBase {
       role: props.role,
       spotPrice: props.spotPrice?.toString(),
       blockDevices: props.blockDevices,
+      userData: props.userData,
     });
 
     this.targetCapacity = parseInt((this.fleet.node.defaultChild as CfnAutoScalingGroup).maxSize, 10);
diff --git a/packages/aws-rfdk/lib/deadline/test/worker-fleet.test.ts b/packages/aws-rfdk/lib/deadline/test/worker-fleet.test.ts
index c8ee5eab4..4d02714d1 100644
--- a/packages/aws-rfdk/lib/deadline/test/worker-fleet.test.ts
+++ b/packages/aws-rfdk/lib/deadline/test/worker-fleet.test.ts
@@ -28,6 +28,7 @@ import {
   Peer,
   SecurityGroup,
   SubnetType,
+  UserData,
   Vpc,
 } from 'aws-cdk-lib/aws-ec2';
 import {
@@ -477,6 +478,25 @@ test.each([
   });
 });
 
+test('worker fleet uses given UserData', () => {
+  // GIVEN
+  const id  = 'workerFleet';
+  const userData = UserData.forLinux();
+
+  // WHEN
+  const workerFleet = new WorkerInstanceFleet(stack, id, {
+    vpc,
+    workerMachineImage: new GenericLinuxImage({
+      'us-east-1': '123',
+    }),
+    renderQueue,
+    userData,
+  });
+
+  // THEN
+  expect(workerFleet.fleet.userData).toBe(userData);
+});
+
 test('default linux worker fleet is created correctly custom subnet values', () => {
   vpc = new Vpc(stack, 'VPC1Az', {
     maxAzs: 1,
diff --git a/packages/aws-rfdk/lib/lambdas/nodejs/lib/x509-certs/test/certificate.test.ts b/packages/aws-rfdk/lib/lambdas/nodejs/lib/x509-certs/test/certificate.test.ts
index a53aeca5f..6002ad96b 100644
--- a/packages/aws-rfdk/lib/lambdas/nodejs/lib/x509-certs/test/certificate.test.ts
+++ b/packages/aws-rfdk/lib/lambdas/nodejs/lib/x509-certs/test/certificate.test.ts
@@ -273,5 +273,7 @@ test('decrypt private key', async () => {
   // THEN
   expect(decryptedKey).toEqual(expectedDecryptedKey);
   // Must have the decrypted private key
-  expect(decryptedKey).toContain('-----BEGIN RSA PRIVATE KEY-----');
+  // OpenSSL 1.0.x: -----BEGIN RSA PRIVATE KEY-----
+  // OpenSSL 1.1.x: -----BEGIN PRIVATE KEY-----
+  expect(decryptedKey).toMatch(/-----BEGIN (?:RSA )?PRIVATE KEY-----/);
 });