Skip to content

Commit 64f2a3f

Browse files
atulpatildbzatulpatildbz
authored
[keymanager] add cloudbuild integration test flow (google#701)
This adds the automated integration tests for the keymanager service. - Adds a test workload container image that runs the end-to-end Python client test - Adds a cloudbuild yaml test definition to spin up a Keymanager VM with the workload - Hooks the test into the main launcher cloudbuild.yaml workflow
1 parent c16668d commit 64f2a3f

File tree

7 files changed

+460
-99
lines changed

7 files changed

+460
-99
lines changed

‎launcher/cloudbuild.yaml‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -220,6 +220,21 @@ steps:
220220
--substitutions _IMAGE_NAME=${OUTPUT_IMAGE_NAME},_IMAGE_PROJECT=${PROJECT_ID}
221221
exit
222222
223+
- name: 'gcr.io/cloud-builders/gcloud'
224+
id: KeymanagerTests
225+
waitFor: ['DebugVgImageBuild']
226+
env:
227+
- 'OUTPUT_IMAGE_NAME=${_OUTPUT_IMAGE_PREFIX}-vg-debug-${_OUTPUT_IMAGE_SUFFIX}'
228+
- 'PROJECT_ID=$PROJECT_ID'
229+
script: |
230+
#!/usr/bin/env bash
231+
232+
cd launcher/image/test
233+
echo "running keymanager tests on ${OUTPUT_IMAGE_NAME}"
234+
gcloud builds submit --config=test_keymanager_cloudbuild.yaml --region us-west1 \
235+
--substitutions _IMAGE_NAME=${OUTPUT_IMAGE_NAME},_IMAGE_PROJECT=${PROJECT_ID}
236+
exit
237+
223238
- name: 'gcr.io/cloud-builders/gcloud'
224239
id: DebugImageTests
225240
waitFor: ['DebugImageBuild']

‎launcher/image/test/scripts/test_keymanager.sh‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
#!/bin/bash
2+
set -euo pipefail
3+
source util/read_serial.sh
4+
5+
# This test requires the workload to run and print
6+
# corresponding messages to the serial console.
7+
SERIAL_OUTPUT=$(read_serial $1 $2)
8+
print_serial=false
9+
10+
if echo "$SERIAL_OUTPUT" | grep -q "Success! Flow completed."; then
11+
echo "- test keymanager"
12+
else
13+
echo "FAILED: Could not find 'Success! Flow completed.' in the serial console"
14+
echo "TEST FAILED. Keymanager flow was expected to pass validation." > /workspace/status.txt
15+
print_serial=true
16+
fi
17+
18+
if $print_serial; then
19+
echo "$SERIAL_OUTPUT"
20+
fi

‎launcher/image/test/test_keymanager_cloudbuild.yaml‎

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,48 @@
1+
# Test that the Keymanager service works for standard KMS-like flows.
2+
# This runs the test workload which calls generate, enumerate, decap,
3+
# and verifies a local encap.
4+
substitutions:
5+
'_IMAGE_NAME': ''
6+
'_IMAGE_PROJECT': ''
7+
'_CLEANUP': 'true'
8+
'_VM_NAME_PREFIX': 'cs-keymanager-test'
9+
'_ZONE': 'asia-east1-a'
10+
# Important notes: the actual real URL shouldn't be hardcoded since it requires building
11+
# The run_cloudbuild.sh script or user's invocation will override _WORKLOAD_IMAGE
12+
# This workload image is built from launcher/image/testworkloads/keymanager
13+
# Registry reference: us-west1-docker.pkg.dev/confidential-space-images-dev/cs-integ-test-images/keymanager_workload:latest
14+
'_WORKLOAD_IMAGE': 'us-west1-docker.pkg.dev/confidential-space-images-dev/cs-integ-test-images/keymanager_workload:latest'
15+
steps:
16+
- name: 'gcr.io/cloud-builders/gcloud'
17+
id: CreateVM
18+
entrypoint: 'bash'
19+
env:
20+
- 'BUILD_ID=$BUILD_ID'
21+
args: ['create_vm.sh','-i', '${_IMAGE_NAME}',
22+
'-p', '${_IMAGE_PROJECT}',
23+
'-m', 'tee-image-reference=${_WORKLOAD_IMAGE},tee-container-log-redirect=true',
24+
'-n', '${_VM_NAME_PREFIX}-${BUILD_ID}',
25+
'-z', '${_ZONE}',
26+
]
27+
- name: 'gcr.io/cloud-builders/gcloud'
28+
id: TestKeymanager
29+
entrypoint: 'bash'
30+
args: ['scripts/test_keymanager.sh', '${_VM_NAME_PREFIX}-${BUILD_ID}', '${_ZONE}']
31+
waitFor: ['CreateVM']
32+
- name: 'gcr.io/cloud-builders/gcloud'
33+
id: CleanUp
34+
entrypoint: 'bash'
35+
env:
36+
- 'CLEANUP=$_CLEANUP'
37+
args: ['cleanup.sh', '${_VM_NAME_PREFIX}-${BUILD_ID}', '${_ZONE}']
38+
# Must come after cleanup.
39+
- name: 'gcr.io/cloud-builders/gcloud'
40+
id: CheckFailure
41+
entrypoint: 'bash'
42+
env:
43+
- 'BUILD_ID=$BUILD_ID'
44+
args: ['check_failure.sh']
45+
46+
options:
47+
pool:
48+
name: 'projects/confidential-space-images-dev/locations/us-west1/workerPools/cs-image-build-vpc'

‎launcher/image/testworkloads/keyclaims/Dockerfile‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
# From current directory:
2+
# 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
3+
FROM python:3.11-slim
4+
5+
WORKDIR /app
6+
7+
RUN pip install --no-cache-dir requests requests-unixsocket
8+
9+
COPY workload.py .
10+
11+
RUN mkdir -p /run/container_launcher
12+
13+
CMD ["python", "workload.py"]

‎launcher/image/testworkloads/keyclaims/workload.py‎

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,73 @@
1+
import requests
2+
import requests_unixsocket
3+
import json
4+
import base64
5+
from urllib.parse import quote
6+
7+
# Configure the session to use the UDS adapter
8+
session = requests_unixsocket.Session()
9+
10+
# Define the socket path
11+
# Note: requests_unixsocket requires URL-encoding the path slashes
12+
# The KMS and AS are mounted at /run/container_launcher in the workload VM.
13+
kms_socket_path = "/run/container_launcher/kmaserver.sock"
14+
kms_encoded_path = quote(kms_socket_path, safe='')
15+
kms_base_url = f"http+unix://{kms_encoded_path}"
16+
as_socket_path = "/run/container_launcher/teeserver.sock"
17+
as_encoded_path = quote(as_socket_path, safe='')
18+
as_base_url = f"http+unix://{as_encoded_path}"
19+
20+
def get_capabilities():
21+
"""Check supported algorithms."""
22+
resp = session.get(f"{kms_base_url}/v1/capabilities")
23+
resp.raise_for_status()
24+
print("Capabilities:", json.dumps(resp.json(), indent=2))
25+
26+
def generate_key():
27+
"""Generate a new KEM key."""
28+
payload = {
29+
"algorithm": {
30+
"type": "kem",
31+
"params": {
32+
"kem_id": "DHKEM_X25519_HKDF_SHA256",
33+
}
34+
},
35+
"lifespan": 3600
36+
}
37+
resp = session.post(f"{kms_base_url}/v1/keys:generate_key", json=payload)
38+
resp.raise_for_status()
39+
key_handle = resp.json()["key_handle"]
40+
print(f"Generated Key: {key_handle['handle']}")
41+
return key_handle
42+
43+
def get_key_endorsement(challenge, key_handle):
44+
"""Generates a fresh key endorsement."""
45+
challange_base64 = base64.b64encode(challenge).decode('utf-8')
46+
payload = {
47+
"challenge": challange_base64,
48+
"key_handle": key_handle
49+
}
50+
resp = session.post(f"{as_base_url}/v1/keys:getEndorsement", json=payload)
51+
print(f"resp: {resp}")
52+
resp.raise_for_status()
53+
return resp.json()
54+
55+
56+
57+
if __name__ == "__main__":
58+
try:
59+
print("--- 1. Checking Capabilities ---")
60+
get_capabilities()
61+
print("\n--- 2. Generating Key ---")
62+
handle = generate_key()
63+
# Simulate some encapsulated key (32 bytes for X25519)
64+
dummy_ciphertext = b'\x00' * 32
65+
print("\n--- 3. Generating Key Endorsement ---")
66+
challenge = b'\x00' * 32
67+
68+
print(get_key_endorsement(challenge, handle))
69+
except requests.exceptions.ConnectionError:
70+
print(f"Error: Could not connect to socket at {kms_socket_path}.")
71+
print("Ensure the Workload Services Daemon is running.")
72+
except Exception as e:
73+
print(f"An error occurred: {e}")

‎launcher/image/testworkloads/keymanager/Dockerfile‎

Lines changed: 14 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,21 @@
11
# From current directory:
2-
# 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
3-
FROM python:3.11-slim
2+
# gcloud builds submit --tag us-west1-docker.pkg.dev/confidential-space-images-dev/cs-integ-test-images/keymanager_workload:latest --project confidential-space-images-dev
43

4+
# Base image with Python support
5+
FROM python:3.10-slim
6+
7+
# Set the working directory
58
WORKDIR /app
69

7-
RUN pip install --no-cache-dir requests requests-unixsocket
10+
# Install the required Python packages
11+
RUN pip install requests requests-unixsocket pyhpke jsonschema
812

9-
COPY workload.py .
13+
# Copy the client script
14+
COPY workload.py /app/workload.py
1015

11-
RUN mkdir -p /run/container_launcher
16+
LABEL "tee.launch_policy.allow_env_override"="ALLOWED_OVERRIDE"
17+
LABEL "tee.launch_policy.allow_cmd_override"="true"
18+
LABEL "tee.launch_policy.log_redirect"="always"
1219

13-
CMD ["python", "workload.py"]
20+
# Set the entrypoint to explicitly use unbuffered output so logs appear instantly
21+
ENTRYPOINT ["python3", "-u", "/app/workload.py"]

0 commit comments

Comments
 (0)