fix(vms): refactoring and code fixes

ctrlc03 · ctrlc03 · commit 0ebe401894be · 2023-06-23T10:59:55.000+01:00
diff --git a/packages/actions/src/helpers/ec2.ts b/packages/actions/src/helpers/ec2.ts
@@ -103,86 +103,15 @@ export const generateVMCommand = (zKeyPath: string, ptauPath: string): string[]
 
 /**
  * Determine the VM specs based on the circuit constraints (TODO)
- * @param circuitConstraints <string> the constraints of the circuit
+ * @param requiredSpace <string> the required disk space
+ * @return <any> the disk space and ram requirements
  */
-// export const determineVMSpecs = async (circuitConstraints: string) => {}
-
-// // RAM -> instanceId
-// const instancesTypes = {
-//     "t3.nano": {
-//         RAM: "0.5 GiB",
-//         VCPU: "2"
-//     },
-//     "t3.micro": {
-//         RAM: "1 GiB",
-//         VCPU: "2"
-//     },
-//     "t3.small": {
-//         RAM: "2 GiB",
-//         VCPU: "2"
-//     },
-//     "t3.medium": {
-//         RAM: "4 GiB",
-//         VCPU: "2"
-//     },
-//     "t3.large": {
-//         RAM: "8 GiB",
-//         VCPU: "2"
-//     },
-//     "t3.xlarge": {
-//         RAM: "16 GiB",
-//         VCPU: "4"
-//     },
-//     "t3.2xlarge": {
-//         RAM: "32 GiB",
-//         VCPU: "8"
-//     },
-//     "c5.9xlarge": {
-//         RAM: "36 GiB",
-//         VCPU: "36"
-//     },
-//     "c5.18xlarge": {
-//         RAM: "72 GiB",
-//         VCPU: "72"
-//     },
-//     "c5a.8xlarge": {
-//         RAM: "64 GiB",
-//         VCPU: "32"
-//     },
-//     "c5.12xlarge": {
-//         RAM: "96 GiB",
-//         VCPU: "48"
-//     },
-//     "c5a.16xlarge": {
-//         RAM: "128 GiB",
-//         VCPU: "64"
-//     },
-//     "c6i.32xlarge": {
-//         RAM: "256 GiB",
-//         VCPU: "128"
-//     },
-//     "m6a.32xlarge": {
-//         RAM: "512 GiB",
-//         VCPU: "128"
-//     }
-// }
-
-// 1. create ssh key in ec2 tab -> save the name
-// 2. IAM role: access to ssh key ("iam:GetSSHPublicKey",)
-// 3. IAM role: ec2 access
-// 4. ec2 give role for s3 access
-// 5. have an api (express) running on the vm
-// 6. have a script that runs on the vm that does the verification
-// 7. JWT Authorization: Bearer <token>
-// each circuit document needs to have the instance id of the vm
-/*
-{
-    bucket: "x",
-    action: "verify/checkStatus",
-    "zKeyIndex": 0,
-    "zKeyStoragePath": /circuit/..,
+export const determineVMSpecs = (requiredSpace: string) => {
+    return {
+        "vm": "c5.x9large",
+        "disk": 32
+    }
 }
-*/
 
 /**
  * Creates a new EC2 instance
@@ -192,6 +121,7 @@ export const generateVMCommand = (zKeyPath: string, ptauPath: string): string[]
  * @param amiId <string> the AMI ID to be used
  * @param keyName <string> the name of the key to be used
  * @param roleArn <string> the ARN of the role to be used
+ * @param volumeSize <number> the size of the volume to be used
  * @returns <Promise<P0tionEC2Instance>> the instance that was created
  */
 export const createEC2Instance = async (
@@ -200,7 +130,8 @@ export const createEC2Instance = async (
     instanceType: string,
     amiId: string,
     keyName: string,
-    roleArn: string
+    roleArn: string,
+    volumeSize: number 
 ): Promise<P0tionEC2Instance> => {
     // create the params
     const params = {
@@ -214,7 +145,17 @@ export const createEC2Instance = async (
             Arn: roleArn
         },
         // how to run commands on startup
-        UserData: Buffer.from(commands.join("\n")).toString("base64")
+        UserData: Buffer.from(commands.join("\n")).toString("base64"),
+        BlockDeviceMappings: [
+            {
+                DeviceName: '/dev/xvda',
+                Ebs: {
+                DeleteOnTermination: true,
+                VolumeSize: volumeSize, // size in GB
+                VolumeType: 'gp2', // change this as per your needs
+                },
+            },
+        ],
     }
 
     // create command
diff --git a/packages/actions/src/index.ts b/packages/actions/src/index.ts
@@ -143,5 +143,6 @@ export {
     getEC2Ip,
     generateVMCommand,
     runCommandOnEC2,
-    retrieveCommandOutput
+    retrieveCommandOutput,
+    determineVMSpecs
 } from "./helpers/ec2"
diff --git a/packages/actions/test/unit/ec2.test.ts b/packages/actions/test/unit/ec2.test.ts
@@ -2,7 +2,7 @@ import chai, { expect } from "chai"
 import chaiAsPromised from "chai-as-promised"
 import { EC2Client } from "@aws-sdk/client-ec2"
 import fetch from "@adobe/node-fetch-retry"
-import { cleanUpMockUsers, cleanUpRecursively, createMockUser, envType, generateUserPasswords, getStorageConfiguration, getTranscriptLocalFilePath, initializeAdminServices, initializeUserServices, sleep } from "../utils"
+import { cleanUpMockUsers, cleanUpRecursively, createMockUser, deleteBucket, deleteObjectFromS3, envType, generateUserPasswords, getStorageConfiguration, getTranscriptLocalFilePath, initializeAdminServices, initializeUserServices, sleep } from "../utils"
 import {  
     checkEC2Status, 
     createEC2Client, 
@@ -191,6 +191,7 @@ describe("VMs", () => {
 
             // 1 create a bucket for the ceremony
             await signInWithEmailAndPassword(userAuth, users[1].data.email, passwords[1])
+            console.log("LOGGED IN")
             await createS3Bucket(userFunctions, ceremonyBucket)
 
             // zkey upload
@@ -204,8 +205,13 @@ describe("VMs", () => {
                 await terminateEC2Instance(ec2, instanceId)
             }
 
+            for (const objectToDelete of objectsToDelete) {
+                await deleteObjectFromS3(ceremonyBucket, objectToDelete)
+            }
+            await deleteBucket(ceremonyBucket)
+
             await cleanUpMockUsers(adminAuth, adminFirestore, users)
-            await cleanUpRecursively(adminFirestore, ceremonyId)
+            // await cleanUpRecursively(adminFirestore, ceremonyId)
         })
 
         it("should create a ceremony and two VMs should spin up", async () => {
@@ -337,5 +343,7 @@ describe("VMs", () => {
                 String(process.env.FIREBASE_CF_URL_VERIFY_CONTRIBUTION)
             )
         })
+
+        it("should terminate the VM(s) when finalizaing the ceremony", async () => {})
     })
 })
diff --git a/packages/backend/src/functions/ceremony.ts b/packages/backend/src/functions/ceremony.ts
@@ -28,10 +28,13 @@ import {
     createEC2Client,
     getAWSVariables,
     sleep,
-    uploadFileToBucket
+    uploadFileToBucket,
+    createTemporaryLocalPath,
+    uploadFileToBucketNoFile
 } from "../lib/utils"
 import { LogLevel } from "../types/enums"
 import { writeFileSync, unlinkSync } from "fs"
+import { determineVMSpecs } from "@p0tion/actions"
 
 dotenv.config()
 
@@ -151,20 +154,20 @@ export const setupCeremony = functions
             )
 
             // upload the instructions file the bucket and clean up
