Skip to content

Out-of-bounds read in merge_bijections when n_plans < n_samples #910

Closed

Description

@xihaian251

Description

ot.bsp.merge_bijections() reads past the end of the perms array whenever fewer plans
are supplied than there are samples (n_plans < n_samples). The number of plans is taken
from the wrong axis, so the C++ merge loop iterates n_samples times over a buffer that
holds n_plans * n_samples integers.

This is a memory-safety bug in the public API, not a numerical difference: on an MSVC build
it segfaults the interpreter, and AddressSanitizer reports a heap-buffer-overflow on the
merge loop.

All line numbers below refer to master at 0b6f8d5.

Root cause

The Python API documents perms as (n_plans, n_samples)
(ot/bsp/bsp_ot.py:125),
and the Cython layer allocates the candidate-plan buffer with that layout
(ot/bsp/bsp_wrap.pyx:52,
np.zeros((n_plans, n))).

merge_bijections_c() however reads the plan count off axis 1 — the number of samples:

# ot/bsp/bsp_wrap.pyx:82-90
cdef int n = X.shape[0]
cdef int d = X.shape[1]
cdef int k = plans.shape[1]          # <-- n_samples, not n_plans; should be plans.shape[0]
cdef np.ndarray[int, ndim=1, mode="c"] plan = np.zeros(n, dtype=np.int32)
...
cost_val = MergeBijections(n, d, <double*>X.data, <double*>Y.data, k, <int*> plans.data, <int*> plan.data, lp_power)

MergeBijections() consumes that value as the iteration bound and reads n integers per
"plan":

