Bounds-check scatter indices in the WASM backend

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);

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) {

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/);
+  });
+
 });

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',

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,