Fix teeserver context reset issue

cloudbuild.yaml

@@ -228,7 +228,20 @@ steps:
     gcloud builds submit --config=test_memory_monitoring.yaml --region us-west1 \
       --substitutions _IMAGE_NAME=${OUTPUT_IMAGE_PREFIX}-hardened-${OUTPUT_IMAGE_SUFFIX},_IMAGE_PROJECT=${PROJECT_ID}
     exit
-
+- name: 'gcr.io/cloud-builders/gcloud'
+  id: ODAWithSignedContainerTest
+  waitFor: ['HardenedImageBuild']
+  env:
+  - 'OUTPUT_IMAGE_PREFIX=$_OUTPUT_IMAGE_PREFIX'
+  - 'OUTPUT_IMAGE_SUFFIX=$_OUTPUT_IMAGE_SUFFIX'
+  - 'PROJECT_ID=$PROJECT_ID'
+  script: |
+    #!/usr/bin/env bash
+    cd launcher/image/test
+    echo "running ODA and signed container tests on ${OUTPUT_IMAGE_PREFIX}-hardened-${OUTPUT_IMAGE_SUFFIX}"
+    gcloud builds submit --config=test_oda_with_signed_container.yaml --region us-west1 \
+      --substitutions _IMAGE_NAME=${OUTPUT_IMAGE_PREFIX}-hardened-${OUTPUT_IMAGE_SUFFIX},_IMAGE_PROJECT=${PROJECT_ID}
+    exit
 options:
   pool:
     name: 'projects/confidential-space-images-dev/locations/us-west1/workerPools/cs-image-build-vpc'

launcher/container_runner.go

