Fix/serialize security

docs/developer_guide/testing.md

@@ -42,6 +42,39 @@ representation, update the pinned values in the same commit and call it out in t
 request, so the change is reviewed rather than hidden.
 ```
 
+## Testing mathematical correctness, not just shape
+
+A representation-heavy library like PreTab has a failure mode that shape and dtype checks
+cannot catch: a basis function, penalty matrix, or encoding can be computed with the wrong
+formula and still produce output of the right shape, the right dtype, and finite values. A
+test that only asserts `X.shape == (n, k)` or `np.isfinite(out).all()` will pass on both the
+correct and the incorrect implementation.
+
+When a representation has a closed-form mathematical property, test that property directly
+instead of (or in addition to) its shape:
+
+- **Known identities.** A B-spline basis should sum to `1` at every point (partition of
+  unity); an M-spline should integrate to `1` over its own support; an I-spline should be
+  monotonically non-decreasing and bounded in `[0, 1]`.
+- **Independent reference values.** A penalty matrix or a hand-derivable formula (a
+  particular basis value at a particular knot, say) can be checked against a value computed
+  a different way, for example a fine numerical quadrature or a direct closed-form
+  substitution, not just re-derived with the same code path the implementation itself uses.
+- **Boundary behaviour.** Values at, or just past, a fitted range's edge are where
+  clipping-versus-extrapolation bugs and off-by-one integration bounds hide. Test a value
+  exactly at the boundary and one just beyond it, not only values safely inside the range.
+- **Realistic missing-data shapes.** A mixed object array with an actual `NaN`/`None` among
+  string categories (the ordinary shape of a pandas column with missing values) is a
+  different code path than an all-numeric array with `NaN`, and needs its own test if a
+  transformer declares `allow_nan=True`.
+
+```{warning}
+Shape/symmetry/finiteness assertions are still useful as a first line of defense, but they
+are not sufficient proof that a mathematical implementation is correct: they pass equally
+well on a subtly wrong formula as on a correct one. Pair them with at least one value-level
+assertion for anything that has a defined mathematical property to check against.
+```
+
 ## Markers
 
 The suite defines a `smoke` marker for fast end-to-end sanity checks that run as a dedicated CI

docs/homepage.md

@@ -1,12 +1,15 @@
 # PreTab
 
-**PreTab** is a modular, extensible, and [scikit-learn](https://scikit-learn.org/)-compatible
-preprocessing library for tabular data. It supports **all `sklearn` transformers** out of the
-box and extends them with a rich set of custom encoders, splines, and neural basis expansions.
+**PreTab** is a modular, [scikit-learn](https://scikit-learn.org/)-compatible representation
+and preprocessing library for tabular data. Its focus is feature representation and
+expansion: splines, neural basis maps, piecewise-linear encoding, kernel approximations, and
+language embeddings, each shipped as a standalone transformer that speaks the standard
+`fit` / `transform` API. Every PreTab transformer subclasses `BaseEstimator` and
+`TransformerMixin`, so it drops into a `Pipeline` or `ColumnTransformer` alongside any
+sklearn-native transformer you already use.
 
 ```{note}
