codekingpro/portable-devtools

1"""Validation utilities for LangChain serialization.2 3Provides escape-based protection against injection attacks in serialized objects. The4approach uses an allowlist design: only dicts explicitly produced by5`Serializable.to_json()` are treated as LC objects during deserialization.6 7## How escaping works8 9During serialization, plain dicts (user data) that contain an `'lc'` key are wrapped:10 11```python12{"lc": 1, ...}  # user data that looks like LC object13# becomes:14{"__lc_escaped__": {"lc": 1, ...}}15```16 17During deserialization, escaped dicts are unwrapped and returned as plain dicts,18NOT instantiated as LC objects.19"""20 21from typing import Any, cast22 23from langchain_core.load.serializable import (24    Serializable,25    to_json_not_implemented,26)27 28_LC_ESCAPED_KEY = "__lc_escaped__"29"""Sentinel key used to mark escaped user dicts during serialization.30 31When a plain dict contains 'lc' key (which could be confused with LC objects),32we wrap it as {"__lc_escaped__": {...original...}}.33"""34 35 36def _needs_escaping(obj: dict[str, Any]) -> bool:37    """Check if a dict needs escaping to prevent confusion with LC objects.38 39    A dict needs escaping if:40 41    1. It has an `'lc'` key (could be confused with LC serialization format)42    2. It has only the escape key (would be mistaken for an escaped dict)43    """44    return "lc" in obj or (len(obj) == 1 and _LC_ESCAPED_KEY in obj)45 46 47def _escape_dict(obj: dict[str, Any]) -> dict[str, Any]:48    """Wrap a dict in the escape marker.49 50    Example:51        ```python52        {"key": "value"}  # becomes {"__lc_escaped__": {"key": "value"}}53        ```54    """55    return {_LC_ESCAPED_KEY: obj}56 57 58def _is_escaped_dict(obj: dict[str, Any]) -> bool:59    """Check if a dict is an escaped user dict.60 61    Example:62        ```python63        {"__lc_escaped__": {...}}  # is an escaped dict64        ```65    """66    return len(obj) == 1 and _LC_ESCAPED_KEY in obj67 68 69def _serialize_value(obj: Any) -> Any:70    """Serialize a value with escaping of user dicts.71 72    Called recursively on kwarg values to escape any plain dicts that could be confused73    with LC objects.74 75    Args:76        obj: The value to serialize.77 78    Returns:79        The serialized value with user dicts escaped as needed.80    """81    if isinstance(obj, Serializable):82        # This is an LC object - serialize it properly (not escaped)83        return _serialize_lc_object(obj)84    if isinstance(obj, dict):85        if not all(isinstance(k, (str, int, float, bool, type(None))) for k in obj):86            # if keys are not json serializable87            return to_json_not_implemented(obj)88        # Check if dict needs escaping BEFORE recursing into values.89        # If it needs escaping, wrap it as-is - the contents are user data that90        # will be returned as-is during deserialization (no instantiation).91        # This prevents re-escaping of already-escaped nested content.92        if _needs_escaping(obj):93            return _escape_dict(obj)94        # Safe dict (no 'lc' key) - recurse into values95        return {k: _serialize_value(v) for k, v in obj.items()}96    if isinstance(obj, (list, tuple)):97        return [_serialize_value(item) for item in obj]98    if isinstance(obj, (str, int, float, bool, type(None))):99        return obj100 101    # Non-JSON-serializable object (datetime, custom objects, etc.)102    return to_json_not_implemented(obj)103 104 105def _get_secret_keys(obj: Serializable) -> set[str]:106    """Return the merged set of constructor kwarg names declared as secrets.107 108    Mirrors the MRO walk in `Serializable.to_json` so the keys returned here109    match the keys whose values `_replace_secrets` rewrites into secret110    markers. Used by `_serialize_lc_object` to decide which kwargs to skip111    when escaping user data.112    """113    secrets: dict[str, str] = {}114    model_fields = type(obj).model_fields115    for cls in [None, *obj.__class__.mro()]:116        if cls is Serializable:117            break118        this = cast("Serializable", obj if cls is None else super(cls, obj))119        secrets.update(this.lc_secrets)120        for key in list(secrets):121            if (key in model_fields) and (alias := model_fields[key].alias) is not None:122                secrets[alias] = secrets[key]123    return set(secrets)124 125 126def _serialize_lc_object(obj: Any) -> dict[str, Any]:127    """Serialize a `Serializable` object with escaping of user data in kwargs.128 129    Args:130        obj: The `Serializable` object to serialize.131 132    Returns:133        The serialized dict with user data in kwargs escaped as needed.134 135    Note:136        Kwargs values are processed with `_serialize_value` to escape user data137        (like metadata) that contains `'lc'` keys. Secret fields are identified138        by the class's declared `lc_secrets` and skipped because `to_json()`139        already converted their values to secret markers.140 141        The check is key-based rather than shape-based. A shape-based check142        ("this dict looks like a secret marker") can be forged by user data,143        letting attacker-controlled free-form dicts bypass escaping and reach144        the Reviver.145    """146    if not isinstance(obj, Serializable):147        msg = f"Expected Serializable, got {type(obj)}"148        raise TypeError(msg)149 150    serialized: dict[str, Any] = dict(obj.to_json())151 152    # Process kwargs to escape user data that could be confused with LC objects.153    # Skip kwargs declared as secrets - `to_json()` already replaced their154    # values with secret markers via `_replace_secrets`.155    if serialized.get("type") == "constructor" and "kwargs" in serialized:156        secret_keys = _get_secret_keys(obj)157        serialized["kwargs"] = {158            k: v if k in secret_keys else _serialize_value(v)159            for k, v in serialized["kwargs"].items()160        }161 162    return serialized163 164 165def _unescape_value(obj: Any) -> Any:166    """Unescape a value, processing escape markers in dict values and lists.167 168    When an escaped dict is encountered (`{"__lc_escaped__": ...}`), it's169    unwrapped and the contents are returned AS-IS (no further processing).170    The contents represent user data that should not be modified.171 172    For regular dicts and lists, we recurse to find any nested escape markers.173 174    Args:175        obj: The value to unescape.176 177    Returns:178        The unescaped value.179    """180    if isinstance(obj, dict):181        if _is_escaped_dict(obj):182            # Unwrap and return the user data as-is (no further unescaping).183            # The contents are user data that may contain more escape keys,184            # but those are part of the user's actual data.185            return obj[_LC_ESCAPED_KEY]186 187        # Regular dict - recurse into values to find nested escape markers188        return {k: _unescape_value(v) for k, v in obj.items()}189    if isinstance(obj, list):190        return [_unescape_value(item) for item in obj]191    return obj192