Add unit tests for IPC server

Add image test

cloudbuild.yaml

@@ -80,6 +80,22 @@ steps:
       --substitutions _IMAGE_NAME=${OUTPUT_IMAGE_PREFIX}-debug-${OUTPUT_IMAGE_SUFFIX},_IMAGE_PROJECT=${PROJECT_ID}
     exit
 
+- name: 'gcr.io/cloud-builders/gcloud'
+  id: HttpServerTests
+  waitFor: ['DebugImageBuild']
+  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 http server tests on ${OUTPUT_IMAGE_PREFIX}-debug-${OUTPUT_IMAGE_SUFFIX}"
+    gcloud builds submit --config=test_http_server.yaml --region us-west1 \
+      --substitutions _IMAGE_NAME=${OUTPUT_IMAGE_PREFIX}-debug-${OUTPUT_IMAGE_SUFFIX},_IMAGE_PROJECT=${PROJECT_ID}
+    exit
+
 - name: 'gcr.io/cloud-builders/gcloud'
   id: DebugImageTests
   waitFor: ['DebugImageBuild']

launcher/image/test/scripts/test_custom_token.sh

@@ -0,0 +1,21 @@
+ #!/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 $2 $3) 
+print_serial=false
+
+if echo $SERIAL_OUTPUT | grep -q "Token valid: $1"
+then
+    echo "- test custom token"
+else
+    echo "FAILED: Could not find 'Token valid: $1' in the serial console"
+    echo "TEST FAILED. Token was expected to pass validation." > /workspace/status.txt
+    print_serial=true
+fi
+
+if $print_serial; then
+    echo $SERIAL_OUTPUT
+fi

launcher/image/test/test_http_server.yaml

@@ -0,0 +1,42 @@
+# Test that the TEE server can accept requests for and serve custom tokens.
+# This is a happy path test.
+substitutions:
+  '_IMAGE_NAME': ''
+  '_IMAGE_PROJECT': ''
+  '_CLEANUP': 'true'
+  '_VM_NAME_PREFIX': 'cs-http-server-test'
+  '_ZONE': 'us-central1-a'
+  '_WORKLOAD_IMAGE': 'us-west1-docker.pkg.dev/confidential-space-images-dev/cs-integ-test-images/ipc/happypath: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: TestCustomToken
+  entrypoint: 'bash'
+  args: ['scripts/test_custom_token.sh', "true", '${_VM_NAME_PREFIX}-${BUILD_ID}', '${_ZONE}']
+- 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/customtoken/happypath/Dockerfile

@@ -0,0 +1,17 @@
+# From current directory:
+# GOOS=linux GOARCH=amd64 CGO_ENABLED=0 go build -o main .
+# gcloud builds submit --tag us-west1-docker.pkg.dev/confidential-space-images-dev/cs-integ-test-images/ipc/happypath:latest
+FROM alpine
+
+COPY main /
+
+ENV env_bar="val_bar"
+
+LABEL "tee.launch_policy.allow_env_override"="ALLOWED_OVERRIDE"
+LABEL "tee.launch_policy.allow_cmd_override"="true"
+LABEL "tee.launch_policy.log_redirect"="always"
+
+ENTRYPOINT ["/main"]
+
+# Can be overridden because of the launch policy.
+CMD ["arg_foo"]

launcher/image/testworkloads/customtoken/happypath/main.go