-These docs are for PreTab {{ version }}. The project is under active development and the
-public API may evolve while the major version is `0`.
+These docs are for PreTab {{ version }}.
 ```
 
 ## Highlights

pretab/compose/serialize.py

@@ -25,12 +25,29 @@
 
 from .._version import __version__ as _PRETAB_VERSION
 from ..core.parameters import UNSET
+from ..core.policy import RepresentationPolicy
+from ..core.representation import FeatureLineage, RepresentationSpec
 from ..exceptions import PretabError, PretabSerializationError
+from ..placement.base import PlacementResult
+from .registry import TransformerSpec
 
 SCHEMA_VERSION = 1
 
-# Top-level packages that a spec is allowed to import classes/types from.
-_ALLOWED_TOP_LEVEL = frozenset({"pretab", "sklearn", "numpy", "scipy", "builtins"})
+# Top-level packages that a spec is allowed to *import a module from*. This is
+# only the first line of defense (see ``_ALLOWED_DATACLASSES`` and the
+# ``BaseEstimator`` check below for the checks that actually gate what gets
+# called/instantiated). ``builtins`` is deliberately excluded: nothing PreTab
+# ever needs to reconstruct from a spec lives there, and allowing it is what let
+# a crafted spec call ``builtins.open`` with attacker-controlled arguments.
+_ALLOWED_TOP_LEVEL = frozenset({"pretab", "sklearn", "numpy", "scipy"})
+
+# The exact, closed set of dataclasses a spec is allowed to reconstruct via the
+# ``__dataclass__`` tag. Reconstruction calls ``cls(**fields)`` (i.e. runs the
+# dataclass's ``__init__``), so this must be an exact allow-list, not merely
+# "any dataclass importable under an allowed module" -- the latter would still
+# let a spec construct an arbitrary sklearn/numpy/scipy dataclass with
+# attacker-controlled fields.
+_ALLOWED_DATACLASSES = frozenset({RepresentationPolicy, RepresentationSpec, FeatureLineage, PlacementResult, TransformerSpec})
 
 
 # --- helpers -------------------------------------------------------------
@@ -152,14 +169,41 @@ def _decode_mapping(mapping: dict) -> dict:
 
 
 def _decode_estimator(payload: dict):
+    """Reconstruct a ``BaseEstimator`` from an ``__estimator__`` tag.
+
+    Bypasses ``__init__``/``__reduce__``/``__setstate__`` by design (a plain
+    ``__new__`` plus a ``__dict__`` update), but only for a class that is
+    actually a registered scikit-learn estimator -- anything else (a crafted
+    ``\"class\"`` naming an unrelated callable) is refused before it is ever
+    instantiated.
+    """
     cls = cast(Any, _resolve(payload["class"]))
+    if not (isinstance(cls, type) and issubclass(cls, BaseEstimator)):
+        raise PretabSerializationError(
+            f"Refusing to reconstruct {payload['class']!r} as an estimator: not a "
+            "scikit-learn BaseEstimator subclass."
+        )
     obj = cls.__new__(cls)
     obj.__dict__.update(_decode_mapping(payload["state"]))
     return obj
 
 
 def _decode_dataclass(payload: dict):
+    """Reconstruct a dataclass from a ``__dataclass__`` tag.
+
+    Unlike :func:`_decode_estimator`, this calls ``cls(**fields)`` (the
+    dataclass's real ``__init__``), so the allow-list must be an *exact* set of
+    approved classes, not merely "any dataclass reachable under an allowed
+    module" -- refusing anything outside ``_ALLOWED_DATACLASSES`` is what stops
+    a crafted spec from constructing an arbitrary callable with
+    attacker-controlled keyword arguments.
+    """
     cls = cast(Any, _resolve(payload["class"]))
+    if not (isinstance(cls, type) and dataclasses.is_dataclass(cls) and cls in _ALLOWED_DATACLASSES):
+        allowed = sorted(c.__qualname__ for c in _ALLOWED_DATACLASSES)
+        raise PretabSerializationError(
+            f"Refusing to reconstruct disallowed dataclass {payload['class']!r}. Only {allowed} are permitted."
+        )
     fields = {k: _decode(v) for k, v in payload["fields"].items()}
     return cls(**fields)
 

tests/integration/test_serialization.py

@@ -197,6 +197,80 @@ def test_from_spec_refuses_disallowed_module(frame, target):
         Preprocessor.from_spec(spec)
 
 
+def test_from_spec_refuses_builtins_module():
+    """``builtins`` must not be resolvable at all (closes the arbitrary-call hole).
+
+    Previously ``builtins`` was an allowed top-level module, so a crafted
+    ``__dataclass__``/``__estimator__`` tag naming ``builtins:open`` could call
+    ``open(**attacker_fields)`` with attacker-controlled keyword arguments,
+    including creating/truncating an arbitrary file. This asserts the module is
+    refused outright, and that no file is created as a side effect of the
+    refused load.
+    """
+    import os
+    import tempfile
+
+    target_path = os.path.join(tempfile.gettempdir(), "pretab_test_should_not_exist.txt")
+    if os.path.exists(target_path):
+        os.remove(target_path)
+
+    malicious_spec = {
+        "schema_version": SCHEMA_VERSION,
+        "state": {
+            "__dict__": [
+                [
+                    "evil",
+                    {
+                        "__dataclass__": {
+                            "class": "builtins:open",
+                            "fields": {"file": target_path, "mode": "w"},
+                        }
+                    },
+                ]
+            ]
+        },
+    }
+    with pytest.raises(PretabSerializationError, match="disallowed module"):
+        Preprocessor.from_spec(malicious_spec)
+    assert not os.path.exists(target_path)
+
+
+def test_from_spec_refuses_estimator_not_a_base_estimator():
+    """An allowed-module class that isn't a ``BaseEstimator`` must still be refused.
+
+    Proves the ``__estimator__`` check is a structural ``issubclass`` check, not
+    just a module-prefix check: ``numpy`` is an allowed module, but
+    ``numpy.ndarray`` is not a scikit-learn estimator.
+    """
+    spec = {
+        "schema_version": SCHEMA_VERSION,
+        "state": {"__dict__": [["evil", {"__estimator__": {"class": "numpy:ndarray", "state": {}}}]]},
+    }
+    with pytest.raises(PretabSerializationError, match="not a scikit-learn BaseEstimator"):
+        Preprocessor.from_spec(spec)
+
+
+def test_from_spec_refuses_dataclass_not_in_allowlist():
+    """An allowed-module, genuinely-a-dataclass class must still be refused unless
+    it is one of the specifically approved dataclasses.
+
+    ``PreprocessorConfig`` is a real, internal PreTab dataclass (module "pretab",
+    would pass a module-prefix check) that is deliberately not part of a fitted
+    ``Preprocessor``'s serialized state, so it must not be reconstructable via a
+    spec either.
+    """
+    spec = {
+        "schema_version": SCHEMA_VERSION,
+        "state": {
+            "__dict__": [
+                ["evil", {"__dataclass__": {"class": "pretab.compose.config:PreprocessorConfig", "fields": {}}}]
+            ]
+        },
+    }
+    with pytest.raises(PretabSerializationError, match="disallowed dataclass"):
+        Preprocessor.from_spec(spec)
+
+
 def test_from_spec_rejects_bad_source_type():
     with pytest.raises(PretabSerializationError):
         Preprocessor.from_spec(12345)