feat: add SSM parameters for the load tests

.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

.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

.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

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
+}

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}",
+    ]
+  }
+}

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
+}

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 {

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"]

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
+```

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()

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

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(

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(

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(