Trusted Execution Environments voor AI-workloads

# Example: Loading a model into an SGX enclave using Gramine (formerly Graphene)
# gramine-manifest.toml configuration for PyTorch inference
 
"""
Gramine manifest for running PyTorch model inference inside SGX enclave.
This configuration specifies trusted files (model weights) and
enclave memory settings.
"""
 
import subprocess
import json
from pathlib import Path
 
def generate_gramine_manifest(model_path: str, enclave_size_mb: int = 4096) -> str:
    """Generate a Gramine manifest for SGX-protected model inference."""
    manifest = f"""
[loader]
entrypoint = "file:{{{{ gramine.libos }}}}"
log_level = "warning"
 
[libos]
entrypoint = "/usr/bin/python3"
 
[loader.env]
LD_LIBRARY_PATH = "/lib:/usr/lib:/usr/local/lib"
PYTHONPATH = "/usr/local/lib/python3.11/site-packages"
 
[loader.argv]
argv0 = "python3"
argv1 = "inference_server.py"
 
[sys]
# Enclave memory size - must accommodate model + runtime
enclave_size = "{enclave_size_mb}M"
# Number of threads inside enclave
thread_num = 8
 
[sgx]
# Enable SGX debug mode (disable in production)
debug = false
# Remote attestation type
ra_client_spid = ""
ra_client_linkable = false
 
# Trusted files - model weights are measured at load time
[[sgx.trusted_files]]
uri = "file:{model_path}"
 
[[sgx.trusted_files]]
uri = "file:inference_server.py"
 
# Allowed (but unmeasured) files
[[sgx.allowed_files]]
uri = "file:/tmp/inference_input"
 
[[sgx.allowed_files]]
uri = "file:/tmp/inference_output"
"""
    return manifest
 
def verify_sgx_attestation(quote: bytes, expected_mrenclave: str) -> bool:
    """
    Verify an SGX attestation quote against expected enclave measurement.
 
    In production, this would contact Intel's Attestation Service (IAS)
    or use DCAP for local verification.
    """
    # Extract MRENCLAVE from quote (offset 112, 32 bytes in SGX quote structure)
    mrenclave_offset = 112
    mrenclave = quote[mrenclave_offset:mrenclave_offset + 32].hex()
 
    if mrenclave != expected_mrenclave:
        raise SecurityError(
            f"MRENCLAVE mismatch: got {mrenclave}, "
            f"expected {expected_mrenclave}"
        )
 
    # In production: verify quote signature via IAS or DCAP
    # result = dcap_verify(quote)
    # if not result.is_valid:
    #     raise SecurityError("Quote signature verification failed")
 
    return True

"""
Demonstration of cache timing side-channel analysis against
TEE-protected neural network inference.
 
This script analyzes timing variations in cache access patterns
to infer which neurons were activated during inference.
"""
 
import numpy as np
import ctypes
import time
from typing import List, Tuple
 
class CacheTimingAnalyzer:
    """
    Analyzes cache timing measurements to extract information
    about TEE-protected model inference.
    """
 
    # Typical L1 cache line size on x86
    CACHE_LINE_SIZE = 64
    # Threshold in CPU cycles to distinguish cache hit vs miss
    # (calibrate per-system; typical values: 100-200 cycles)
    HIT_MISS_THRESHOLD = 150
 
    def __init__(self, num_samples: int = 10000):
        self.num_samples = num_samples
        self.timing_data: List[np.ndarray] = []
 
    def prime_probe_round(
        self,
        monitored_cache_sets: List[int],
        wait_callback: callable,
    ) -> np.ndarray:
        """
        Execute one round of Prime+Probe on specified cache sets.
 
        1. PRIME: Fill target cache sets with attacker's data
        2. WAIT: Let the victim (enclave) execute
        3. PROBE: Measure access time to attacker's data
           - Slow access = victim evicted the line (accessed that cache set)
           - Fast access = victim did not use that cache set
        """
        num_sets = len(monitored_cache_sets)
        timings = np.zeros(num_sets, dtype=np.int64)
 
        # In a real attack, this would use rdtsc and carefully
        # crafted eviction sets. This is a conceptual demonstration.
        for i, cache_set in enumerate(monitored_cache_sets):
            # Prime phase: access attacker's buffer at addresses
            # mapping to the target cache set
            # (omitted: actual memory access with precise addressing)
 
            # Wait for victim computation
            wait_callback()
 
            # Probe phase: time re-access to attacker's buffer
            start = time.perf_counter_ns()
            # (omitted: actual memory access)
            elapsed = time.perf_counter_ns() - start
            timings[i] = elapsed
 
        return timings
 
    def collect_traces(
        self,
        monitored_sets: List[int],
        trigger_inference: callable,
    ) -> np.ndarray:
        """Collect multiple timing traces across inference invocations."""
        traces = []
        for _ in range(self.num_samples):
            timings = self.prime_probe_round(
                monitored_sets,
                wait_callback=trigger_inference,
            )
            traces.append(timings)
        return np.array(traces)
 
    def analyze_neuron_activation(
        self,
        traces: np.ndarray,
    ) -> np.ndarray:
        """
        Infer neuron activation patterns from cache timing traces.
 
        Cache sets that consistently show evictions (slow access)
        during inference correspond to weight matrix regions that
        were accessed, indicating activated neurons.
        """
        # Classify each measurement as hit (0) or miss (1)
        hit_miss = (traces > self.HIT_MISS_THRESHOLD).astype(np.float32)
 
        # Average miss rate per cache set across all samples
        miss_rates = hit_miss.mean(axis=0)
 
        # High miss rate = frequently accessed by enclave = active neuron
        # Low miss rate = not accessed = inactive neuron (e.g., ReLU zeroed)
        activation_threshold = 0.5
        inferred_activations = (miss_rates > activation_threshold).astype(int)
 
        return inferred_activations
 
    def reconstruct_weight_signs(
        self,
        traces_positive_input: np.ndarray,
        traces_negative_input: np.ndarray,
    ) -> np.ndarray:
        """
        Infer weight signs by comparing activation patterns
        for positive vs negative inputs.
 
        If a neuron activates for positive input but not negative,
        the corresponding weight is likely positive (assuming ReLU).
        """
        pos_activations = self.analyze_neuron_activation(traces_positive_input)
        neg_activations = self.analyze_neuron_activation(traces_negative_input)
 
        # +1 = positive weight, -1 = negative weight, 0 = indeterminate
        weight_signs = pos_activations.astype(int) - neg_activations.astype(int)
        return weight_signs

