feat: add SSM parameters for the load tests (#872)

Add AWS SSM parameters needed by the load tests to run. This includes changes to the Lambda so that it can read these parameters and load them into the execution environment. The reason these are not being passed in directly as Lambda environment variables is because these are stored as plain text and can be leaked by AWS CLI commands.
cds-snc · Oct 17, 2024 · b6204ff · b6204ff
1 parent 15c4d56
commit b6204ff
Show file tree

Hide file tree

Showing 14 changed files with 175 additions and 47 deletions.
diff --git a/.github/workflows/terragrunt-apply-staging.yml b/.github/workflows/terragrunt-apply-staging.yml
@@ -45,6 +45,9 @@ env:
   TF_VAR_cognito_code_template_id: 12a18f84-062c-4a67-8310-bf114af051ea
   TF_VAR_email_address_contact_us: ${{ vars.STAGING_CONTACT_US_EMAIL }}
   TF_VAR_email_address_support: ${{ vars.STAGING_SUPPORT_EMAIL }}
+  TF_VAR_load_testing_form_id: ${{ vars.STAGING_LOAD_TESTING_FORM_ID }}
+  TF_VAR_load_testing_form_private_key: ${{ vars.STAGING_LOAD_TESTING_FORM_PRIVATE_KEY }}
+  TF_VAR_load_testing_zitadel_app_private_key: ${{ vars.STAGING_ZITADEL_APPLICATION_KEY }}
   TF_VAR_zitadel_provider: ${{ vars.STAGING_ZITADEL_PROVIDER }}
   TF_VAR_zitadel_administration_key: ${{ secrets.STAGING_ZITADEL_ADMINISTRATION_KEY }}
   # IdP

diff --git a/.github/workflows/terragrunt-plan-all-staging.yml b/.github/workflows/terragrunt-plan-all-staging.yml
@@ -37,6 +37,9 @@ env:
   TF_VAR_cognito_code_template_id: 12a18f84-062c-4a67-8310-bf114af051ea
   TF_VAR_email_address_contact_us: ${{ vars.STAGING_CONTACT_US_EMAIL }}
   TF_VAR_email_address_support: ${{ vars.STAGING_SUPPORT_EMAIL }}
+  TF_VAR_load_testing_form_id: ${{ vars.STAGING_LOAD_TESTING_FORM_ID }}
+  TF_VAR_load_testing_form_private_key: ${{ vars.STAGING_LOAD_TESTING_FORM_PRIVATE_KEY }}
+  TF_VAR_load_testing_zitadel_app_private_key: ${{ vars.STAGING_ZITADEL_APPLICATION_KEY }}
   TF_VAR_zitadel_provider: ${{ vars.STAGING_ZITADEL_PROVIDER }}
   TF_VAR_zitadel_administration_key: ${{ secrets.STAGING_ZITADEL_ADMINISTRATION_KEY }}
   # IdP

diff --git a/.github/workflows/terragrunt-plan-staging.yml b/.github/workflows/terragrunt-plan-staging.yml
@@ -47,6 +47,9 @@ env:
   TF_VAR_cognito_code_template_id: 12a18f84-062c-4a67-8310-bf114af051ea
   TF_VAR_email_address_contact_us: ${{ vars.STAGING_CONTACT_US_EMAIL }}
   TF_VAR_email_address_support: ${{ vars.STAGING_SUPPORT_EMAIL }}
+  TF_VAR_load_testing_form_id: ${{ vars.STAGING_LOAD_TESTING_FORM_ID }}
+  TF_VAR_load_testing_form_private_key: ${{ vars.STAGING_LOAD_TESTING_FORM_PRIVATE_KEY }}
+  TF_VAR_load_testing_zitadel_app_private_key: ${{ vars.STAGING_ZITADEL_APPLICATION_KEY }}
   TF_VAR_zitadel_provider: ${{ vars.STAGING_ZITADEL_PROVIDER }}
   TF_VAR_zitadel_administration_key: ${{ secrets.STAGING_ZITADEL_ADMINISTRATION_KEY }}
   # IdP

diff --git a/aws/load_testing/inputs.tf b/aws/load_testing/inputs.tf
@@ -2,3 +2,26 @@ variable "ecr_repository_url_load_testing_lambda" {
   description = "URL of the Load Testing Lambda ECR"
   type        = string
 }
+
+variable "lambda_submission_function_name" {
+  description = "Name of the Submission Lambda function."
+  type        = string
+}
+
+variable "load_testing_form_id" {
+  description = "Form ID that will be used to generate, retrieve and confirm responses."
+  type        = string
+  sensitive   = true
+}
+
+variable "load_testing_form_private_key" {
+  description = "Private key JSON of the form that will be used to authenticate the API requests.  This must be a key from the `var.load_testing_form_id` form."
+  type        = string
+  sensitive   = true
+}
+
+variable "load_testing_zitadel_app_private_key" {
+  description = "Private key JSON of the Zitadel application to perform access token introspection requests."
+  type        = string
+  sensitive   = true
+}
diff --git a/aws/load_testing/lambda.tf b/aws/load_testing/lambda.tf
@@ -6,8 +6,8 @@ resource "aws_lambda_function" "load_testing" {
   image_uri     = "${var.ecr_repository_url_load_testing_lambda}:latest"
   function_name = "load-testing"
   role          = aws_iam_role.load_test_lambda.arn
-  timeout       = 300
-  memory_size   = 200
+  timeout       = 900
+  memory_size   = 1024
   package_type  = "Image"
   description   = "A function that runs a locust load test"
 
@@ -52,7 +52,44 @@ data "aws_iam_policy_document" "lambda_assume_policy" {
   }
 }
 
+resource "aws_iam_policy" "load_test_lambda" {
+  name        = "LoadTestLambda"
+  description = "Allow access to resources needed by the load testing Lambda function"
+  policy      = data.aws_iam_policy_document.load_test_lambda.json
+}
+
+resource "aws_iam_role_policy_attachment" "load_test_lambda" {
+  role       = aws_iam_role.load_test_lambda.name
+  policy_arn = aws_iam_policy.load_test_lambda.arn
+}
+
 resource "aws_iam_role_policy_attachment" "load_test_lambda_basic_access" {
   role       = aws_iam_role.load_test_lambda.name
   policy_arn = "arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole"
 }
+
+data "aws_iam_policy_document" "load_test_lambda" {
+  statement {
+    sid    = "GetSSMParameters"
+    effect = "Allow"
+    actions = [
+      "ssm:GetParameters",
+    ]
+    resources = [
+      aws_ssm_parameter.load_testing_form_id.arn,
+      aws_ssm_parameter.load_testing_form_private_key.arn,
+      aws_ssm_parameter.load_testing_zitadel_app_private_key.arn,
+    ]
+  }
+
+  statement {
+    sid    = "InvokeSubmissionLambda"
+    effect = "Allow"
+    actions = [
+      "lambda:InvokeFunction",
+    ]
+    resources = [
+      "arn:aws:lambda:${var.region}:${var.account_id}:function:${var.lambda_submission_function_name}",
+    ]
+  }
+}
diff --git a/aws/load_testing/parameters.tf b/aws/load_testing/parameters.tf
@@ -0,0 +1,23 @@
+resource "aws_ssm_parameter" "load_testing_form_id" {
+  # checkov:skip=CKV_AWS_337: default service encryption key is acceptable
+  name        = "/load-testing/form-id"
+  description = "Form ID that will be used to generate, retrieve and confirm responses."
+  type        = "SecureString"
+  value       = var.load_testing_form_id
+}
+
+resource "aws_ssm_parameter" "load_testing_form_private_key" {
+  # checkov:skip=CKV_AWS_337: default service encryption key is acceptable
+  name        = "/load-testing/form-private-key"
+  description = "Private key JSON of the form that will be used to authenticate the API requests.  This must be a key for the `/load-testing/form-id` form."
+  type        = "SecureString"
+  value       = var.load_testing_form_private_key
+}
+
+resource "aws_ssm_parameter" "load_testing_zitadel_app_private_key" {
+  # checkov:skip=CKV_AWS_337: default service encryption key is acceptable
+  name        = "/load-testing/zitadel-app-private-key"
+  description = "Private key JSON of the Zitadel application to perform access token introspection requests."
+  type        = "SecureString"
+  value       = var.load_testing_zitadel_app_private_key
+}
diff --git a/env/cloud/load_testing/terragrunt.hcl b/env/cloud/load_testing/terragrunt.hcl
@@ -3,7 +3,7 @@ terraform {
 }
 
 dependencies {
-  paths = ["../ecr"]
+  paths = ["../ecr", "../lambdas"]
 }
 
 dependency "ecr" {
@@ -16,8 +16,19 @@ dependency "ecr" {
   }
 }
 
+dependency "lambdas" {
+  config_path = "../lambdas"
+
+  mock_outputs_allowed_terraform_commands = ["init", "fmt", "validate", "plan", "show"]
+  mock_outputs_merge_strategy_with_state  = "shallow"
+  mock_outputs = {
+    lambda_submission_function_name = "Submission"
+  }
+}
+
 inputs = {
   ecr_repository_url_load_testing_lambda = dependency.ecr.outputs.ecr_repository_url_load_testing_lambda
+  lambda_submission_function_name        = dependency.lambdas.outputs.lambda_submission_function_name
 }
 
 include {

diff --git a/lambda-code/load-testing/Dockerfile b/lambda-code/load-testing/Dockerfile
@@ -1,12 +1,9 @@
-FROM amazon/aws-lambda-python:3.11@sha256:99cadc3bd9674a32a4ef694ff2e27f0b3d6c7f369b174db792b0099699fa0da4
-COPY main.py .
-COPY tests ./tests
-COPY requirements.txt .
-
-RUN yum -y groupinstall "Development Tools"
-
-RUN pip3 install --upgrade pip
+FROM amazon/aws-lambda-python:3.12@sha256:37b95206c4c78331f6d5cb0e8389ef573f39cfea01f73c530f28f3ac6f6493c7
 
+COPY requirements.txt .
 RUN pip3 install -r requirements.txt
 
+COPY main.py .
+COPY tests ./tests
+
 CMD ["main.handler"]
diff --git a/lambda-code/load-testing/README.md b/lambda-code/load-testing/README.md
@@ -0,0 +1,27 @@
+# Load testing
+Locust load tests that can be run in a Lambda function or locally.
+
+## Lambda
+Invoke the function using an event that looks like so:
+```json
+{
+  "locustfile": "./tests/locust_test_file.py",
+  "host": "https://forms-staging.cdssandbox.xyz",
+  "num_users": "5",
+  "spawn_rate": "1",
+  "run_time": "5m"
+}
+```
+
+## Locally
+You will need AWS access credentials for the target environment, along with the following environment variables set:
+```sh
+FORM_ID                 # Form ID to use for load testing
+FORM_PRIVATE_KEY        # JSON private key for the form (must be from the `FORM_ID` form)
+ZITADEL_APP_PRIVATE_KEY # JSON private key for the Zitadel application that is used for access token introspection
+```
+Once the variables are set, you can start the tests like so:
+```sh
+make install
+make locust
+```
diff --git a/lambda-code/load-testing/main.py b/lambda-code/load-testing/main.py
@@ -1,51 +1,52 @@
+import invokust
 import logging
 import os
 import boto3
-from invokust.aws_lambda import get_lambda_runtime_info
-from invokust import LocustLoadTest, create_settings
 
 logging.basicConfig(level=logging.INFO)
 
 ssm_client = boto3.client("ssm")
 
 
-def get_ssm_parameter(client, parameter_name):
-    response = client.get_parameter(Name=parameter_name, WithDecryption=True)
-    return response["Parameter"]["Value"]
+def get_ssm_parameters(client, parameter_names):
+    response = client.get_parameters(Names=parameter_names, WithDecryption=True)
+    return {param["Name"]: param["Value"] for param in response["Parameters"]}
+
 
 # Load required environment variables from AWS SSM
-os.environ["FORM_ID"] = get_ssm_parameter(ssm_client, "load-testing/form-id")
-os.environ["PRIVATE_API_KEY_APP_JSON"] = get_ssm_parameter(
-    ssm_client, "load-testing/private-api-key-app"
-)
-os.environ["PRIVATE_API_KEY_USER_JSON"] = get_ssm_parameter(
-    ssm_client, "load-testing/private-api-key-user"
+params = get_ssm_parameters(
+    ssm_client,
+    [
+        "/load-testing/form-id",
+        "/load-testing/form-private-key",
+        "/load-testing/zitadel-app-private-key",
+    ],
 )
-
+os.environ["FORM_ID"] = params["/load-testing/form-id"]
+os.environ["FORM_PRIVATE_KEY"] = params["/load-testing/form-private-key"]
+os.environ["ZITADEL_APP_PRIVATE_KEY"] = params["/load-testing/zitadel-app-private-key"]
 
 def handler(event=None, context=None):
 
     # Check for required environment variables
     required_env_vars = [
         "FORM_ID",
-        "PRIVATE_API_KEY_APP_JSON",
-        "PRIVATE_API_KEY_USER_JSON",
+        "FORM_PRIVATE_KEY",
+        "ZITADEL_APP_PRIVATE_KEY",
     ]
     for env_var in required_env_vars:
         if env_var not in os.environ:
             raise ValueError(f"Missing required environment variable: {env_var}")
 
     try:
         settings = (
-            create_settings(**event)
+            invokust.create_settings(**event)
             if event
-            else create_settings(from_environment=True)
+            else invokust.create_settings(from_environment=True)
         )
-        loadtest = LocustLoadTest(settings)
+        loadtest = invokust.LocustLoadTest(settings)
         loadtest.run()
     except Exception as e:
         logging.error("Exception running locust tests {0}".format(repr(e)))
     else:
-        locust_stats = loadtest.stats()
-        locust_stats.update(get_lambda_runtime_info(context))
-        return locust_stats
+        return loadtest.stats()
diff --git a/lambda-code/load-testing/tests/behaviours/api.py b/lambda-code/load-testing/tests/behaviours/api.py
@@ -20,20 +20,20 @@ def __init__(self, parent: HttpUser) -> None:
         self.form_decrypted_submissions = {}
         self.form_new_submissions = None
         self.headers = None
-        self.jwt_user = None
+        self.jwt_form = None
 
     def on_start(self) -> None:
-        self.jwt_user = JwtGenerator.generate(self.idp_url, self.private_api_key_user)
+        self.jwt_form = JwtGenerator.generate(self.idp_url, self.form_private_key)
         data = {
             "grant_type": "urn:ietf:params:oauth:grant-type:jwt-bearer",
-            "assertion": self.jwt_user,
+            "assertion": self.jwt_form,
             "scope": f"openid profile urn:zitadel:iam:org:project:id:{self.idp_project_id}:aud",
         }
         response = self.request_with_failure_check(
             "post", f"{self.idp_url}/oauth/v2/token", 200, data=data
         )
         self.headers = {
-            "Authorization": f"Bearer {response["access_token"]}",
+            "Authorization": f"Bearer {response['access_token']}",
             "Content-Type": "application/json",
         }
 
@@ -69,7 +69,7 @@ def get_submission_by_name(self) -> None:
         )
         encrypted_submission = EncryptedFormSubmission.from_json(response)
         decrypted_submission = FormSubmissionDecrypter.decrypt(
-            encrypted_submission, self.private_api_key_user
+            encrypted_submission, self.form_private_key
         )
         self.form_decrypted_submissions[submission["name"]] = json.loads(
             decrypted_submission

diff --git a/lambda-code/load-testing/tests/behaviours/idp.py b/lambda-code/load-testing/tests/behaviours/idp.py
@@ -11,18 +11,18 @@ class AccessTokenBehaviour(SequentialTaskSetWithFailure):
     def __init__(self, parent: HttpUser) -> None:
         super().__init__(parent)
         self.jwt_app = None
-        self.jwt_user = None
+        self.jwt_form = None
         self.access_token = None
 
     def on_start(self) -> None:
-        self.jwt_app = JwtGenerator.generate(self.idp_url, self.private_api_key_app)
-        self.jwt_user = JwtGenerator.generate(self.idp_url, self.private_api_key_user)
+        self.jwt_app = JwtGenerator.generate(self.idp_url, self.zitadel_app_private_key)
+        self.jwt_form = JwtGenerator.generate(self.idp_url, self.form_private_key)
 
     @task
     def request_access_token(self) -> None:
         data = {
             "grant_type": "urn:ietf:params:oauth:grant-type:jwt-bearer",
-            "assertion": self.jwt_user,
+            "assertion": self.jwt_form,
             "scope": f"openid profile urn:zitadel:iam:org:project:id:{self.idp_project_id}:aud",
         }
         response = self.request_with_failure_check(

diff --git a/lambda-code/load-testing/tests/behaviours/submit.py b/lambda-code/load-testing/tests/behaviours/submit.py
@@ -15,16 +15,16 @@ class FormSubmitBehaviour(SequentialTaskSetWithFailure):
     def __init__(self, parent: HttpUser) -> None:
         super().__init__(parent)
         self.access_token = None
-        self.jwt_user = None
+        self.jwt_form = None
         self.form_id = os.getenv("FORM_ID")
         self.form_template = None
         self.form_submission_generator = None
 
     def on_start(self) -> None:
-        self.jwt_user = JwtGenerator.generate(self.idp_url, self.private_api_key_user)
+        self.jwt_form = JwtGenerator.generate(self.idp_url, self.form_private_key)
         data = {
             "grant_type": "urn:ietf:params:oauth:grant-type:jwt-bearer",
-            "assertion": self.jwt_user,
+            "assertion": self.jwt_form,
             "scope": f"openid profile urn:zitadel:iam:org:project:id:{self.idp_project_id}:aud",
         }
         response = self.request_with_failure_check(

diff --git a/lambda-code/load-testing/tests/utils/task_set.py b/lambda-code/load-testing/tests/utils/task_set.py
@@ -16,11 +16,11 @@ def __init__(self, parent) -> None:
         self.api_url = f"{parsed_url.scheme}://api.{parsed_url.netloc}"
         self.idp_url = f"{parsed_url.scheme}://auth.{parsed_url.netloc}"
         self.idp_project_id = os.getenv("IDP_PROJECT_ID", "275372254274006635")
-        self.private_api_key_app = PrivateApiKey.from_json(
-            json.loads(os.getenv("PRIVATE_API_KEY_APP_JSON"))
+        self.form_private_key = PrivateApiKey.from_json(
+            json.loads(os.getenv("FORM_PRIVATE_KEY").replace('\n', '\\n'))
         )
-        self.private_api_key_user = PrivateApiKey.from_json(
-            json.loads(os.getenv("PRIVATE_API_KEY_USER_JSON"))
+        self.zitadel_app_private_key = PrivateApiKey.from_json(
+            json.loads(os.getenv("ZITADEL_APP_PRIVATE_KEY").replace('\n', '\\n'))
         )
 
     def request_with_failure_check(