How to Use Confidential Computing for Secure AI Inference

# Verify your VM supports SEV-SNP
dmesg | grep -i "sev\|snp\|ccp"
# Expected output includes: AMD Secure Encrypted Virtualization (SEV) active

# Dockerfile.confidential
FROM python:3.12-slim

# Minimize attack surface
RUN apt-get update && apt-get install -y --no-install-recommends \
    libgomp1 \
    && rm -rf /var/lib/apt/lists/*

WORKDIR /app

# Install dependencies with pinned versions for reproducible measurement
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

# Copy model and server (weights come from encrypted volume at runtime)
COPY server.py .

# Non-root execution — required for attestation policies in most TEE setups
RUN useradd -m inferenceuser
USER inferenceuser

CMD ["uvicorn", "server:app", "--host", "0.0.0.0", "--port", "8443", "--ssl-keyfile", "/certs/key.pem", "--ssl-certfile", "/certs/cert.pem"]

# server.py
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
import os

app = FastAPI()

# Model path is injected at runtime via encrypted env or secret manager
MODEL_PATH = os.environ.get("MODEL_PATH", "/models/mistral-7b-q4")

# Load once at startup — inside the encrypted enclave
tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_PATH,
    torch_dtype=torch.float16,
    device_map="auto",
    low_cpu_mem_usage=True,  # Critical for large models in constrained enclave memory
)

class InferenceRequest(BaseModel):
    prompt: str
    max_new_tokens: int = 256

@app.post("/infer")
async def infer(req: InferenceRequest):
    if len(req.prompt) > 4096:
        raise HTTPException(status_code=400, detail="Prompt too long")

    inputs = tokenizer(req.prompt, return_tensors="pt").to(model.device)
    
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=req.max_new_tokens,
            do_sample=False,  # Deterministic — important for reproducibility in audits
        )
    
    result = tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
    return {"response": result}

@app.get("/health")
async def health():
    return {"status": "ok"}

# attestation_client.py — run this on the CLIENT before sending data to the enclave

import requests
import json
import base64
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import padding
from cryptography import x509
import hashlib

MAA_ENDPOINT = "https://sharedeus2.eus2.attest.azure.net"

def get_attestation_token(enclave_url: str) -> dict:
    """
    Request a signed attestation token from the enclave.
    The enclave must expose a /attest endpoint that returns an MAA JWT.
    """
    response = requests.get(f"{enclave_url}/attest", timeout=10)
    response.raise_for_status()
    return response.json()

def verify_attestation(token: str, expected_measurement: str) -> bool:
    """
    Verify the MAA JWT against the AMD Root of Trust.
    expected_measurement is the SHA-384 hash of your container image.
    """
    # Fetch MAA signing certificates
    jwks = requests.get(f"{MAA_ENDPOINT}/certs").json()
    
    # Decode and verify the JWT (use a proper JWT library in production)
    # This checks: signature validity, expiry, and enclave measurement
    parts = token.split(".")
    claims = json.loads(base64.urlsafe_b64decode(parts[1] + "=="))
    
    actual_measurement = claims.get("x-ms-isolation-tee", {}).get("x-ms-runtime", {}).get("vm-configuration", {}).get("secure-boot", "")
    
    # Critical: verify the enclave is running exactly the code you expect
    if claims.get("x-ms-compliance-status") != "azure-compliant-cvm":
        return False
    
    print(f"Attestation valid. Enclave is running verified code.")
    print(f"CVM compliance status: {claims.get('x-ms-compliance-status')}")
    return True

# Usage
token_data = get_attestation_token("https://your-enclave.example.com")
is_valid = verify_attestation(token_data["token"], expected_measurement="your-image-sha384")

if is_valid:
    # Now safe to send sensitive data
    response = requests.post(
        "https://your-enclave.example.com/infer",
        json={"prompt": "Patient data: ..."},
        verify=True  # Always verify TLS in production
    )

# Create a key in Azure Key Vault with a release policy
# The key is ONLY released to VMs that pass attestation

az keyvault key create \
  --vault-name your-vault \
  --name model-encryption-key \
  --kty RSA \
  --size 4096 \
  --exportable true \
  --policy @key-release-policy.json

# key-release-policy.json — only release to attested CVMs

{
  "anyOf": [
    {
      "allOf": [
        {
          "claim": "x-ms-compliance-status",
          "equals": "azure-compliant-cvm"
        },
        {
          "claim": "x-ms-isolation-tee.x-ms-attestation-type",
          "equals": "sevsnpvm"
        }
      ]
    }
  ],
  "version": "1.0.0"
}

# Encrypt model weights before uploading to storage
openssl enc -aes-256-gcm \
  -in ./mistral-7b-q4.bin \
  -out ./mistral-7b-q4.bin.enc \
  -pass env:MODEL_ENCRYPTION_KEY \
  -pbkdf2

# Upload encrypted weights to Azure Blob
az storage blob upload \
  --file ./mistral-7b-q4.bin.enc \
  --container-name models \
  --name mistral-7b-q4.bin.enc \
  --auth-mode login

# 1. Check the container image measurement matches your expected value
docker inspect your-inference-image --format='{{.Id}}'
# Record this SHA — it's what you'll verify during attestation

# 2. Test the full attestation flow end-to-end
python attestation_client.py --enclave-url https://your-enclave.example.com --verbose

# 3. Verify inference works with a non-sensitive prompt
curl -X POST https://your-enclave.example.com/infer \
  -H "Content-Type: application/json" \
  -d '{"prompt": "What is 2+2?", "max_new_tokens": 10}'

{
  "attestation_status": "valid",
  "compliance": "azure-compliant-cvm",
  "response": "4"
}