{"title":"Arbitrary Code Execution via marshal.loads() on Untrusted Data","language":"Python","severity":"Critical","cwe":"CWE-502","source_lines":[9],"flow_lines":[9,10,11],"sink_lines":[11],"vulnerable_code":"import marshal\nimport base64\nfrom flask import Flask, request, jsonify\n\napp = Flask(__name__)\n\n@app.route('/ai/model/restore', methods=['POST'])\ndef restore_ml_checkpoint():\n    checkpoint_data = request.json.get('serialized_weights')\n    decoded_checkpoint = base64.b64decode(checkpoint_data)\n    model_state = marshal.loads(decoded_checkpoint)\n    weights = model_state['layers']\n    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":"import json\nimport base64\nimport hmac\nimport hashlib\nfrom flask import Flask, request, jsonify\n\napp = Flask(__name__)\n\n# Secret key for HMAC verification of checkpoint integrity\n# In production, load from secure configuration/environment variable\nCHECKPOINT_SECRET_KEY = app.config.get('CHECKPOINT_SECRET_KEY', 'change-this-in-production')\n\ndef verify_checkpoint_signature(data: bytes, signature: str) -> bool:\n    \"\"\"Verify HMAC signature of checkpoint data to ensure integrity.\"\"\"\n    expected_sig = hmac.new(\n        CHECKPOINT_SECRET_KEY.encode('utf-8'),\n        data,\n        hashlib.sha256\n    ).hexdigest()\n    return hmac.compare_digest(expected_sig, signature)\n\ndef validate_weights_structure(data: dict) -> bool:\n    \"\"\"Validate that deserialized data conforms to expected model weight structure.\"\"\"\n    if not isinstance(data, dict):\n        return False\n    if 'layers' not in data:\n        return False\n    if not isinstance(data['layers'], list):\n        return False\n    # Validate each layer contains only numeric weight data\n    for layer in data['layers']:\n        if not isinstance(layer, dict):\n            return False\n        for key, value in layer.items():\n            if not isinstance(key, str):\n                return False\n            if not isinstance(value, (list, int, float)):\n                return False\n            if isinstance(value, list):\n                if not all(isinstance(v, (int, float)) for v in value):\n                    return False\n    return True\n\n@app.route('/ai/model/restore', methods=['POST'])\ndef restore_ml_checkpoint():\n    checkpoint_data = request.json.get('serialized_weights')\n    signature = request.json.get('signature')\n\n    if not checkpoint_data or not signature:\n        return jsonify({'error': 'Missing checkpoint data or signature'}), 400\n\n    try:\n        decoded_checkpoint = base64.b64decode(checkpoint_data)\n    except Exception:\n        return jsonify({'error': 'Invalid base64 encoding'}), 400\n\n    # Verify the signature to ensure data was produced by a trusted source\n    if not verify_checkpoint_signature(decoded_checkpoint, signature):\n        return jsonify({'error': 'Invalid checkpoint signature'}), 403\n\n    # Use JSON for safe deserialization instead of marshal\n    try:\n        model_state = json.loads(decoded_checkpoint)\n    except (json.JSONDecodeError, UnicodeDecodeError):\n        return jsonify({'error': 'Invalid checkpoint format'}), 400\n\n    # Validate structure before processing\n    if not validate_weights_structure(model_state):\n        return jsonify({'error': 'Invalid model weight structure'}), 400\n\n    weights = model_state['layers']\n    return jsonify({'status': 'restored', 'layer_count': len(weights)})"}