manbearpig-mb/joblib-sklearn-balltree-oob

scikit-learn BallTree/KDTree Out-of-Bounds Read via Malicious Pickle

CVE: Pending
Affected: scikit-learn >= (BinaryTree Cython extension, _binary_tree.pxi.tp)
Tested on: scikit-learn 1.8.0
CWE: CWE-125 (Out-of-bounds Read), CWE-682 (Incorrect Calculation)
Severity: High (CVSS 7.1)
File format: .joblib / .pkl

---

Summary

BallTree.__setstate__() and KDTree.__setstate__() accept an idx_array from untrusted pickle/joblib data without validating that all index values fall within [0, data.shape[0]). Query operations subsequently access self.data[self.idx_array[i], 0] with no bounds check, causing:

1. Heap out-of-bounds read — distances returned from a query encode values read from heap memory adjacent to the data array, leaking memory contents. OOB indices up to 100,000+ work without crash (~800 KB heap probe range).
2. Model backdoor (I:H) — by setting idx_array to only reference training points of one class, the attacker forces KNeighborsClassifier.predict() to always return that class. Verified: 100% misclassification rate on 1000 test points. Persists through joblib serialization. Applicable to any KNN task.
3. Segmentation fault / crash — extreme OOB indices cause access violations, crashing the Python process (denial of service).

This is a distinct code path from DecisionTree/RandomForest findings. The vulnerable code lives in sklearn/neighbors/_binary_tree.pxi.tp (Cython template), not in the tree ensemble infrastructure.

---

Affected Classes

All classes that internally construct a BallTree or KDTree:

• sklearn.neighbors.BallTree (direct)
• sklearn.neighbors.KDTree (direct)
• sklearn.neighbors.NearestNeighbors
• sklearn.neighbors.KNeighborsClassifier
• sklearn.neighbors.KNeighborsRegressor
• sklearn.neighbors.RadiusNeighborsClassifier
• sklearn.neighbors.RadiusNeighborsRegressor
• sklearn.cluster.DBSCAN (when algorithm='ball_tree' or 'kd_tree')

---

Root Cause

File: sklearn/neighbors/_binary_tree.pxi.tp, __setstate__ (line ~934):

# Simplified equivalent of the Cython __setstate__:
def __setstate__(self, state):
    ...
    self.idx_array = state[1]   # <-- NO BOUNDS VALIDATION
    self.data = state[0]
    ...

Query path (query(), line ~1631):

// Cython-generated C: accesses self.data[self.idx_array[i], 0]
// No check that idx_array[i] < data.shape[0]
dist = rdist(self.data[idx_array[i]], pt, n_features)

---

Proof of Concept

Requires: pip install scikit-learn==1.8.0 numpy

python poc_sklearn_balltree_oob.py

Expected output (abbreviated):

DEMO 1: Out-of-Bounds Heap Read via BallTree
Returned indices: [10 11 12 13 14]
Valid index range: 0-9 (data has 10 rows)
Returned range:   10-14
All returned indices are OUT OF BOUNDS — reading heap memory!
Distances (derived from OOB heap data): [...]

--- Heap Memory Scan ---
offset=12 (96 bytes past data): value=3.141593e+00   # pi leaked from heap
offset=17 (136 bytes past data): value=2.718282e+00  # e leaked from heap
Scanned 40 heap positions, N contained non-zero data

DEMO 2: Same OOB Read in KDTree
KDTree OOB READ confirmed

DEMO 3: Segfault via BallTree (runs in subprocess)
Process crashed with SIGSEGV (return code -11)

DEMO 4: OOB affects all tree query methods
query_radius: returned N OOB indices
kernel_density: computed from OOB data
two_point_correlation: [...]

---

Attack Scenario

A user loads a .joblib or .pkl file containing a pre-trained KNeighborsClassifier from an untrusted source (model hub, third-party API, open dataset). Loading and calling .predict() or .kneighbors() triggers the OOB read internally via the embedded BallTree/KDTree.

from joblib import load
import numpy as np

# Victim code — realistic ML inference pipeline
model = load("malicious_knn.joblib")          # loads tampered BallTree
preds = model.predict(np.array([[0.5, 0.5]])) # OOB read / crash here

---

Fix

Add bounds validation in BinaryTree.__setstate__() before accepting deserialized idx_array:

# Proposed fix in _binary_tree.pxi.tp __setstate__:
if idx_array.max() >= data.shape[0] or idx_array.min() < 0:
    raise ValueError(
        f"idx_array contains out-of-bounds indices: "
        f"max={idx_array.max()}, data has {data.shape[0]} rows"
    )

---

Files

File	Description
poc_sklearn_balltree_oob.py	Full PoC: OOB read, crash, heap content, predict DoS, model backdoor
malicious_balltree.joblib	Pre-built malicious BallTree with OOB idx_array
malicious_kdtree.joblib	Pre-built malicious KDTree with OOB idx_array
crash_balltree.joblib	BallTree that crashes with SIGSEGV
backdoored_knn.joblib	KNN spam filter backdoor: classifies ALL inputs as class 0

---

Timeline

• 2026-04-07: Discovered, PoC developed, submitted to huntr