Replace teeserver binary with mock_wsd for mocking WorkloadAttestationService.

Author: yanchengchen7

cmd/mock_wsd/README.md

@@ -0,0 +1,38 @@
+# mock_wsd
+
+`mock_wsd` is a standalone binary that mocks the `WorkloadAttestationService` specifically for the `GetKeyEndorsement` endpoint. 
+
+The binary exposes UDS located at `/run/workload_attestation.sock`. It serves RESTful JSON API that answers requests formatted as `GetKeyEndorsementRequest` and returns `GetKeyEndorsementResponse`.
+
+Internally, it spins up the `teeserver` logic and accesses the `AttestationAgent` directly to fetch a standalone VMAttestation, wraps it in the `GetKeyEndorsementResponse` struct, and sets the label to `WORKLOAD_ATTESTATION`.
+
+> **Note on Attestation Evidence**
+> By default, `mock_wsd` is configured with `launchSpec.Experiments.EnableAttestationEvidence = false`, flip this value to `true` before compiling.
+
+### 1. Build the binary
+
+From the root of `go-tpm-tools`:
+
+```bash
+go build -o mock_wsd_bin ./cmd/mock_wsd
+```
+
+### 2. Start the `mock_wsd`
+
+(Requires root privileges to listen on `/run/workload_attestation.sock` and access `/dev/tpmrm0`)
+
+```bash
+sudo ./mock_wsd_bin
+```
+
+### 3. Query `mock_wsd`
+
+Once the server is running, you can query `mock_wsd` from another terminal. You can write the output to a JSON file (e.g., `evidence.json`) as follows:
+```bash
+curl --unix-socket /run/workload_attestation.sock \
+	-H "Content-Type: application/json" \
+	-d '{"challenge": "Y2hhbGxlbmdl", "key_handle": {"handle": "some_handle"}}' \
+	http://localhost/v1/workload/attestation/key_endorsement | jq . > evidence.json
+```
+
+This will save the JSON API response (containing the nested `KeyEndorsement` -> `VmProtectedKeyEndorsement` -> `KeyAttestation` -> `VMAttestation`) into `evidence.json`.

cmd/mock_wsd/main.go