// ot/bsp/bsp_wrapper.cpp:126-138
double MergeBijections(int n, int d, double *X, double *Y, uint64_t nb_plans, int *plans_ptr, int *final_plan_ptr, int lp_power) {
    std::vector<BijectiveMatching> plans(nb_plans);
    for (std::size_t i = 0; i < nb_plans; ++i) {
        std::vector<int> bij(n);
        std::copy(
            plans_ptr + i * n,
            plans_ptr + (i + 1) * n,   // <-- nb_plans * n integers read in total
            bij.begin());
        plans[i] = BijectiveMatching(bij);

With nb_plans == n_samples == n, the loop reads n * n integers from an allocation of
n_plans * n, i.e. (n - n_plans) * n integers beyond its end. The values read are then
handed to BijectiveMatching as if they were permutations.

The two off-by-shape conditions are asymmetric:

  • n_plans < n_samples — out-of-bounds read, as above.
  • n_plans > n_samples — no overrun, but the plans at index >= n_samples are silently
    dropped, so the merged bijection ignores part of its input.

For comparison, bsp_solve_c() passes n_plans explicitly and is not affected
(ot/bsp/bsp_wrap.pyx:60).
The Python wrapper checks that X and Y agree but never validates perms against
n_samples
(ot/bsp/bsp_ot.py:147-155),
so nothing catches the mismatch before it reaches C++.

Minimal reproduction

The natural call sequence — a solver result fed straight back into the merger:

import numpy as np
import ot
import ot.bsp

n = 32
rng = np.random.RandomState(0)
X = rng.randn(n, 2)
Y = rng.randn(n, 2)

cost, perm, perms = ot.bsp.compute_bspot_bijection(X, Y, n_plans=1, seed=0)
print(perms.shape)                                          # (1, 32)

cost2, merged = ot.bsp.merge_bijections(X, Y, perms, p=2)   # segmentation fault

A second probe separates the overrun from the crash, using only documented inputs. Two arrays
are passed whose content is identical — one row, the identity permutation, shape (1, 32) —
but whose trailing memory differs, because each is a view of the first row of a larger
allocation whose remaining rows hold two different filler permutations:

passed array logical content trailing memory reported cost
view A (1, 32), row 0 = identity rows of roll(arange(32), 1) 3.733570766539
view B (1, 32), row 0 = identity rows of roll(arange(32), 7) 3.285768445257
exact-size copy of that row (1, 32), row 0 = identity unmanaged heap segmentation fault

The first two rows return different answers for the same input, which is only possible if the
function reads data the caller never passed. The third shows what happens when the overrun
leaves readable memory: (32 - 1) * 32 = 992 integers past a 128-byte allocation.

The dropped-input case, with n_samples = 8 and n_plans = 9 (all rows valid permutations):
changing row 8 alone leaves both the cost and the merged permutation unchanged, so the ninth
plan is never read.

Expected behavior

merge_bijections(X, Y, perms, p) merges all perms.shape[0] supplied plans and reads no more
than perms.shape[0] * perms.shape[1] integers, for any n_plans and n_samples. With
n_plans > n_samples every plan should still take part in the merge.

Actual behavior

n_plans < n_samples reads (n_samples - n_plans) * n_samples integers out of bounds, ending
in a crash or a silently wrong answer; n_plans > n_samples ignores the plans at index
>= n_samples.

Relevant source locations

  • ot/bsp/bsp_wrap.pyx:84 — cdef int k = plans.shape[1], the dimension mix-up.
  • ot/bsp/bsp_wrap.pyx:90 — k passed as nb_plans.
  • ot/bsp/bsp_wrapper.cpp:126-138 — loop reading nb_plans * n integers.
  • ot/bsp/bsp_ot.py:114-155 — public merge_bijections, no shape check against n_samples.

Introduced by 048f3ae (#768). Unrelated to PR #898, which only changes seeding in
ot/bsp/bsp_wrapper.cpp; the scope of this report is limited to the nb_plans mix-up.

Environment / reproduction status

Reproduced, not hypothesized, and verified twice.

AddressSanitizer, against ot/bsp/bsp_wrapper.cpp taken unmodified from master
(0b6f8d5) and compiled into a standalone harness that calls MergeBijections() exactly as the
Cython layer does — n = 32, d = 2, a plan buffer of 1 * 32 ints, nb_plans = 32:

n=32 d=2 n_plans=1 buffer_ints=32 nb_plans_passed=32 ints_read=1024
=================================================================
==431==ERROR: AddressSanitizer: heap-buffer-overflow on address 0x50c0000000c0
READ of size 128 at 0x50c0000000c0 thread T0
    #0 in memmove sanitizer_common_interceptors_memintrinsics.inc:98
    #5 std::copy<int*, ...> /usr/include/c++/13/bits/stl_algobase.h:633
    #6 MergeBijections(int, int, double*, double*, unsigned long, int*, int*, int) bsp_wrapper.cpp:134
    #7 main asan_harness.cpp:42
0x50c0000000c0 is located 0 bytes after 128-byte region [0x50c000000040,0x50c0000000c0)
allocated by thread T0 here:
    #1 main asan_harness.cpp:31

The allocation at asan_harness.cpp:31 is new int[n_plans * n], i.e. the 128-byte plan
buffer, and the faulting read begins exactly at its end. The same harness with the correct
plan count (nb_plans = 1) exits cleanly with cost=0.124175 and the identity permutation,
which localizes the fault to the value passed rather than to the merge itself.

  • Ubuntu 24.04 under WSL2, gcc 13.3.0, -std=c++17 -fsanitize=address -g -O0 -fopenmp,
    Eigen from deps/eigen

The segmentation fault and both probe tables above were measured against the built Python
extension, not a harness:

  • Windows 11 (10.0.26200), MSVC 14.44 with /openmp
  • Python 3.13.1, NumPy 2.4.6, SciPy 1.16.3, POT 0.9.8.dev
  • Installed from source with python setup.py build_ext --inplace

Suggested direction

plans.shape[0] for the plan count, plus a perms.shape[1] == X.shape[0] check in the Python
wrapper so a malformed perms raises instead of reaching C++. Happy to open a PR if that would
help.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions