diff --git a/paddle/phi/kernels/cpu/embedding_kernel.cc b/paddle/phi/kernels/cpu/embedding_kernel.cc index fa0fb09f42802..0c768f6df6143 100644 --- a/paddle/phi/kernels/cpu/embedding_kernel.cc +++ b/paddle/phi/kernels/cpu/embedding_kernel.cc @@ -49,27 +49,31 @@ struct EmbeddingCPUFunctor { dev_ctx_.template Alloc(out_); auto* output = out_->data(); + // Validate every id unconditionally. The range check must not be gated on + // padding_idx: `kNoPadding` is -1, so when the layer is created with any + // padding_idx (e.g. the common padding_idx=0) the previous guard skipped + // validation for all ids, allowing an out-of-bounds read in the copy loop + // below (CWE-125). Only the padding-row zeroing should depend on + // padding_idx. for (int64_t i = 0; i < ids_numel; ++i) { - if (padding_idx_ == kNoPadding && ids[i] != padding_idx_) { - PADDLE_ENFORCE_LT( - ids[i], - row_number, - common::errors::InvalidArgument( - "Variable value (input) of OP(embedding) " - "expected >= 0 and < %ld, but got %ld. Please check input " - "value.", - row_number, - ids[i])); - PADDLE_ENFORCE_GE( - ids[i], - 0, - common::errors::InvalidArgument( - "Variable value (input) of OP(embedding) " - "expected >= 0 and < %ld, but got %ld. Please check input " - "value.", - row_number, - ids[i])); - } + PADDLE_ENFORCE_LT( + ids[i], + row_number, + common::errors::InvalidArgument( + "Variable value (input) of OP(embedding) " + "expected >= 0 and < %ld, but got %ld. Please check input " + "value.", + row_number, + ids[i])); + PADDLE_ENFORCE_GE( + ids[i], + 0, + common::errors::InvalidArgument( + "Variable value (input) of OP(embedding) " + "expected >= 0 and < %ld, but got %ld. Please check input " + "value.", + row_number, + ids[i])); } #if defined(_OPENMP) && !defined(PADDLE_WITH_CUDA) diff --git a/test/ai_edited_test/test_ai_nn_functional_input.py b/test/ai_edited_test/test_ai_nn_functional_input.py index b97cf12659108..8ca2072f7d0c1 100644 --- a/test/ai_edited_test/test_ai_nn_functional_input.py +++ b/test/ai_edited_test/test_ai_nn_functional_input.py @@ -134,6 +134,27 @@ def test_embedding_invalid_padding_idx(self): with self.assertRaises(ValueError): paddle.nn.functional.embedding(x, weight, padding_idx=10) + def test_embedding_out_of_range_id_with_padding_idx(self): + """测试设置 padding_idx 时越界输入 id 会报错(CPU 内核越界读回归) + When padding_idx is set, an out-of-range input id must raise instead + of reading out of bounds (CWE-125 regression for the CPU kernel). + Forces the CPU place so the test is independent of the build target.""" + orig_device = paddle.device.get_device() + paddle.set_device('cpu') + try: + weight = paddle.randn([16, 8]) + # id 40 is out of range for a 16-row table; padding_idx=0 must not + # cause the range check to be skipped. + x = paddle.to_tensor([0, 40], dtype='int64') + with self.assertRaises(ValueError): + paddle.nn.functional.embedding(x, weight, padding_idx=0) + # A negative out-of-range id must also raise. + x_neg = paddle.to_tensor([0, -3], dtype='int64') + with self.assertRaises(ValueError): + paddle.nn.functional.embedding(x_neg, weight, padding_idx=0) + finally: + paddle.set_device(orig_device) + def test_embedding_sparse(self): """测试 sparse 模式的 embedding Test embedding in sparse mode"""