{"title":"Pickle Deserialization via __reduce__","language":"Python","severity":"Critical","cwe":"CWE-502","source_lines":[9],"flow_lines":[9,11],"sink_lines":[11],"vulnerable_code":"import pickle\nimport base64\nfrom flask import request, jsonify\n\ndef restore_ml_model_state():\n    encrypted_state = request.headers.get('X-Model-Checkpoint')\n    if encrypted_state:\n        model_bytes = base64.b64decode(encrypted_state)\n        restored_model = pickle.loads(model_bytes)\n        restored_model.initialize_weights()\n        return jsonify({'status': 'model_restored', 'accuracy': restored_model.get_metrics()})\n    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":"import json\nimport base64\nimport hmac\nimport hashlib\nfrom flask import request, jsonify, current_app\n\nALLOWED_MODEL_CLASSES = {'LinearRegression', 'NeuralNetwork', 'DecisionTree', 'RandomForest'}\n\ndef verify_checkpoint_signature(data_bytes, signature):\n    \"\"\"Verify HMAC signature of the checkpoint data.\"\"\"\n    secret_key = current_app.config['CHECKPOINT_SIGNING_KEY']\n    expected_sig = hmac.HMAC(secret_key.encode(), data_bytes, hashlib.sha256).hexdigest()\n    return hmac.compare_digest(expected_sig, signature)\n\ndef restore_ml_model_state():\n    encrypted_state = request.headers.get('X-Model-Checkpoint')\n    signature = request.headers.get('X-Checkpoint-Signature')\n\n    if not encrypted_state:\n        return jsonify({'error': 'no_checkpoint'}), 400\n\n    if not signature:\n        return jsonify({'error': 'missing_signature'}), 403\n\n    try:\n        model_bytes = base64.b64decode(encrypted_state)\n    except Exception:\n        return jsonify({'error': 'invalid_base64_encoding'}), 400\n\n    if not verify_checkpoint_signature(model_bytes, signature):\n        return jsonify({'error': 'invalid_checkpoint_signature'}), 403\n\n    try:\n        model_state = json.loads(model_bytes)\n    except (json.JSONDecodeError, ValueError):\n        return jsonify({'error': 'invalid_checkpoint_format'}), 400\n\n    if not isinstance(model_state, dict):\n        return jsonify({'error': 'invalid_checkpoint_structure'}), 400\n\n    model_class = model_state.get('model_class')\n    if model_class not in ALLOWED_MODEL_CLASSES:\n        return jsonify({'error': 'unsupported_model_class'}), 400\n\n    weights = model_state.get('weights')\n    hyperparams = model_state.get('hyperparameters')\n\n    if not isinstance(weights, list) or not isinstance(hyperparams, dict):\n        return jsonify({'error': 'invalid_model_data'}), 400\n\n    restored_model = load_model_from_safe_state(model_class, weights, hyperparams)\n    restored_model.initialize_weights()\n    return jsonify({'status': 'model_restored', 'accuracy': restored_model.get_metrics()})\n\n\ndef load_model_from_safe_state(model_class, weights, hyperparams):\n    \"\"\"Safely reconstruct a model from validated parameters without deserialization.\"\"\"\n    from models import MODEL_REGISTRY\n    model_cls = MODEL_REGISTRY[model_class]\n    model = model_cls(**hyperparams)\n    model.set_weights(weights)\n    return model"}