# 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():
    encrypted_state = request.headers.get('X-Model-Checkpoint')
    if encrypted_state:
        model_bytes = base64.b64decode(encrypted_state)
        restored_model = pickle.loads(model_bytes)
        restored_model.initialize_weights()
        return jsonify({'status': 'model_restored', 'accuracy': restored_model.get_metrics()})
    return jsonify({'error': 'no_checkpoint'}), 400
```

## Explanation

The application accepts arbitrary serialized data from an HTTP header (X-Model-Checkpoint), decodes it, and directly deserializes it using pickle.loads() without any validation. An attacker can craft a malicious pickle payload that executes arbitrary code during deserialization by exploiting Python's __reduce__ method, leading to Remote Code Execution on the server.

## Remediation

The fix eliminates pickle deserialization entirely, replacing it with JSON-based safe deserialization that only accepts structured data (model class name, weights as lists, and hyperparameters as dicts). Additionally, HMAC signature verification ensures that only checkpoints signed by the server's secret key are accepted, preventing tampering. The model is reconstructed programmatically from validated parameters using an allowlist of model classes.

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

ALLOWED_MODEL_CLASSES = {'LinearRegression', 'NeuralNetwork', 'DecisionTree', 'RandomForest'}

def verify_checkpoint_signature(data_bytes, signature):
    """Verify HMAC signature of the checkpoint data."""
    secret_key = current_app.config['CHECKPOINT_SIGNING_KEY']
    expected_sig = hmac.HMAC(secret_key.encode(), data_bytes, hashlib.sha256).hexdigest()
    return hmac.compare_digest(expected_sig, signature)

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

    if not encrypted_state:
        return jsonify({'error': 'no_checkpoint'}), 400

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

    try:
        model_bytes = base64.b64decode(encrypted_state)
    except Exception:
        return jsonify({'error': 'invalid_base64_encoding'}), 400

    if not verify_checkpoint_signature(model_bytes, signature):
        return jsonify({'error': 'invalid_checkpoint_signature'}), 403

    try:
        model_state = json.loads(model_bytes)
    except (json.JSONDecodeError, ValueError):
        return jsonify({'error': 'invalid_checkpoint_format'}), 400

    if not isinstance(model_state, dict):
        return jsonify({'error': 'invalid_checkpoint_structure'}), 400

    model_class = model_state.get('model_class')
    if model_class not in ALLOWED_MODEL_CLASSES:
        return jsonify({'error': 'unsupported_model_class'}), 400

    weights = model_state.get('weights')
    hyperparams = model_state.get('hyperparameters')

    if not isinstance(weights, list) or not isinstance(hyperparams, dict):
        return jsonify({'error': 'invalid_model_data'}), 400

    restored_model = load_model_from_safe_state(model_class, weights, hyperparams)
    restored_model.initialize_weights()
    return jsonify({'status': 'model_restored', 'accuracy': restored_model.get_metrics()})


def load_model_from_safe_state(model_class, weights, hyperparams):
    """Safely reconstruct a model from validated parameters without deserialization."""
    from models import MODEL_REGISTRY
    model_cls = MODEL_REGISTRY[model_class]
    model = model_cls(**hyperparams)
    model.set_weights(weights)
    return model
```