@@ -0,0 +1,234 @@
+package main
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"log"
+	"net"
+	"net/http"
+	"os"
+	"os/signal"
+	"strings"
+	"syscall"
+
+	clogging "cloud.google.com/go/logging"
+	"github.com/google/go-tpm-tools/client"
+	"github.com/google/go-tpm-tools/launcher/agent"
+	"github.com/google/go-tpm-tools/launcher/spec"
+	"github.com/google/go-tpm-tools/launcher/teeserver"
+	"github.com/google/go-tpm-tools/launcher/teeserver/models"
+	"github.com/google/go-tpm-tools/verifier/util"
+)
+
+const (
+	teeserverSocketPath = "/run/container_launcher/teeserver.sock"
+	mockWsdSocketPath   = "/run/workload_attestation.sock"
+)
+
+type cmdLogger struct {
+	*log.Logger
+}
+
+func (l *cmdLogger) Log(severity clogging.Severity, msg string, args ...any) {
+	l.Printf("%v: %s %v\n", severity, msg, args)
+}
+
+func (l *cmdLogger) Info(msg string, args ...any) {
+	l.Printf("INFO: %s %v\n", msg, args)
+}
+
+func (l *cmdLogger) Warn(msg string, args ...any) {
+	l.Printf("WARN: %s %v\n", msg, args)
+}
+
+func (l *cmdLogger) Error(msg string, args ...any) {
+	l.Printf("ERROR: %s %v\n", msg, args)
+}
+
+func (l *cmdLogger) SerialConsoleFile() *os.File {
+	return nil
+}
+
+func (l *cmdLogger) Close() {
+}
+
+type KeyHandle struct {
+	Handle string `json:"handle"`
+}
+
+type GetKeyEndorsementRequest struct {
+	Challenge []byte    `json:"challenge"`
+	KeyHandle KeyHandle `json:"key_handle"`
+}
+
+type GetKeyEndorsementResponse struct {
+	Endorsement KeyEndorsement `json:"endorsement"`
+}
+
+type KeyEndorsement struct {
+	VmProtectedKeyEndorsement VmProtectedKeyEndorsement `json:"vm_protected_key_endorsement"`
+}
+
+type VmProtectedKeyEndorsement struct {
+	BindingKeyAttestation   *KeyAttestation `json:"binding_key_attestation,omitempty"`
+	ProtectedKeyAttestation *KeyAttestation `json:"protected_key_attestation,omitempty"`
+}
+
+type KeyAttestation struct {
+	Attestation *models.VMAttestation `json:"attestation"`
+}
+
+func main() {
+	if os.Getuid() != 0 {
+		log.Println("Warning: mock_wsd usually requires root privileges to create sockets in /run")
+	}
+
+	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
+	defer stop()
+
+	logger := &cmdLogger{log.New(os.Stdout, "mock_wsd ", log.LstdFlags)}
+
+	// 1. Initialize Attestation Agent for teeserver
+	launchSpec := spec.LaunchSpec{}
+	launchSpec.Experiments.EnableAttestationEvidence = true
+
+	vTPM, err := os.OpenFile("/dev/tpmrm0", os.O_RDWR, 0)
+	if err != nil {
+		logger.Error(fmt.Sprintf("Failed to open vTPM: %v", err))
+	} else {
+		defer vTPM.Close()
+	}
+	var tpmCloser io.ReadWriteCloser
+	if vTPM != nil {
+		tpmCloser = vTPM
+	}
+
+	var akFetcher util.TpmKeyFetcher
+	if tpmCloser != nil {
+		akFetcher = client.GceAttestationKeyECC
+	} else {
+		akFetcher = func(_ io.ReadWriter) (*client.Key, error) {
+			return nil, fmt.Errorf("no vTPM available")
+		}
+	}
+
+	attestAgent, err := agent.CreateAttestationAgent(
+		tpmCloser, akFetcher, nil, nil, nil, launchSpec, logger,
+	)
+	if err != nil {
+		logger.Error(fmt.Sprintf("failed to create attestation agent: %v", err))
+		os.Exit(1)
+	}
+	defer attestAgent.Close()
+
+	// 2. Start teeserver
+	if err := os.MkdirAll("/run/container_launcher", 0755); err != nil { logger.Error(fmt.Sprintf("failed to create directory /run/container_launcher: %v", err)); os.Exit(1) }
+	if err := os.RemoveAll(teeserverSocketPath); err != nil {
+		logger.Error(fmt.Sprintf("Failed to remove existing socket %s: %v", teeserverSocketPath, err))
+	}
+
+	clients := teeserver.AttestClients{GCA: nil, ITA: nil}
+	tServer, err := teeserver.New(ctx, teeserverSocketPath, attestAgent, logger, launchSpec, clients)
+	if err != nil {
+		logger.Error(fmt.Sprintf("failed to create tee server: %v", err))
+		os.Exit(1)
+	}
+
+	if err := os.Chmod(teeserverSocketPath, 0777); err != nil {
+		logger.Warn(fmt.Sprintf("failed to chmod socket %s: %v", teeserverSocketPath, err))
+	}
+
+	logger.Info("Starting TEE Server", "socket", teeserverSocketPath)
+	errChan := make(chan error, 2)
+	go func() {
+		errChan <- tServer.Serve()
+	}()
+
+	// 3. Start mock_wsd server
+	mux := http.NewServeMux()
+	// Pass the attestAgent directly to avoid the loopback HTTP call
+	mux.HandleFunc("/v1/workload/attestation/key_endorsement", func(w http.ResponseWriter, r *http.Request) {
+		handleGetKeyEndorsement(w, r, attestAgent, logger)
+	})
+
+	if err := os.RemoveAll(mockWsdSocketPath); err != nil {
+		logger.Error(fmt.Sprintf("Failed to remove existing socket %s: %v", mockWsdSocketPath, err))
+	}
+
+	listener, err := net.Listen("unix", mockWsdSocketPath)
+	if err != nil {
+		logger.Error(fmt.Sprintf("Failed to listen on %s: %v", mockWsdSocketPath, err))
+		os.Exit(1)
+	}
+	defer listener.Close()
+
+	if err := os.Chmod(mockWsdSocketPath, 0777); err != nil {
+		logger.Warn(fmt.Sprintf("failed to chmod socket %s: %v", mockWsdSocketPath, err))
+	}
+
+	logger.Info("Starting mock_wsd server attached to TEE Server", "socket", mockWsdSocketPath)
+	go func() {
+		errChan <- http.Serve(listener, mux)
+	}()
+
+	// 4. Wait for termination
+	select {
+	case err := <-errChan:
+		if err != nil {
+			logger.Error(fmt.Sprintf("server error: %v", err))
+			os.Exit(1)
+		}
+	case <-ctx.Done():
+		logger.Info("Shutting down servers")
+		if err := tServer.Shutdown(context.Background()); err != nil {
+			if !strings.Contains(err.Error(), "use of closed network connection") {
+				logger.Error(fmt.Sprintf("failed to shutdown tee server: %v", err))
+			}
+		}
+	}
+}
+
+func handleGetKeyEndorsement(w http.ResponseWriter, r *http.Request, attestAgent agent.AttestationAgent, logger *cmdLogger) {
+	if r.Method != http.MethodPost {
+		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
+		return
+	}
+
+	var req GetKeyEndorsementRequest
+	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
+		http.Error(w, fmt.Sprintf("Failed to decode request: %v", err), http.StatusBadRequest)
+		return
+	}
+
+	if len(req.Challenge) == 0 {
+		http.Error(w, "challenge is required", http.StatusBadRequest)
+		return
+	}
+
+	// Call the generic AttestationEvidence function directly from the agent
+	attestation, err := attestAgent.AttestationEvidence(r.Context(), req.Challenge, nil)
+	if err != nil {
+		logger.Error(fmt.Sprintf("Error getting evidence from agent: %v", err))
+		http.Error(w, fmt.Sprintf("Internal error: %v", err), http.StatusInternalServerError)
+		return
+	}
+
+	// Construct the response
+	// For the mock, we only include the attestation in BindingKeyAttestation
+	resp := GetKeyEndorsementResponse{
+		Endorsement: KeyEndorsement{
+			VmProtectedKeyEndorsement: VmProtectedKeyEndorsement{
+				BindingKeyAttestation: &KeyAttestation{
+					Attestation: attestation,
+				},
+			},
+		},
+	}
+
+	w.Header().Set("Content-Type", "application/json")
+	if err := json.NewEncoder(w).Encode(resp); err != nil {
+		logger.Error(fmt.Sprintf("Error encoding response: %v", err))
+	}
+}