[keymanager] add cloudbuild integration test flow

launcher/cloudbuild.yaml

@@ -220,6 +220,21 @@ steps:
       --substitutions _IMAGE_NAME=${OUTPUT_IMAGE_NAME},_IMAGE_PROJECT=${PROJECT_ID}
     exit
 
+- name: 'gcr.io/cloud-builders/gcloud'
+  id: KeymanagerTests
+  waitFor: ['DebugVgImageBuild']
+  env:
+  - 'OUTPUT_IMAGE_NAME=${_OUTPUT_IMAGE_PREFIX}-vg-debug-${_OUTPUT_IMAGE_SUFFIX}'
+  - 'PROJECT_ID=$PROJECT_ID'
+  script: |
+    #!/usr/bin/env bash
+
+    cd launcher/image/test
+    echo "running keymanager tests on ${OUTPUT_IMAGE_NAME}"
+    gcloud builds submit --config=test_keymanager_cloudbuild.yaml --region us-west1 \
+      --substitutions _IMAGE_NAME=${OUTPUT_IMAGE_NAME},_IMAGE_PROJECT=${PROJECT_ID}
+    exit
+
 - name: 'gcr.io/cloud-builders/gcloud'
   id: DebugImageTests
   waitFor: ['DebugImageBuild']

launcher/image/test/scripts/test_keymanager.sh

@@ -0,0 +1,20 @@
+#!/bin/bash
+set -euo pipefail
+source util/read_serial.sh
+
+# This test requires the workload to run and print
+# corresponding messages to the serial console.
+SERIAL_OUTPUT=$(read_serial $1 $2) 
+print_serial=false
+
+if echo "$SERIAL_OUTPUT" | grep -q "Success! Flow completed."; then
+    echo "- test keymanager"
+else
+    echo "FAILED: Could not find 'Success! Flow completed.' in the serial console"
+    echo "TEST FAILED. Keymanager flow was expected to pass validation." > /workspace/status.txt
+    print_serial=true
+fi
+
+if $print_serial; then
+    echo "$SERIAL_OUTPUT"
+fi

launcher/image/test/test_keymanager_cloudbuild.yaml

@@ -0,0 +1,46 @@
+# Test that the Keymanager service works for standard KMS-like flows.
+# This runs the test workload which calls generate, enumerate, decap,
+# and verifies a local encap.
+substitutions:
+  '_IMAGE_NAME': ''
+  '_IMAGE_PROJECT': ''
+  '_CLEANUP': 'true'
+  '_VM_NAME_PREFIX': 'cs-keymanager-test'
+  '_ZONE': 'asia-east1-a'
+  # Important notes: the actual real URL shouldn't be hardcoded since it requires building
+  # The run_cloudbuild.sh script or user's invocation will override _WORKLOAD_IMAGE
+  '_WORKLOAD_IMAGE': 'us-west1-docker.pkg.dev/confidential-space-images-dev/cs-integ-test-images/keymanager_workload:latest'
+steps:
+- name: 'gcr.io/cloud-builders/gcloud'
+  id: CreateVM
+  entrypoint: 'bash'
+  env:
+  - 'BUILD_ID=$BUILD_ID'
+  args: ['create_vm.sh','-i', '${_IMAGE_NAME}',
+          '-p', '${_IMAGE_PROJECT}',
+          '-m', 'tee-image-reference=${_WORKLOAD_IMAGE},tee-container-log-redirect=true',
+          '-n', '${_VM_NAME_PREFIX}-${BUILD_ID}',
+          '-z', '${_ZONE}',
+        ]
+- name: 'gcr.io/cloud-builders/gcloud'
+  id: TestKeymanager
+  entrypoint: 'bash'
+  args: ['scripts/test_keymanager.sh', '${_VM_NAME_PREFIX}-${BUILD_ID}', '${_ZONE}']
+  waitFor: ['CreateVM']
+- name: 'gcr.io/cloud-builders/gcloud'
+  id: CleanUp
+  entrypoint: 'bash'
+  env:
+  - 'CLEANUP=$_CLEANUP'
+  args: ['cleanup.sh', '${_VM_NAME_PREFIX}-${BUILD_ID}', '${_ZONE}']
+# Must come after cleanup.
+- name: 'gcr.io/cloud-builders/gcloud'
+  id: CheckFailure
+  entrypoint: 'bash'
+  env:
+  - 'BUILD_ID=$BUILD_ID'
+  args: ['check_failure.sh']
+
+options:
+  pool:
+    name: 'projects/confidential-space-images-dev/locations/us-west1/workerPools/cs-image-build-vpc'