@@ -502,7 +502,7 @@ func (r *ContainerRunner) Run(ctx context.Context) error {
 	// create and start the TEE server behind the experiment
 	if r.launchSpec.Experiments.EnableOnDemandAttestation {
 		r.logger.Println("EnableOnDemandAttestation is enabled: initializing TEE server.")
-		teeServer, err := teeserver.New(path.Join(launcherfile.HostTmpPath, teeServerSocket), r.attestAgent, r.logger)
+		teeServer, err := teeserver.New(ctx, path.Join(launcherfile.HostTmpPath, teeServerSocket), r.attestAgent, r.logger)
 		if err != nil {
 			return fmt.Errorf("failed to create the TEE server: %v", err)
 		}

launcher/image/test/scripts/test_launcher_workload_discover_signatures.sh

@@ -7,11 +7,12 @@ source util/read_serial.sh
 SERIAL_OUTPUT=$(read_serial $1 $2) 
 print_serial=false
 
-if echo $SERIAL_OUTPUT | grep -q 'Found container image signatures'
-then
-    echo "- container image signatures found"
+# Check how many times "Found container image signatures" is being logged.
+counts=$(echo $SERIAL_OUTPUT | grep -o 'Found container image signatures' | wc -l)
+if [ $counts -eq $3 ]; then
+    echo "- container image signatures found with expected counts: $3"
 else
-    echo "FAILED: container image signatures not found"
+    echo "FAILED: container image signatures want $3 counts, but got $counts"
     echo 'TEST FAILED.' > /workspace/status.txt
     print_serial=true
 fi

launcher/image/test/test_discover_signatures.yaml

@@ -36,7 +36,7 @@ steps:
 - name: 'gcr.io/cloud-builders/gcloud'
   id: BasicDiscoverSignaturesTest
   entrypoint: 'bash'
-  args: ['scripts/test_launcher_workload_discover_signatures.sh', '${_VM_NAME_PREFIX}-${BUILD_ID}', '${_ZONE}']
+  args: ['scripts/test_launcher_workload_discover_signatures.sh', '${_VM_NAME_PREFIX}-${BUILD_ID}', '${_ZONE}', '1']
 - name: 'gcr.io/cloud-builders/gcloud'
   id: CleanUp
   entrypoint: 'bash'

launcher/image/test/test_oda_with_signed_container.yaml

@@ -0,0 +1,75 @@
+substitutions:
+  '_IMAGE_NAME': ''
+  '_IMAGE_PROJECT': ''
+  '_CLEANUP': 'true'
+  '_VM_NAME_PREFIX': 'oda-signedcontainer'
+  '_ZONE': 'us-east1-b'
+  '_WORKLOAD_IMAGE': 'us-west1-docker.pkg.dev/confidential-space-images-dev/cs-integ-test-images/ipc/happypath:latest'
+  '_SIGNATURE_REPO': 'us-docker.pkg.dev/confidential-space-images-dev/cs-cosign-tests/hardened'
+
+steps:
+- name: 'gcr.io/projectsigstore/cosign:v2.2.0'
+  id: SignContainer
+  entrypoint: 'sh'
+  env:
+  - 'BUILD_ID=$BUILD_ID'
+  args:
+  - -c
+  - |
+    # Unpadded base64 encoding on the CloudKMS public key
+    pub=$(cosign public-key --key gcpkms://projects/confidential-space-images-dev/locations/global/keyRings/cosign-test/cryptoKeys/ecdsa/cryptoKeyVersions/1 | openssl base64)
+    pub=$(echo $pub | tr -d '[:space:]' | sed 's/[=]*$//')
+    # Use cosign sign
+    export COSIGN_REPOSITORY=${_SIGNATURE_REPO}
+    cosign sign --key gcpkms://projects/confidential-space-images-dev/locations/global/keyRings/cosign-test/cryptoKeys/ecdsa/cryptoKeyVersions/1 ${_WORKLOAD_IMAGE} -a dev.cosignproject.cosign/sigalg=ECDSA_P256_SHA256 -a dev.cosignproject.cosign/pub=$pub
+- 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,tee-signed-image-repos=${_SIGNATURE_REPO},tee-env-ALLOWED_OVERRIDE=overridden,tee-cmd=["newCmd"]',
+          '-n', '${_VM_NAME_PREFIX}-${BUILD_ID}',
+          '-z', '${_ZONE}',
+        ]
+- name: 'gcr.io/cloud-builders/gcloud'
+  id: TestCustomToken
+  entrypoint: 'bash'
+  args: ['scripts/test_custom_token.sh', "true", '${_VM_NAME_PREFIX}-${BUILD_ID}', '${_ZONE}']
+- name: 'gcr.io/cloud-builders/gcloud'
+  id: BasicDiscoverSignaturesTest
+  entrypoint: 'bash'
+  # Check how many times `Found container image signatures` is being logged. 
+  # Since signature discovery will occur on refresh the default token, and on teeserver receiving ODA requests, so the expected number should be 2.
+  args: ['scripts/test_launcher_workload_discover_signatures.sh', '${_VM_NAME_PREFIX}-${BUILD_ID}', '${_ZONE}', '2']
+- name: 'gcr.io/cloud-builders/gcloud'
+  id: CleanUp
+  entrypoint: 'bash'
+  env:
+  - 'CLEANUP=$_CLEANUP'
+  args: ['cleanup.sh', '${_VM_NAME_PREFIX}-${BUILD_ID}', '${_ZONE}']
+- name: 'gcr.io/cloud-builders/gcloud'
+  id: DeleteContainerSignatures
+  env:
+  - 'BUILD_ID=$BUILD_ID'
+  entrypoint: 'bash'
+  args:
+  - -c
+  - |
+    echo "Deleting container signatures..."
+    digest=$(gcloud artifacts docker images describe ${_WORKLOAD_IMAGE} --format 'value(image_summary.digest)')
+    tag=${digest/":"/"-"}.sig
+    # Delete container signature by its tag
+    gcloud artifacts docker images delete -q ${_SIGNATURE_REPO}:${tag}
+# 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/teeserver/tee_server.go

@@ -17,6 +17,7 @@ import (
 )
 
 type attestHandler struct {
+	ctx              context.Context
 	attestAgent      agent.AttestationAgent
 	defaultTokenFile string
 	logger           *log.Logger
@@ -36,7 +37,7 @@ type TeeServer struct {
 }
 
 // New takes in a socket and start to listen to it, and create a server
-func New(unixSock string, a agent.AttestationAgent, logger *log.Logger) (*TeeServer, error) {
+func New(ctx context.Context, unixSock string, a agent.AttestationAgent, logger *log.Logger) (*TeeServer, error) {
 	var err error
 	nl, err := net.Listen("unix", unixSock)
 	if err != nil {
@@ -47,6 +48,7 @@ func New(unixSock string, a agent.AttestationAgent, logger *log.Logger) (*TeeSer
 		netListener: nl,
 		server: &http.Server{
 			Handler: (&attestHandler{
+				ctx:              ctx,
 				attestAgent:      a,
 				defaultTokenFile: filepath.Join(launcherfile.HostTmpPath, launcherfile.AttestationVerifierTokenFilename),
 				logger:           logger,
@@ -76,7 +78,7 @@ func (a *attestHandler) getToken(w http.ResponseWriter, r *http.Request) {
 
 	switch r.Method {
 	case "GET":
-		// this could call Attest(context.Background()) directly later.
+		// this could call Attest(ctx) directly later.
 		data, err := os.ReadFile(a.defaultTokenFile)
 
 		if err != nil {
@@ -113,7 +115,7 @@ func (a *attestHandler) getToken(w http.ResponseWriter, r *http.Request) {
 			return
 		}
 
-		tok, err := a.attestAgent.Attest(context.Background(),
+		tok, err := a.attestAgent.Attest(a.ctx,
 			agent.AttestAgentOpts{
 				Aud:       tokenReq.Audience,
 				Nonces:    tokenReq.Nonces,