From 350dddedd837b3b75de7e7dbc51dd9bbe210056b Mon Sep 17 00:00:00 2001 From: Gideon Koenig Date: Mon, 27 Jul 2026 22:41:44 +0200 Subject: [PATCH] [tfjs-node] Prevent tensor ID rollover leaks and collisions --- tfjs-node/binding/tfjs_backend.cc | 31 ++++++++++++++++-- tfjs-node/binding/tfjs_backend.h | 4 +++ tfjs-node/binding/tfjs_binding.cc | 34 ++++++++++++++++++++ tfjs-node/src/nodejs_kernel_backend.ts | 12 +++++-- tfjs-node/src/nodejs_kernel_backend_test.ts | 25 +++++++++++++++ tfjs-node/src/tfjs_binding_test.ts | 35 +++++++++++++++++++++ 6 files changed, 137 insertions(+), 4 deletions(-) diff --git a/tfjs-node/binding/tfjs_backend.cc b/tfjs-node/binding/tfjs_backend.cc index 4267cf5b960..c5eaa1142be 100644 --- a/tfjs-node/binding/tfjs_backend.cc +++ b/tfjs-node/binding/tfjs_backend.cc @@ -19,6 +19,7 @@ #include #include +#include #include #include #include @@ -34,6 +35,8 @@ namespace tfnodejs { +constexpr int32_t kTensorIdNotOnDevice = -1; + // Used to hold strings beyond the lifetime of a JS call. static std::set ATTR_NAME_SET; @@ -717,9 +720,33 @@ TFJSBackend::~TFJSBackend() { TFJSBackend *TFJSBackend::Create(napi_env env) { return new TFJSBackend(env); } +int32_t TFJSBackend::NextTensorId() { + const int32_t tensor_id = next_tensor_id_; + next_tensor_id_ = tensor_id == std::numeric_limits::max() + ? std::numeric_limits::min() + : tensor_id + 1; + + // -1 represents tensor data that has not been uploaded to the native + // backend. + if (next_tensor_id_ == kTensorIdNotOnDevice) { + next_tensor_id_ = 0; + } + return tensor_id; +} + int32_t TFJSBackend::InsertHandle(TFE_TensorHandle *tfe_handle) { - return tfe_handle_map_.insert(std::make_pair(next_tensor_id_++, tfe_handle)) - .first->first; + while (true) { + const int32_t tensor_id = NextTensorId(); + const auto insert_result = tfe_handle_map_.emplace(tensor_id, tfe_handle); + if (insert_result.second) { + return tensor_id; + } + } +} + +void TFJSBackend::SetNextTensorIdForTest(int32_t next_tensor_id) { + next_tensor_id_ = + next_tensor_id == kTensorIdNotOnDevice ? 0 : next_tensor_id; } int32_t TFJSBackend::InsertSavedModel(TF_Session *tf_session, diff --git a/tfjs-node/binding/tfjs_backend.h b/tfjs-node/binding/tfjs_backend.h index fee71d06074..2559ccceb85 100644 --- a/tfjs-node/binding/tfjs_backend.h +++ b/tfjs-node/binding/tfjs_backend.h @@ -88,9 +88,13 @@ class TFJSBackend { // Get number of tensor handles in the backend: napi_value GetNumOfTensors(napi_env env); + // Sets the next tensor ID. Intended for testing only. + void SetNextTensorIdForTest(int32_t next_tensor_id); + private: TFJSBackend(napi_env env); + int32_t NextTensorId(); int32_t InsertHandle(TFE_TensorHandle *tfe_handle); int32_t InsertSavedModel(TF_Session *tf_session, TF_Graph *tf_graph); napi_value GenerateOutputTensorInfo(napi_env env, TFE_TensorHandle *handle); diff --git a/tfjs-node/binding/tfjs_binding.cc b/tfjs-node/binding/tfjs_binding.cc index ad00b79e78f..ea0e9870637 100644 --- a/tfjs-node/binding/tfjs_binding.cc +++ b/tfjs-node/binding/tfjs_binding.cc @@ -274,6 +274,38 @@ static napi_value GetNumOfSavedModels(napi_env env, napi_callback_info info) { return backend->GetNumOfSavedModels(env); } +static napi_value SetNextTensorIdForTest(napi_env env, + napi_callback_info info) { + napi_status nstatus; + + size_t argc = 1; + napi_value args[1]; + napi_value js_this; + nstatus = napi_get_cb_info(env, info, &argc, args, &js_this, nullptr); + ENSURE_NAPI_OK_RETVAL(env, nstatus, js_this); + + if (argc < 1) { + NAPI_THROW_ERROR( + env, + "Invalid number of args passed to setNextTensorIdForTest(). " + "Expecting 1 arg but got %d.", + argc); + return js_this; + } + + ENSURE_VALUE_IS_NUMBER_RETVAL(env, args[0], js_this); + + int32_t next_tensor_id; + nstatus = napi_get_value_int32(env, args[0], &next_tensor_id); + ENSURE_NAPI_OK_RETVAL(env, nstatus, js_this); + + TFJSBackend *const backend = GetTFJSBackend(env); + if (!backend) return nullptr; + + backend->SetNextTensorIdForTest(next_tensor_id); + return js_this; +} + /** * Called by Node to cleanup our instance data, which is * the TFJSBackend allocated in `InitTFNodeJSBinding`. @@ -323,6 +355,8 @@ static napi_value InitTFNodeJSBinding(napi_env env, napi_value exports) { nullptr, napi_default, nullptr}, {"getNumOfTensors", nullptr, GetNumOfTensors, nullptr, nullptr, nullptr, napi_default, nullptr}, + {"setNextTensorIdForTest", nullptr, SetNextTensorIdForTest, nullptr, + nullptr, nullptr, napi_default, nullptr}, }; nstatus = napi_define_properties(env, exports, ARRAY_SIZE(exports_properties), diff --git a/tfjs-node/src/nodejs_kernel_backend.ts b/tfjs-node/src/nodejs_kernel_backend.ts index a192f313854..0d2f49a8342 100644 --- a/tfjs-node/src/nodejs_kernel_backend.ts +++ b/tfjs-node/src/nodejs_kernel_backend.ts @@ -32,6 +32,8 @@ type TensorData = { refCount: number; }; +const TENSOR_ID_NOT_ON_DEVICE = -1; + export class NodeJSKernelBackend extends KernelBackend { binding: TFJSBinding; isGPUPackage: boolean; @@ -245,7 +247,7 @@ export class NodeJSKernelBackend extends KernelBackend { return false; } - if (id != null && id >= 0) { + if (id != null && id !== TENSOR_ID_NOT_ON_DEVICE) { this.binding.deleteTensor(id); } this.tensorMap.delete(dataId); @@ -269,7 +271,13 @@ export class NodeJSKernelBackend extends KernelBackend { dataId: DataId, values: backend_util.BackendValues, shape: number[], dtype: DataType, refCount: number): void { this.tensorMap.set( - dataId, {shape, dtype: getTFDType(dtype), values, id: -1, refCount}); + dataId, { + shape, + dtype: getTFDType(dtype), + values, + id: TENSOR_ID_NOT_ON_DEVICE, + refCount + }); } write(values: backend_util.BackendValues, shape: number[], dtype: DataType): diff --git a/tfjs-node/src/nodejs_kernel_backend_test.ts b/tfjs-node/src/nodejs_kernel_backend_test.ts index 19718cbe54d..54976cd2d8b 100644 --- a/tfjs-node/src/nodejs_kernel_backend_test.ts +++ b/tfjs-node/src/nodejs_kernel_backend_test.ts @@ -21,6 +21,10 @@ import {TestKernelBackend} from '@tensorflow/tfjs-core/dist/jasmine_util'; import {createTensorsTypeOpAttr, ensureTensorflowBackend, getTFDType, nodeBackend, NodeJSKernelBackend} from './nodejs_kernel_backend'; +type TestBinding = NodeJSKernelBackend['binding']&{ + setNextTensorIdForTest(nextTensorId: number): void; +}; + describe('delayed upload', () => { it('should handle data before op execution', async () => { const t = tf.tensor1d([1, 2, 3]); @@ -40,6 +44,27 @@ describe('delayed upload', () => { }); }); +describe('tensor ID rollover', () => { + it('disposes native tensors with negative IDs', () => { + const binding = nodeBackend().binding as TestBinding; + const restoreId = binding.createTensor( + [1], binding.TF_INT32, new Int32Array([0])); + const initialNumTensors = binding.getNumOfTensors(); + + try { + binding.setNextTensorIdForTest(2147483647); + tf.tidy(() => { + const result = tf.tensor1d([1, 2]).square(); + tf.test_util.expectArraysClose(result.dataSync(), [1, 4]); + }); + expect(binding.getNumOfTensors()).toBe(initialNumTensors); + } finally { + binding.setNextTensorIdForTest(restoreId + 1); + binding.deleteTensor(restoreId); + } + }); +}); + describe('type casting', () => { it('exp support int32', () => { tf.exp(tf.scalar(2, 'int32')); diff --git a/tfjs-node/src/tfjs_binding_test.ts b/tfjs-node/src/tfjs_binding_test.ts index c02c0e3253e..599d8c0e739 100644 --- a/tfjs-node/src/tfjs_binding_test.ts +++ b/tfjs-node/src/tfjs_binding_test.ts @@ -24,6 +24,9 @@ const bindingPath = // tslint:disable-next-line:no-require-imports const bindings = require(bindingPath); const binding = bindings as TFJSBinding; +const testBinding = bindings as TFJSBinding&{ + setNextTensorIdForTest(nextTensorId: number): void; +}; describe('Exposes TF_DataType enum values', () => { it('contains TF_FLOAT', () => { @@ -116,6 +119,38 @@ describe('tensor management', () => { expect(binding.tensorDataSync(outputMetadata[0].id)) .toEqual(new Int32Array([3])); }); + it('wraps tensor IDs, skips -1, and avoids live ID collisions', () => { + const initialNumTensors = binding.getNumOfTensors(); + const values = new Int32Array([42]); + const persistentId = + binding.createTensor([1], binding.TF_INT32, values); + const createdIds: number[] = []; + + try { + testBinding.setNextTensorIdForTest(persistentId); + const collisionId = + binding.createTensor([1], binding.TF_INT32, new Int32Array([7])); + expect(collisionId).not.toBe(persistentId); + binding.deleteTensor(collisionId); + expect(binding.tensorDataSync(persistentId)).toEqual(values); + + testBinding.setNextTensorIdForTest(2147483647); + createdIds.push(binding.createTensor([1], binding.TF_INT32, values)); + createdIds.push(binding.createTensor([1], binding.TF_INT32, values)); + expect(createdIds).toEqual([2147483647, -2147483648]); + + testBinding.setNextTensorIdForTest(-2); + createdIds.push(binding.createTensor([1], binding.TF_INT32, values)); + createdIds.push(binding.createTensor([1], binding.TF_INT32, values)); + expect(createdIds[2]).toBe(-2); + expect(createdIds[3]).not.toBe(-1); + } finally { + createdIds.forEach(id => binding.deleteTensor(id)); + binding.deleteTensor(persistentId); + testBinding.setNextTensorIdForTest(persistentId + 1); + } + expect(binding.getNumOfTensors()).toBe(initialNumTensors); + }); }); describe('executeOp', () => {