-            const startupScriptPath = `/var/tmp/${startupScript}`
-            writeFileSync(startupScriptPath, vmCommands.join("\n"))
-            await uploadFileToBucket(bucketName, startupScript, startupScriptPath)
-            unlinkSync(startupScriptPath)
+            await uploadFileToBucketNoFile(bucketName, startupScript, vmCommands.join("\n"))
 
             const { amiId, keyName, roleArn } = getAWSVariables()
+
+            const vmSpecs = determineVMSpecs("32")
             // as well as the VM configuration 
             const instance = await createEC2Instance(
                 ec2Client,
                 userData,
-                "t3.small",
+                "t3.xlarge",
                 amiId,
                 keyName,
-                roleArn
+                roleArn,
+                Number(vmSpecs.disk)
             )
 
             // html encode circuit data.
diff --git a/packages/backend/src/functions/storage.ts b/packages/backend/src/functions/storage.ts
@@ -146,9 +146,9 @@ export const createBucket = functions
         // Prepare S3 command.
         const command = new CreateBucketCommand({
             Bucket: data.bucketName,
-            CreateBucketConfiguration: {
-                LocationConstraint: String(process.env.AWS_REGION)
-            },
+            // CreateBucketConfiguration: {
+            //     LocationConstraint: String(process.env.AWS_REGION)
+            // },
             ObjectOwnership: "BucketOwnerPreferred"
         })
 
diff --git a/packages/backend/src/lib/utils.ts b/packages/backend/src/lib/utils.ts
@@ -258,6 +258,37 @@ export const uploadFileToBucket = async (
     if (response.status !== 200 || !response.ok) logAndThrowError(SPECIFIC_ERRORS.SE_STORAGE_UPLOAD_FAILED)
 }
 
+export const uploadFileToBucketNoFile = async (
+    bucketName: string,
+    objectKey: string,
+    data: string,
+    isPublic: boolean = false
+) => {
+    // Prepare AWS S3 client instance.
+    const client = await getS3Client()
+
+    // Prepare command.
+    const command = new PutObjectCommand({
+        Bucket: bucketName,
+        Key: objectKey,
+        ContentType: "text/plain",
+        ACL: isPublic ? "public-read" : "private"
+    })
+
+    // Generate a pre-signed url for uploading the file.
+    const url = await getSignedUrl(client, command, { expiresIn: Number(process.env.AWS_PRESIGNED_URL_EXPIRATION) })
+
+    // Execute upload request.
+    // @ts-ignore
+    const response = await fetch(url, {
+        method: "PUT",
+        body: data,
+        headers: { "Content-Type": "text/plain" }
+    })
+
+    if (response.status !== 200 || !response.ok) logAndThrowError(SPECIFIC_ERRORS.SE_STORAGE_UPLOAD_FAILED)
+}
+
 /**
  * Upload an artifact from the AWS S3 bucket.
  * @param bucketName <string> - the name of the bucket.
