{"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    encoded_state = request.headers.get('X-Model-Checkpoint')\n    if encoded_state:\n        checkpoint_data = base64.b64decode(encoded_state)\n        model_params = pickle.loads(checkpoint_data)\n        return jsonify({'status': 'restored', 'params': str(model_params)})\n    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":"import json\nimport base64\nimport hmac\nimport hashlib\nfrom flask import request, jsonify, current_app\n\nALLOWED_PARAM_KEYS = {'weights', 'biases', 'learning_rate', 'epochs', 'layers', 'optimizer', 'batch_size'}\nALLOWED_PARAM_TYPES = (int, float, str, list, dict, bool, type(None))\n\ndef validate_model_params(params):\n    \"\"\"Validate that model parameters contain only safe types and expected keys.\"\"\"\n    if not isinstance(params, dict):\n        return False\n    for key in params:\n        if key not in ALLOWED_PARAM_KEYS:\n            return False\n        if not _is_safe_value(params[key]):\n            return False\n    return True\n\ndef _is_safe_value(value):\n    \"\"\"Recursively check that a value contains only safe types.\"\"\"\n    if isinstance(value, (int, float, str, bool, type(None))):\n        return True\n    if isinstance(value, list):\n        return all(_is_safe_value(item) for item in value)\n    if isinstance(value, dict):\n        return all(isinstance(k, str) and _is_safe_value(v) for k, v in value.items())\n    return False\n\ndef verify_signature(data_bytes, signature):\n    \"\"\"Verify HMAC signature of the checkpoint data.\"\"\"\n    secret_key = current_app.config.get('MODEL_CHECKPOINT_SECRET')\n    if not secret_key:\n        return False\n    expected_sig = hmac.new(secret_key.encode(), data_bytes, hashlib.sha256).hexdigest()\n    return hmac.compare_digest(expected_sig, signature)\n\ndef restore_ml_model_state():\n    encoded_state = request.headers.get('X-Model-Checkpoint')\n    signature = request.headers.get('X-Checkpoint-Signature')\n\n    if not encoded_state:\n        return jsonify({'error': 'no checkpoint'}), 400\n\n    if not signature:\n        return jsonify({'error': 'missing signature'}), 403\n\n    try:\n        checkpoint_data = base64.b64decode(encoded_state)\n    except Exception:\n        return jsonify({'error': 'invalid base64 encoding'}), 400\n\n    if not verify_signature(checkpoint_data, signature):\n        return jsonify({'error': 'invalid signature'}), 403\n\n    try:\n        model_params = json.loads(checkpoint_data)\n    except (json.JSONDecodeError, UnicodeDecodeError):\n        return jsonify({'error': 'invalid checkpoint format'}), 400\n\n    if not validate_model_params(model_params):\n        return jsonify({'error': 'invalid model parameters'}), 400\n\n    return jsonify({'status': 'restored', 'params': str(model_params)})"}