# Pickle Deserialization via __reduce__

Language: Python
Severity: Critical
CWE: CWE-502

## Source
9

## Flow
9-11

## Sink
11

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

def restore_ml_model_state():
    encoded_state = request.headers.get('X-Model-Checkpoint')
    if encoded_state:
        checkpoint_data = base64.b64decode(encoded_state)
        model_params = pickle.loads(checkpoint_data)
        return jsonify({'status': 'restored', 'params': str(model_params)})
    return jsonify({'error': 'no checkpoint'}), 400
```

## Explanation

The application accepts untrusted pickle data from an HTTP header (X-Model-Checkpoint), base64-decodes it, and deserializes it using pickle.loads(). Pickle deserialization of untrusted data allows arbitrary code execution because pickle can reconstruct malicious objects with __reduce__ methods that execute system commands during deserialization.

## Remediation

The fix replaces unsafe pickle deserialization with JSON parsing, which cannot execute arbitrary code during deserialization. Additionally, it adds HMAC signature verification to ensure checkpoint data originates from a trusted source, and validates the deserialized data structure against an allowlist of expected parameter keys and safe types.

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

ALLOWED_PARAM_KEYS = {'weights', 'biases', 'learning_rate', 'epochs', 'layers', 'optimizer', 'batch_size'}
ALLOWED_PARAM_TYPES = (int, float, str, list, dict, bool, type(None))

def validate_model_params(params):
    """Validate that model parameters contain only safe types and expected keys."""
    if not isinstance(params, dict):
        return False
    for key in params:
        if key not in ALLOWED_PARAM_KEYS:
            return False
        if not _is_safe_value(params[key]):
            return False
    return True

def _is_safe_value(value):
    """Recursively check that a value contains only safe types."""
    if isinstance(value, (int, float, str, bool, type(None))):
        return True
    if isinstance(value, list):
        return all(_is_safe_value(item) for item in value)
    if isinstance(value, dict):
        return all(isinstance(k, str) and _is_safe_value(v) for k, v in value.items())
    return False

def verify_signature(data_bytes, signature):
    """Verify HMAC signature of the checkpoint data."""
    secret_key = current_app.config.get('MODEL_CHECKPOINT_SECRET')
    if not secret_key:
        return False
    expected_sig = hmac.new(secret_key.encode(), data_bytes, hashlib.sha256).hexdigest()
    return hmac.compare_digest(expected_sig, signature)

def restore_ml_model_state():
    encoded_state = request.headers.get('X-Model-Checkpoint')
    signature = request.headers.get('X-Checkpoint-Signature')

    if not encoded_state:
        return jsonify({'error': 'no checkpoint'}), 400

    if not signature:
        return jsonify({'error': 'missing signature'}), 403

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

    if not verify_signature(checkpoint_data, signature):
        return jsonify({'error': 'invalid signature'}), 403

    try:
        model_params = json.loads(checkpoint_data)
    except (json.JSONDecodeError, UnicodeDecodeError):
        return jsonify({'error': 'invalid checkpoint format'}), 400

    if not validate_model_params(model_params):
        return jsonify({'error': 'invalid model parameters'}), 400

    return jsonify({'status': 'restored', 'params': str(model_params)})
```