launcher/image/testworkloads/keyclaims/Dockerfile

@@ -0,0 +1,13 @@
+# From current directory:
+# gcloud builds submit --tag us-west1-docker.pkg.dev/confidential-space-images-dev/cs-integ-test-images/keyclaims-workload:latest --project confidential-space-images-dev
+FROM python:3.11-slim
+
+WORKDIR /app
+
+RUN pip install --no-cache-dir requests requests-unixsocket
+
+COPY workload.py .
+
+RUN mkdir -p /run/container_launcher
+
+CMD ["python", "workload.py"]

launcher/image/testworkloads/keyclaims/workload.py

@@ -0,0 +1,73 @@
+import requests
+import requests_unixsocket
+import json
+import base64
+from urllib.parse import quote
+
+# Configure the session to use the UDS adapter
+session = requests_unixsocket.Session()
+
+# Define the socket path
+# Note: requests_unixsocket requires URL-encoding the path slashes
+# The KMS and AS are mounted at /run/container_launcher in the workload VM.
+kms_socket_path = "/run/container_launcher/kmaserver.sock"
+kms_encoded_path = quote(kms_socket_path, safe='')
+kms_base_url = f"http+unix://{kms_encoded_path}"
+as_socket_path = "/run/container_launcher/teeserver.sock"
+as_encoded_path = quote(as_socket_path, safe='')
+as_base_url = f"http+unix://{as_encoded_path}"
+
+def get_capabilities():
+    """Check supported algorithms."""
+    resp = session.get(f"{kms_base_url}/v1/capabilities")
+    resp.raise_for_status()
+    print("Capabilities:", json.dumps(resp.json(), indent=2))
+
+def generate_key():
+    """Generate a new KEM key."""
+    payload = {
+        "algorithm": {
+            "type": "kem",
+            "params": {
+                "kem_id": "DHKEM_X25519_HKDF_SHA256",
+                }
+            },
+            "lifespan": 3600
+    }
+    resp = session.post(f"{kms_base_url}/v1/keys:generate_key", json=payload)
+    resp.raise_for_status()
+    key_handle = resp.json()["key_handle"]
+    print(f"Generated Key: {key_handle['handle']}")
+    return key_handle
+
+def get_key_endorsement(challenge, key_handle):
+    """Generates a fresh key endorsement."""
+    challange_base64 = base64.b64encode(challenge).decode('utf-8')
+    payload = {
+        "challenge": challange_base64,
+        "key_handle": key_handle
+    }
+    resp = session.post(f"{as_base_url}/v1/keys:getEndorsement", json=payload)
+    print(f"resp: {resp}")
+    resp.raise_for_status()
+    return resp.json()
+
+
+
+if __name__ == "__main__":
+    try:
+        print("--- 1. Checking Capabilities ---")
+        get_capabilities()
+        print("\n--- 2. Generating Key ---")
+        handle = generate_key()
+        # Simulate some encapsulated key (32 bytes for X25519)
+        dummy_ciphertext = b'\x00' * 32
+        print("\n--- 3. Generating Key Endorsement ---")
+        challenge = b'\x00' * 32 
+
+        print(get_key_endorsement(challenge, handle))
+    except requests.exceptions.ConnectionError:
+        print(f"Error: Could not connect to socket at {kms_socket_path}.")
+        print("Ensure the Workload Services Daemon is running.")
+    except Exception as e:
+        print(f"An error occurred: {e}")

launcher/image/testworkloads/keymanager/Dockerfile

@@ -1,13 +1,18 @@
-# From current directory:
-# gcloud builds submit --tag us-west1-docker.pkg.dev/confidential-space-images-dev/cs-integ-test-images/kms-workload:latest --project confidential-space-images-dev
-FROM python:3.11-slim
+# Base image with Python support
+FROM python:3.10-slim
 
+# Set the working directory
 WORKDIR /app
 
-RUN pip install --no-cache-dir requests requests-unixsocket
+# Install the required Python packages
+RUN pip install requests requests-unixsocket pyhpke jsonschema
 
-COPY workload.py .
+# Copy the client script
+COPY workload.py /app/workload.py
 
-RUN mkdir -p /run/container_launcher
+LABEL "tee.launch_policy.allow_env_override"="ALLOWED_OVERRIDE"
+LABEL "tee.launch_policy.allow_cmd_override"="true"
+LABEL "tee.launch_policy.log_redirect"="always"
 
-CMD ["python", "workload.py"]
+# Set the entrypoint to explicitly use unbuffered output so logs appear instantly
+ENTRYPOINT ["python3", "-u", "/app/workload.py"]