@@ -0,0 +1,232 @@
+// package main is a binary that will print out the validation status of a custom attestation token.
+package main
+
+import (
+	"context"
+	"crypto/rsa"
+	"encoding/base64"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"io"
+	"math/big"
+	"net"
+	"net/http"
+	"strings"
+
+	"github.com/golang-jwt/jwt/v4"
+)
+
+// BUILD:
+// CGO_ENABLED=0 go build
+// docker build -t <your_repo> .
+
+const (
+	socketPath     = "/run/container_launcher/teeserver.sock"
+	expectedIssuer = "https://confidentialcomputing.googleapis.com"
+	wellKnownPath  = "/.well-known/openid-configuration"
+)
+
+type jwksFile struct {
+	Keys []jwk `json:"keys"`
+}
+
+type jwk struct {
+	N   string `json:"n"`   // "nMMTBwJ7H6Id8zUCZd-L7uoNyz9b7lvoyse9izD9l2rtOhWLWbiG-7pKeYJyHeEpilHP4KdQMfUo8JCwhd-OMW0be_XtEu3jXEFjuq2YnPSPFk326eTfENtUc6qJohyMnfKkcOcY_kTE11jM81-fsqtBKjO_KiSkcmAO4wJJb8pHOjue3JCP09ZANL1uN4TuxbM2ibcyf25ODt3WQn54SRQTV0wn098Y5VDU-dzyeKYBNfL14iP0LiXBRfHd4YtEaGV9SBUuVhXdhx1eF0efztCNNz0GSLS2AEPLQduVuFoUImP4s51YdO9TPeeQ3hI8aGpOdC0syxmZ7LsL0rHE1Q",
+	E   string `json:"e"`   // "AQAB" or 65537 as an int
+	Kid string `json:"kid"` // "1f12fa916c3a0ef585894b4b420ad17dc9d6cdf5",
+
+	// Unused fields:
+	// Alg string `json:"alg"` // "RS256",
+	// Kty string `json:"kty"` // "RSA",
+	// Use string `json:"use"` // "sig",
+}
+
+type wellKnown struct {
+	JwksURI string `json:"jwks_uri"` // "https://www.googleapis.com/service_accounts/v1/metadata/jwk/[email protected]"
+
+	// Unused fields:
+	// Iss                                   string `json:"issuer"`                                // "https://confidentialcomputing.googleapis.com"
+	// Subject_types_supported               string `json:"subject_types_supported"`               // [ "public" ]
+	// Response_types_supported              string `json:"response_types_supported"`              // [ "id_token" ]
+	// Claims_supported                      string `json:"claims_supported"`                      // [ "sub", "aud", "exp", "iat", "iss", "jti", "nbf", "dbgstat", "eat_nonce", "google_service_accounts", "hwmodel", "oemid", "secboot", "submods", "swname", "swversion" ]
+	// Id_token_signing_alg_values_supported string `json:"id_token_signing_alg_values_supported"` // [ "RS256" ]
+	// Scopes_supported                      string `json:"scopes_supported"`                      // [ "openid" ]
+}
+
+func getCustomTokenBytes(body string) ([]byte, error) {
+	httpClient := http.Client{
+		Transport: &http.Transport{
+			// Set the DialContext field to a function that creates
+			// a new network connection to a Unix domain socket
+			DialContext: func(_ context.Context, _, _ string) (net.Conn, error) {
+				return net.Dial("unix", socketPath)
+			},
+		},
+	}
+
+	// Get the token from the IPC endpoint
+	url := "http://localhost/v1/token"
+
+	resp, err := httpClient.Post(url, "application/json", strings.NewReader(body))
+	if err != nil {
+		return nil, fmt.Errorf("failed to get raw custom token response: %w", err)
+	}
+	tokenbytes, err := io.ReadAll(resp.Body)
+	if err != nil {
+		return nil, fmt.Errorf("failed to read custom token body: %w", err)
+	}
+	fmt.Println(string(tokenbytes))
+	return tokenbytes, nil
+}
+
+func getWellKnownFile() (wellKnown, error) {
+	httpClient := http.Client{}
+	resp, err := httpClient.Get(expectedIssuer + wellKnownPath)
+	if err != nil {
+		return wellKnown{}, fmt.Errorf("failed to get raw .well-known response: %w", err)
+	}
+
+	wellKnownJSON, err := io.ReadAll(resp.Body)
+	if err != nil {
+		return wellKnown{}, fmt.Errorf("failed to read .well-known response: %w", err)
+	}
+
+	wk := wellKnown{}
+	json.Unmarshal(wellKnownJSON, &wk)
+	return wk, nil
+}
+
+func getJWKFile() (jwksFile, error) {
+	wk, err := getWellKnownFile()
+	if err != nil {
+		return jwksFile{}, fmt.Errorf("failed to get .well-known json: %w", err)
+	}
+
+	// Get JWK URI from .wellknown
+	uri := wk.JwksURI
+	fmt.Printf("jwks URI: %v\n", uri)
+
+	httpClient := http.Client{}
+	resp, err := httpClient.Get(uri)
+	if err != nil {
+		return jwksFile{}, fmt.Errorf("failed to get raw JWK response: %w", err)
+	}
+
+	jwkbytes, err := io.ReadAll(resp.Body)
+	if err != nil {
+		return jwksFile{}, fmt.Errorf("failed to read JWK body: %w", err)
+	}
+
+	file := jwksFile{}
+	err = json.Unmarshal(jwkbytes, &file)
+	if err != nil {
+		return jwksFile{}, fmt.Errorf("failed to unmarshall JWK content: %w", err)
+	}
+
+	return file, nil
+}
+
+// N and E are 'base54urlUInt' encoded: https://www.rfc-editor.org/rfc/rfc7518#section-6.3
+func base64urlUIntDecode(s string) (*big.Int, error) {
+	b, err := base64.RawURLEncoding.DecodeString(s)
+	if err != nil {
+		return nil, err
+	}
+	z := new(big.Int)
+	z.SetBytes(b)
+	return z, nil
+}
+
+func getRSAPublicKeyFromJWKsFile(t *jwt.Token) (any, error) {
+	keysfile, err := getJWKFile()
+	if err != nil {
+		return nil, fmt.Errorf("failed to fetch the JWK file: %w", err)
+	}
+
+	// Talk about key rotation + how we know about which one is good
+	// Token is passed in with the kid of the key used to sign.
+	kid := t.Header["kid"]
+	for _, key := range keysfile.Keys {
+		if key.Kid != kid {
+			continue // Select the key used for signing
+		}
+
+		n, err := base64urlUIntDecode(key.N)
+		if err != nil {
+			return nil, fmt.Errorf("failed to decode key.N %w", err)
+		}
+		e, err := base64urlUIntDecode(key.E)
+		if err != nil {
+			return nil, fmt.Errorf("failed to decode key.E %w", err)
+		}
+
+		// The parser expects an rsa.PublicKey: https://github.com/golang-jwt/jwt/blob/main/rsa.go#L53
+		return &rsa.PublicKey{
+			N: n,
+			E: int(e.Int64()),
+		}, nil
+	}
+
+	return nil, fmt.Errorf("failed to find key with kid '%v' from well-known endpoint", kid)
+}
+
+func decodeAndValidateToken(tokenBytes []byte, keyFunc func(t *jwt.Token) (any, error)) (*jwt.Token, error) {
+	var err error
+	fmt.Println("Unmarshalling token and checking its validity...")
+	token, err := jwt.NewParser().Parse(string(tokenBytes), keyFunc)
+
+	fmt.Printf("Token valid: %v", token.Valid)
+	if token.Valid {
+		return token, nil
+	}
+	if ve, ok := err.(*jwt.ValidationError); ok {
+		if ve.Errors&jwt.ValidationErrorMalformed != 0 {
+			return nil, fmt.Errorf("token format invalid. Please contact the Confidential Space team for assistance")
+		}
+		if ve.Errors&(jwt.ValidationErrorNotValidYet) != 0 {
+			// If device time is not synchronized with the Attestation Service you may need to account for that here.
+			return nil, errors.New("token is not active yet")
+		}
+		if ve.Errors&(jwt.ValidationErrorExpired) != 0 {
+			return nil, fmt.Errorf("token is expired")
+		}
+		return nil, fmt.Errorf("unknown validation error: %v", err)
+	}
+
+	return nil, fmt.Errorf("couldn't handle this token or couldn't read a validation error: %v", err)
+}
+
+func main() {
+	// Format token request
+	body := `{
+        "audience": "<YOURAUDIENCE>",
+        "nonces": ["thisIsAcustomNonce", "thisIsAMuchLongerCustomNonceWithPaddingFor74Bytes0000000000000000000000000"],
+		"tokenType": "OIDC",
+    }`
+
+	// The following code could be run in a Confidential Space workload container to generate a
+	// custom attestation intended to be sent to a remote party for verification.
+	tokenbytes, err := getCustomTokenBytes(body)
+	if err != nil {
+		panic(err)
+	}
+
+	// Write a method to return a public key from the well-known endpoint
+	keyFunc := getRSAPublicKeyFromJWKsFile
+
+	// The following code could be run by a remote party (not necessarily in a
+	// Confidential Space workload) in order to verify properties of the original
+	// Confidential Space workload that generated the attestation.
+	token, err := decodeAndValidateToken(tokenbytes, keyFunc)
+	if err != nil {
+		panic(err)
+	}
+
+	claimsString, err := json.MarshalIndent(token.Claims, "", "  ")
+	if err != nil {
+		panic(err)
+	}
+
+	fmt.Println(string(claimsString))
+}

launcher/teeserver/tee_server.go

@@ -23,8 +23,9 @@ type attestHandler struct {
 }
 
 type customTokenRequest struct {
-	Audience string   `json:"audience"`
-	Nonces   []string `json:"nonces"`
+	Audience  string   `json:"audience"`
+	Nonces    []string `json:"nonces"`
+	TokenType string   `json:"tokenType"`
 }
 
 // TeeServer is a server that can be called from a container through a unix
@@ -106,6 +107,12 @@ func (a *attestHandler) getToken(w http.ResponseWriter, r *http.Request) {
 			return
 		}
 
+		if tokenReq.TokenType == "" {
+			w.WriteHeader(http.StatusBadRequest)
+			w.Write([]byte("tokenType is a required parameter"))
+			return
+		}
+
 		tok, err := a.attestAgent.Attest(context.Background(),
 			agent.AttestAgentOpts{
 				Aud:    tokenReq.Audience,