#!/usr/bin/env bash
# Controlled-channel attack setup using SGX-Step framework
# SGX-Step provides precise single-stepping of enclave execution
# Reference: Van Bulck et al., "SGX-Step: A Practical Attack Framework
# for Precise Enclave Execution Control", SysTEX 2017
 
# Clone and build SGX-Step (requires SGX hardware and Linux kernel module)
git clone https://github.com/jovanbulck/sgx-step.git
cd sgx-step
 
# Build the kernel module for precise interrupt control
cd kernel
make
sudo insmod sgx-step.ko
 
# The attack uses APIC timer interrupts to single-step enclave execution.
# Each step reveals the instruction page and data page accessed.
# For a neural network inference:
#   - Code page sequence -> model architecture (layer types, order)
#   - Data page sequence -> weight access pattern -> weight values (partial)
 
# Example: Configure page-level monitoring for a target enclave
cd ../app
cat > config.h << 'CEOF'
#define VICTIM_ENCLAVE_PATH "/opt/ai-service/enclave.signed.so"
// Monitor data pages in the weight matrix region
#define MONITOR_DATA_START  0x7f0000000000ULL
#define MONITOR_DATA_END    0x7f0040000000ULL  // 1GB region for model weights
// Interrupt frequency: once per ~100 instructions for coarse tracing
// or once per instruction for fine-grained extraction
#define SINGLE_STEP_MODE    0  // Set to 1 for instruction-level granularity
CEOF
 
echo "Build the controlled-channel monitor:"
make
echo "Run with: sudo ./app/monitor --enclave-pid <target_pid>"

"""
Reconstruct neural network architecture from page fault traces
collected during controlled-channel attack on SGX enclave.
 
Each page fault reveals a 4KB-aligned memory access. By mapping
these to code and data regions of the inference binary, we can
determine layer types, sizes, and execution order.
"""
 
from dataclasses import dataclass
from enum import Enum
from typing import List, Optional
import struct
 
class LayerType(Enum):
    CONV2D = "conv2d"
    LINEAR = "linear"
    BATCHNORM = "batchnorm"
    RELU = "relu"
    ATTENTION = "attention"
    UNKNOWN = "unknown"
 
@dataclass
class InferredLayer:
    layer_type: LayerType
    estimated_params: int
    page_range: tuple[int, int]
    code_pages: List[int]
 
# Known code page signatures for common operations in PyTorch/ONNX Runtime
# These are determined by reverse-engineering the inference binary
LAYER_SIGNATURES = {
    # GEMM kernel touches specific code pages in MKL/oneDNN
    LayerType.LINEAR: {
        "code_pattern": [0x1000, 0x1001, 0x1002],  # MKL SGEMM pages
        "data_pattern": "sequential",  # Linear weight access is sequential
    },
    LayerType.CONV2D: {
        "code_pattern": [0x2000, 0x2001, 0x2002, 0x2003],  # im2col + GEMM
        "data_pattern": "strided",  # Filter weights accessed with stride
    },
    LayerType.ATTENTION: {
        "code_pattern": [0x1000, 0x1001, 0x3000, 0x3001],  # GEMM + softmax
        "data_pattern": "qkv_triple",  # Three weight matrices accessed
    },
}
 
