feat(vms): implement SNS topic command to trigger Lambda that stops t…

…he VM after initialization
privacy-scaling-explorations · Jun 29, 2023 · f5f73bb · f5f73bb
1 parent 2c426cb
commit f5f73bb
Show file tree

Hide file tree

Showing 4 changed files with 26 additions and 6 deletions.
diff --git a/packages/actions/src/helpers/vm.ts b/packages/actions/src/helpers/vm.ts
@@ -73,17 +73,24 @@ export const vmBootstrapCommand = (bucketName: string): Array<string> => [
  * Return the list of Node environment (and packages) installation plus artifact caching for contribution verification.
  * @param zKeyPath <string> - the path to zKey artifact inside AWS S3 bucket.
  * @param potPath <string> - the path to ptau artifact inside AWS S3 bucket.
+ * @param snsTopic <string> - the SNS topic ARN.
  * @returns <Array<string>> - the array of commands to be run by the EC2 instance.
  */
-export const vmDependenciesAndCacheArtifactsCommand = (zKeyPath: string, potPath: string): Array<string> => [
+export const vmDependenciesAndCacheArtifactsCommand = (
+    zKeyPath: string, 
+    potPath: string, 
+    snsTopic: string,
+    ): Array<string> => [
     "#!/bin/bash",
     "sudo yum update -y",
     "sudo yum install -y nodejs",
     "npm install -g snarkjs",
     `aws s3 cp s3://${zKeyPath} /var/tmp/genesisZkey.zkey`,
     `aws s3 cp s3://${potPath} /var/tmp/pot.ptau`,
     "wget https://github.com/BLAKE3-team/BLAKE3/releases/download/1.4.0/b3sum_linux_x64_bin -O /var/tmp/blake3.bin",
-    "chmod +x /var/tmp/blake3.bin"
+    "chmod +x /var/tmp/blake3.bin",
+    "INSTANCE_ID=$(curl -s http://169.254.169.254/latest/meta-data/instance-id)",
+    `aws sns publish --topic-arn ${snsTopic} --message "$INSTANCE_ID"`
 ]
 
 /**
@@ -164,6 +171,10 @@ export const createEC2Instance = async (
                     {
                         Key: "Name",
                         Value: ec2InstanceTag
+                    },
+                    {
+                        Key: "Initialized",
+                        Value: "false" 
                     }
                 ]
             }

diff --git a/packages/backend/.default.env b/packages/backend/.default.env
@@ -16,6 +16,8 @@ AWS_CEREMONY_BUCKET_POSTFIX="-ph2-ceremony"
 AWS_AMI_ID="ami-022e1a32d3f742bd8"
 # The EC2 instance role to access S3
 AWS_ROLE_ARN="YOUR-AWS-ROLE-ARN"
+# The SNS topic ARN to publish notifications 
+AWS_SNS_TOPIC_ARN="YOUR-AWS-SNS-TOPIC-ARN"
 
 ### GENERIC ###
 ### These configs are generic and not tied to Firebase or AWS services.

diff --git a/packages/backend/src/functions/ceremony.ts b/packages/backend/src/functions/ceremony.ts
@@ -30,7 +30,8 @@ import {
     getFinalContribution,
     htmlEncodeCircuitData,
     createEC2Client,
-    uploadFileToBucketNoFile
+    uploadFileToBucketNoFile,
+    getAWSVariables
 } from "../lib/utils"
 import { LogLevel } from "../types/enums"
 
@@ -148,10 +149,14 @@ export const setupCeremony = functions
                 // Get EC2 client.
                 const ec2Client = await createEC2Client()
 
+                // Get AWS variables.
+                const { snsTopic } = getAWSVariables()
+
                 // Prepare dependencies and cache artifacts command.
                 const vmCommands = vmDependenciesAndCacheArtifactsCommand(
                     `${bucketName}/${circuit.files?.initialZkeyStoragePath!}`,
-                    `${bucketName}/${circuit.files?.potStoragePath!}`
+                    `${bucketName}/${circuit.files?.potStoragePath!}`,
+                    snsTopic
                 )
 
                 printLog(vmCommands.join("\n"), LogLevel.DEBUG)

diff --git a/packages/backend/src/lib/utils.ts b/packages/backend/src/lib/utils.ts
@@ -400,7 +400,8 @@ export const getAWSVariables = (): any => {
         !process.env.AWS_ACCESS_KEY_ID ||
         !process.env.AWS_SECRET_ACCESS_KEY ||
         !process.env.AWS_ROLE_ARN ||
-        !process.env.AWS_AMI_ID
+        !process.env.AWS_AMI_ID ||
+        !process.env.AWS_SNS_TOPIC
     )
         logAndThrowError(COMMON_ERRORS.CM_WRONG_CONFIGURATION)
 
@@ -409,7 +410,8 @@ export const getAWSVariables = (): any => {
         secretAccessKey: process.env.AWS_SECRET_ACCESS_KEY!,
         region: process.env.AWS_REGION || "us-east-1",
         roleArn: process.env.AWS_ROLE_ARN!,
-        amiId: process.env.AWS_AMI_ID!
+        amiId: process.env.AWS_AMI_ID!,
+        snsTopic: process.env.AWS_SNS_TOPIC!
     }
 }