Skip to content

fix: validate embedding ids unconditionally in CPU kernel to prevent OOB read (CWE-125) #79507

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Pandya-mayur wants to merge 2 commits into
base: develop
Choose a base branch
from
Open

fix: validate embedding ids unconditionally in CPU kernel to prevent OOB read (CWE-125) #79507

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
f566ba0
fix: validate embedding ids unconditionally in CPU kernel (CWE-125)
Pandya-mayur Jul 18, 2026
4bd7ff7
test: add CPU embedding out-of-range id regression test (CWE-125)
Pandya-mayur Jul 18, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
44 changes: 24 additions & 20 deletions paddle/phi/kernels/cpu/embedding_kernel.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,27 +49,31 @@ struct EmbeddingCPUFunctor {
dev_ctx_.template Alloc<T>(out_);
auto* output = out_->data<T>();

// 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(

@risemeup1111 risemeup1111 Jul 18, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 优先级：P1

这次把检查条件改对了，但没有补回归测试。现有 test_ai_nn_functional_input.py / test_ai_embedding_onehot.py 只覆盖合法 padding_idx、负 padding_idx 和 padding_idx 越界，没有覆盖“padding_idx 已设置时输入 id 越界”的 CPU 路径。没有这个 case，后续很容易把这段保护逻辑改回去，CI 也拦不住这类 OOB 回归。

建议补一个 CPU 动态图用例，至少验证 padding_idx=0 时越界 id 会抛 ValueError，例如：

paddle.disable_static(paddle.CPUPlace())
x = paddle.to_tensor([0, 6], dtype='int64')
w = paddle.randn([5, 3])
with self.assertRaises(ValueError):
    paddle.nn.functional.embedding(x, w, padding_idx=0)

处理要求：请针对该评论修复并提交新的 commit。

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)
Expand Down
21 changes: 21 additions & 0 deletions test/ai_edited_test/test_ai_nn_functional_input.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"""
Expand Down
Loading