Skip to content

Commit 36051ef

Browse files
authored
Add size check on partial_ops x input vector (#34595)
### Details: - *Add input count validation in partial_ops() to prevent out-of-bounds access* ### Tickets: - *CVS-182184*
1 parent b8ca92f commit 36051ef

File tree

3 files changed

+60
-0
lines changed

3 files changed

+60
-0
lines changed

src/frontends/paddle/src/op/partial_ops.hpp

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ namespace op {
1212

1313
NamedOutputs partial_ops(const NodeContext& node, const std::string type) {
1414
auto x = node.get_ng_inputs("X");
15+
PADDLE_OP_CHECK(node, x.size() == 2, "partial_ops requires exactly 2 inputs in X.");
1516
const auto start_index = node.get_attribute<int>("start_index");
1617
const auto length = node.get_attribute<int>("length");
1718
PADDLE_OP_CHECK(node, x[0].get_partial_shape().rank().get_length() == 2, "partial ops only support 2-D Tensor");

src/frontends/paddle/tests/read_paddle_model_test.cpp

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -307,3 +307,18 @@ TEST(Paddle_Reader_Tests, ImportBasicModelToCoreWstring) {
307307
ASSERT_TRUE(res.valid) << res.message;
308308
}
309309
#endif
310+
311+
TEST(Paddle_Reader_Tests, LoadModelWithPartialOpsInsufficientInputs) {
312+
auto model =
313+
FrontEndTestUtils::make_model_path(std::string(TEST_PADDLE_MODELS_DIRNAME) + "partial_sum_oob/partial_sum_oob" +
314+
std::string(TEST_PADDLE_MODEL_EXT));
315+
316+
ov::Core core;
317+
try {
318+
core.read_model(model);
319+
FAIL() << "Expected load to fail due to insufficient X inputs for partial_sum";
320+
} catch (const std::exception& ex) {
321+
const std::string msg = ex.what();
322+
ASSERT_NE(msg.find("partial_ops requires exactly 2 inputs in X."), std::string::npos) << msg;
323+
}
324+
}

src/frontends/paddle/tests/test_models/gen_scripts/generate_partial_sum.py

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,11 +5,48 @@
55
# partial_sum paddle model generator
66
#
77
import numpy as np
8+
import os
89
from save_model import saveModel
910
import paddle
1011
import sys
1112

1213

14+
def _get_framework_pb2():
15+
try:
16+
from paddle.fluid.proto import framework_pb2
17+
return framework_pb2
18+
except Exception:
19+
pass
20+
try:
21+
from paddle.base.proto import framework_pb2
22+
return framework_pb2
23+
except Exception:
24+
from paddle.framework.proto import framework_pb2
25+
return framework_pb2
26+
27+
28+
def _corrupt_partial_sum_inputs(model_path: str):
29+
framework_pb2 = _get_framework_pb2()
30+
prog = framework_pb2.ProgramDesc()
31+
with open(model_path, "rb") as f:
32+
prog.ParseFromString(f.read())
33+
34+
modified = False
35+
for block in prog.blocks:
36+
for op in block.ops:
37+
if op.type == "partial_sum":
38+
for inp in op.inputs:
39+
if inp.parameter == "X" and len(inp.arguments) > 1:
40+
del inp.arguments[1:]
41+
modified = True
42+
43+
if not modified:
44+
raise RuntimeError("Failed to modify partial_sum X inputs in model")
45+
46+
with open(model_path, "wb") as f:
47+
f.write(prog.SerializeToString())
48+
49+
1350
def partial_sum(name: str, x, y, start_index=0, length=-1):
1451
paddle.enable_static()
1552

@@ -64,5 +101,12 @@ def main():
64101
y = np.random.randint(-10, 10, [8, 10]).astype(dtype)
65102
partial_sum("partial_sum_3", x, y, start_index=1, length=5)
66103

104+
dtype = 'float32'
105+
x = np.random.randn(4, 8).astype(dtype)
106+
y = np.random.randn(4, 8).astype(dtype)
107+
partial_sum("partial_sum_oob", x, y, start_index=0, length=4)
108+
model_path = os.path.join(sys.argv[1], "partial_sum_oob", "partial_sum_oob.pdmodel")
109+
_corrupt_partial_sum_inputs(model_path)
110+
67111
if __name__ == "__main__":
68112
main()

0 commit comments

Comments
 (0)