# Arbitrary Code Execution via marshal.loads() on Untrusted Data

Language: Python
Severity: Critical
CWE: CWE-502

## Source
9

## Flow
9-10-11

## Sink
11

## Vulnerable Code
```python
import marshal
import base64
from flask import Flask, request, jsonify

app = Flask(__name__)

@app.route('/ai/model/restore', methods=['POST'])
def restore_ml_checkpoint():
    checkpoint_data = request.json.get('serialized_weights')
    decoded_checkpoint = base64.b64decode(checkpoint_data)
    model_state = marshal.loads(decoded_checkpoint)
    weights = model_state['layers']
    return jsonify({'status': 'restored', 'layer_count': len(weights)})
```

## Explanation

The application accepts untrusted serialized data from user input via the 'serialized_weights' parameter and directly deserializes it using marshal.loads(). The marshal module is explicitly documented as unsafe for untrusted data because it can execute arbitrary code during deserialization, leading to Remote Code Execution (RCE).

## Remediation

The fix replaces the unsafe marshal.loads() deserialization with json.loads(), which cannot execute arbitrary code during parsing. Additionally, HMAC signature verification ensures only checkpoints produced by trusted sources are accepted, and strict schema validation confirms the deserialized data conforms to the expected model weight structure before processing.

## Secure Code
```python
import json
import base64
import hmac
import hashlib
from flask import Flask, request, jsonify

app = Flask(__name__)

# Secret key for HMAC verification of checkpoint integrity
# In production, load from secure configuration/environment variable
CHECKPOINT_SECRET_KEY = app.config.get('CHECKPOINT_SECRET_KEY', 'change-this-in-production')

def verify_checkpoint_signature(data: bytes, signature: str) -> bool:
    """Verify HMAC signature of checkpoint data to ensure integrity."""
    expected_sig = hmac.new(
        CHECKPOINT_SECRET_KEY.encode('utf-8'),
        data,
        hashlib.sha256
    ).hexdigest()
    return hmac.compare_digest(expected_sig, signature)

def validate_weights_structure(data: dict) -> bool:
    """Validate that deserialized data conforms to expected model weight structure."""
    if not isinstance(data, dict):
        return False
    if 'layers' not in data:
        return False
    if not isinstance(data['layers'], list):
        return False
    # Validate each layer contains only numeric weight data
    for layer in data['layers']:
        if not isinstance(layer, dict):
            return False
        for key, value in layer.items():
            if not isinstance(key, str):
                return False
            if not isinstance(value, (list, int, float)):
                return False
            if isinstance(value, list):
                if not all(isinstance(v, (int, float)) for v in value):
                    return False
    return True

@app.route('/ai/model/restore', methods=['POST'])
def restore_ml_checkpoint():
    checkpoint_data = request.json.get('serialized_weights')
    signature = request.json.get('signature')

    if not checkpoint_data or not signature:
        return jsonify({'error': 'Missing checkpoint data or signature'}), 400

    try:
        decoded_checkpoint = base64.b64decode(checkpoint_data)
    except Exception:
        return jsonify({'error': 'Invalid base64 encoding'}), 400

    # Verify the signature to ensure data was produced by a trusted source
    if not verify_checkpoint_signature(decoded_checkpoint, signature):
        return jsonify({'error': 'Invalid checkpoint signature'}), 403

    # Use JSON for safe deserialization instead of marshal
    try:
        model_state = json.loads(decoded_checkpoint)
    except (json.JSONDecodeError, UnicodeDecodeError):
        return jsonify({'error': 'Invalid checkpoint format'}), 400

    # Validate structure before processing
    if not validate_weights_structure(model_state):
        return jsonify({'error': 'Invalid model weight structure'}), 400

    weights = model_state['layers']
    return jsonify({'status': 'restored', 'layer_count': len(weights)})
```