def parse_page_fault_log(log_path: str) -> List[tuple[int, int, str]]:
    """
    Parse page fault log from SGX-Step controlled-channel attack.
    Returns list of (timestamp, page_address, access_type) tuples.
    """
    events = []
    with open(log_path, "r") as f:
        for line in f:
            parts = line.strip().split(",")
            if len(parts) != 3:
                continue
            timestamp = int(parts[0])
            page_addr = int(parts[1], 16)
            access_type = parts[2]  # "code" or "data"
            events.append((timestamp, page_addr, access_type))
    return events
 
def identify_layer_boundaries(
    events: List[tuple[int, int, str]],
) -> List[tuple[int, int]]:
    """
    Identify layer boundaries by detecting transitions between
    code regions. Each layer type uses a distinct set of code pages.
    """
    boundaries = []
    current_code_set = set()
    boundary_start = 0
 
    for i, (ts, addr, atype) in enumerate(events):
        if atype == "code":
            if addr not in current_code_set and len(current_code_set) > 2:
                # Significant change in code pages = new layer
                boundaries.append((boundary_start, i))
                boundary_start = i
                current_code_set = set()
            current_code_set.add(addr)
 
    if boundary_start < len(events):
        boundaries.append((boundary_start, len(events)))
 
    return boundaries
 
def infer_architecture(log_path: str) -> List[InferredLayer]:
    """
    Full pipeline: parse page faults -> identify layers -> classify.
    """
    events = parse_page_fault_log(log_path)
    boundaries = identify_layer_boundaries(events)
    layers = []
 
    for start, end in boundaries:
        segment = events[start:end]
        code_pages = sorted(set(
            addr for _, addr, atype in segment if atype == "code"
        ))
        data_pages = sorted(set(
            addr for _, addr, atype in segment if atype == "data"
        ))
 
        # Match code pages against known signatures
        layer_type = LayerType.UNKNOWN
        for ltype, sig in LAYER_SIGNATURES.items():
            pattern = sig["code_pattern"]
            if all(p in code_pages for p in pattern):
                layer_type = ltype
                break
 
        # Estimate parameter count from number of unique data pages
        # Each page is 4KB; fp32 weights are 4 bytes each
        estimated_params = len(data_pages) * (4096 // 4)
 
        layers.append(InferredLayer(
            layer_type=layer_type,
            estimated_params=estimated_params,
            page_range=(data_pages[0] if data_pages else 0,
                       data_pages[-1] if data_pages else 0),
            code_pages=code_pages,
        ))
 
    return layers
 
if __name__ == "__main__":
    import sys
    if len(sys.argv) != 2:
        print(f"Usage: {sys.argv[0]} <page_fault_log>")
        sys.exit(1)
 
    layers = infer_architecture(sys.argv[1])
    print(f"Inferred {len(layers)} layers:")
    for i, layer in enumerate(layers):
        print(
            f"  Layer {i}: {layer.layer_type.value}, "
            f"~{layer.estimated_params:,} params, "
            f"pages {layer.page_range[0]:#x}-{layer.page_range[1]:#x}"
        )

"""
Audit script for verifying AMD SEV-SNP attestation reports
in confidential AI VM deployments.
 
Checks that the VM is running with expected security properties:
- SNP is enabled (not just SEV or SEV-ES)
- Firmware version meets minimum requirements
- Guest policy enforces migration restrictions
- The TCB (Trusted Computing Base) version is current
"""
 
import hashlib
import struct
import json
from pathlib import Path
from dataclasses import dataclass
from typing import Optional
 
@dataclass
class SNPAttestationReport:
    """Parsed AMD SEV-SNP attestation report fields."""
    version: int
    guest_svn: int
    policy: int
    family_id: bytes
    image_id: bytes
    vmpl: int
    signature_algo: int
    current_tcb: int
    platform_info: int
    author_key_en: int
    report_data: bytes
    measurement: bytes  # 384-bit hash of initial guest memory
    host_data: bytes
    id_key_digest: bytes
    author_key_digest: bytes
    report_id: bytes
    chip_id: bytes
    committed_tcb: int
    current_build: int
    committed_build: int
    launch_tcb: int
    signature: bytes
 
def parse_attestation_report(raw_report: bytes) -> SNPAttestationReport:
    """
    Parse a raw SEV-SNP attestation report (0x4A0 bytes).
    Reference: AMD SEV-SNP ABI Specification, Table 21.
    """
    if len(raw_report) < 0x4A0:
        raise ValueError(
            f"Report too short: {len(raw_report)} bytes, "
            f"expected at least 0x4A0"
        )
 
    version = struct.unpack_from("<I", raw_report, 0x0)[0]
    guest_svn = struct.unpack_from("<I", raw_report, 0x4)[0]
    policy = struct.unpack_from("<Q", raw_report, 0x8)[0]
    family_id = raw_report[0x10:0x20]
    image_id = raw_report[0x20:0x30]
    vmpl = struct.unpack_from("<I", raw_report, 0x30)[0]
    signature_algo = struct.unpack_from("<I", raw_report, 0x34)[0]
    current_tcb = struct.unpack_from("<Q", raw_report, 0x38)[0]
    platform_info = struct.unpack_from("<Q", raw_report, 0x40)[0]
    author_key_en = struct.unpack_from("<I", raw_report, 0x48)[0]
    report_data = raw_report[0x50:0x90]
    measurement = raw_report[0x90:0xC0]
    host_data = raw_report[0xC0:0xE0]
    id_key_digest = raw_report[0xE0:0x110]
    author_key_digest = raw_report[0x110:0x140]
    report_id = raw_report[0x140:0x160]
    chip_id = raw_report[0x1A0:0x1E0]
    committed_tcb = struct.unpack_from("<Q", raw_report, 0x1E0)[0]
    current_build = struct.unpack_from("<I", raw_report, 0x1E8)[0]
    committed_build = struct.unpack_from("<I", raw_report, 0x1EC)[0]
    launch_tcb = struct.unpack_from("<Q", raw_report, 0x1F0)[0]
    signature = raw_report[0x2A0:0x4A0]
 
    return SNPAttestationReport(
        version=version, guest_svn=guest_svn, policy=policy,
        family_id=family_id, image_id=image_id, vmpl=vmpl,
        signature_algo=signature_algo, current_tcb=current_tcb,
        platform_info=platform_info, author_key_en=author_key_en,
        report_data=report_data, measurement=measurement,
        host_data=host_data, id_key_digest=id_key_digest,
        author_key_digest=author_key_digest, report_id=report_id,
        chip_id=chip_id, committed_tcb=committed_tcb,
        current_build=current_build, committed_build=committed_build,
        launch_tcb=launch_tcb, signature=signature,
    )
 
def audit_snp_report(
    report: SNPAttestationReport,
    expected_measurement: str,
    min_tcb_version: int = 0x03000000000008,
) -> list[str]:
    """
    Audit an SEV-SNP attestation report for security issues.
    Returns a list of findings.
    """
    findings = []
 
    # Check report version
    if report.version < 2:
        findings.append(
            "CRITICAL: Report version < 2. SNP attestation requires v2+."
        )
 
    # Verify measurement matches expected value
    measurement_hex = report.measurement.hex()
    if measurement_hex != expected_measurement:
        findings.append(
            f"CRITICAL: Measurement mismatch. "
            f"Expected {expected_measurement}, got {measurement_hex}. "
            f"The VM image may have been tampered with."
        )
 
    # Check VMPL (Virtual Machine Privilege Level)
    if report.vmpl != 0:
        findings.append(
            f"WARNING: VMPL is {report.vmpl}, not 0. "
            f"Higher VMPL has reduced privileges."
        )
 
    # Check guest policy flags
    # Bit 0: SMT allowed (potential side-channel risk)
    # Bit 1: Reserved
    # Bit 2: Migration allowed (confidentiality risk)
    if report.policy & 0x1:
        findings.append(
            "WARNING: SMT is allowed by guest policy. "
            "Simultaneous multithreading enables side-channel attacks."
        )
    if report.policy & 0x4:
        findings.append(
            "CRITICAL: Migration is allowed by guest policy. "
            "A malicious hypervisor could migrate the VM to extract secrets."
        )
 
    # Check TCB version
    if report.current_tcb < min_tcb_version:
        findings.append(
            f"WARNING: TCB version {report.current_tcb:#x} is below "
            f"minimum {min_tcb_version:#x}. Firmware may have known "
            f"vulnerabilities."
        )
 
    # Check platform info
    if report.platform_info & 0x1:
        findings.append(
            "INFO: Platform reports SMT enabled at hardware level."
        )
 
    if not findings:
        findings.append("PASS: All attestation checks passed.")
 
    return findings
 
if __name__ == "__main__":
    import sys
 
    if len(sys.argv) != 3:
        print(f"Usage: {sys.argv[0]} <report_file> <expected_measurement>")
        sys.exit(1)
 
    raw = Path(sys.argv[1]).read_bytes()
    report = parse_attestation_report(raw)
    findings = audit_snp_report(report, sys.argv[2])
 
    for finding in findings:
        print(finding)