From 144ee2a5e27a281f914909f51cacf1b97c1fb9b5 Mon Sep 17 00:00:00 2001 From: shaggy Date: Sat, 29 Aug 2026 10:22:09 +0530 Subject: [PATCH] Bounds-check scatter indices in the wasm backend The wasm scatter kernels computed an output pointer directly from attacker-controlled index values with no range validation, so an out-of-range index (including a negative one, which sign-extends) wrote outside the output allocation in wasm linear memory. The CPU backend throws 'Invalid indices' for the same inputs. Validate indices in the ScatterNd/TensorScatterUpdate wrappers before entering wasm, matching the CPU backend's error, and guard the C++ kernels as defense in depth. Mirrors the CropAndResize bounds fix. --- .../src/cc/kernels/TensorScatterUpdate.cc | 7 ++++++ tfjs-backend-wasm/src/cc/scatter_impl.cc | 10 ++++++++ tfjs-backend-wasm/src/index_test.ts | 14 +++++++++++ tfjs-backend-wasm/src/kernels/ScatterNd.ts | 25 +++++++++++++++++++ .../src/kernels/TensorScatterUpdate.ts | 24 ++++++++++++++++++ 5 files changed, 80 insertions(+) diff --git a/tfjs-backend-wasm/src/cc/kernels/TensorScatterUpdate.cc b/tfjs-backend-wasm/src/cc/kernels/TensorScatterUpdate.cc index 15ba0b273e7..64a9de4a203 100644 --- a/tfjs-backend-wasm/src/cc/kernels/TensorScatterUpdate.cc +++ b/tfjs-backend-wasm/src/cc/kernels/TensorScatterUpdate.cc @@ -40,6 +40,13 @@ void scatter(const int* indices_ptr, const T* updates_ptr, indices_ptr++; } + // Skip out-of-range indices instead of writing outside the output + // tensor (negative int indices sign-extend and are caught here too). + if (flattened_index >= output_size / slice_size) { + updates_ptr += slice_size; + continue; + } + out_buf_ptr += flattened_index * slice_size; memcpy(out_buf_ptr, updates_ptr, slice_size * dtype_size); diff --git a/tfjs-backend-wasm/src/cc/scatter_impl.cc b/tfjs-backend-wasm/src/cc/scatter_impl.cc index e1b638266f0..1a1a4274a16 100644 --- a/tfjs-backend-wasm/src/cc/scatter_impl.cc +++ b/tfjs-backend-wasm/src/cc/scatter_impl.cc @@ -40,6 +40,16 @@ void scatter(const int* indices_ptr, const T* updates_ptr, indices_ptr++; } + // Skip out-of-range indices instead of writing outside the output + // tensor. flattened_index is computed from int indices, so negative + // values sign-extend to huge size_t values and are caught here too. + if (flattened_index >= output_size / slice_size) { + if (!update_as_scalar) { + updates_ptr += slice_size; + } + continue; + } + T* out = out_buf_ptr + flattened_index * slice_size; for (size_t k = 0; k < slice_size; ++k) { diff --git a/tfjs-backend-wasm/src/index_test.ts b/tfjs-backend-wasm/src/index_test.ts index d9428cb99f9..dff64ded03b 100644 --- a/tfjs-backend-wasm/src/index_test.ts +++ b/tfjs-backend-wasm/src/index_test.ts @@ -265,4 +265,18 @@ describeWithFlags('wasm init', BROWSER_ENVS, () => { expect(() => setWasmPath('too/late')) .toThrowError(/The WASM backend was already initialized. Make sure/); }); + + it('scatterNd rejects out-of-range indices instead of writing OOB', async () => { + const indices = tf.tensor2d([2000000000], [1, 1], 'int32'); + const updates = tf.tensor1d([42], 'float32'); + expect(() => tf.scatterNd(indices, updates, [100])).toThrowError( + /Invalid indices/); + // Negative indices sign-extend in the wasm kernel; they must be rejected + // too, not written before the output allocation. + const negIndices = tf.tensor2d([-1], [1, 1], 'int32'); + expect(() => tf.tensorScatterUpdate( + tf.zeros([100]), negIndices, updates)) + .toThrowError(/Invalid indices/); + }); + }); diff --git a/tfjs-backend-wasm/src/kernels/ScatterNd.ts b/tfjs-backend-wasm/src/kernels/ScatterNd.ts index 264476f2e39..2934330069d 100644 --- a/tfjs-backend-wasm/src/kernels/ScatterNd.ts +++ b/tfjs-backend-wasm/src/kernels/ScatterNd.ts @@ -64,6 +64,10 @@ function scatterNd( const stridesBytes = new Uint8Array(new Int32Array(strides).buffer); + // Match the CPU backend: reject out-of-range indices before entering wasm, + // where the scatter kernel would otherwise write outside the output tensor. + validateScatterIndices(indices, strides, outputSize, sliceSize); + const outId = backend.dataIdMap.get(out.dataId).id; wasmScatterNd( indicesId, updatesId, CppDType[updates.dtype], sliceRank, numUpdates, @@ -72,6 +76,27 @@ function scatterNd( return out; } + +function validateScatterIndices( + indices: Tensor, strides: number[], outputSize: number, + sliceSize: number): void { + const numSlices = outputSize / sliceSize; + const indicesVals = indices.dataSync(); + const sliceRank = strides.length; + const numUpdates = indicesVals.length / sliceRank; + for (let i = 0; i < numUpdates; ++i) { + let flattenedIndex = 0; + for (let j = 0; j < sliceRank; ++j) { + flattenedIndex += indicesVals[i * sliceRank + j] * strides[j]; + } + if (flattenedIndex < 0 || flattenedIndex >= numSlices) { + throw new Error( + `Invalid indices: ${flattenedIndex} does not index into ` + + `${numSlices}`); + } + } +} + export const scatterNdConfig: KernelConfig = { kernelName: ScatterNd, backendName: 'wasm', diff --git a/tfjs-backend-wasm/src/kernels/TensorScatterUpdate.ts b/tfjs-backend-wasm/src/kernels/TensorScatterUpdate.ts index 700985ba5fe..7756eda92d4 100644 --- a/tfjs-backend-wasm/src/kernels/TensorScatterUpdate.ts +++ b/tfjs-backend-wasm/src/kernels/TensorScatterUpdate.ts @@ -1,3 +1,24 @@ + +function validateScatterIndices( + indices: Tensor, strides: number[], outputSize: number, + sliceSize: number): void { + const numSlices = outputSize / sliceSize; + const indicesVals = indices.dataSync(); + const sliceRank = strides.length; + const numUpdates = indicesVals.length / sliceRank; + for (let i = 0; i < numUpdates; ++i) { + let flattenedIndex = 0; + for (let j = 0; j < sliceRank; ++j) { + flattenedIndex += indicesVals[i * sliceRank + j] * strides[j]; + } + if (flattenedIndex < 0 || flattenedIndex >= numSlices) { + throw new Error( + `Invalid indices: ${flattenedIndex} does not index into ` + + `${numSlices}`); + } + } +} + /** * @license * Copyright 2022 Google LLC. All Rights Reserved. @@ -70,6 +91,9 @@ function tensorScatterUpdate(args: { const stridesBytes = new Uint8Array(new Int32Array(strides).buffer); + // Match the CPU backend: reject out-of-range indices before entering wasm. + validateScatterIndices(indices, strides, outputSize, sliceSize); + const outId = backend.dataIdMap.get(out.dataId).id; wasmTensorScatterUpdate( indicesId, updatesId, CppDType[updates.dtype], sliceRank, numUpdates,