From 92aa92fa8474087da2d874363d43ef48ee702604 Mon Sep 17 00:00:00 2001
From: wanghuancoder <wanghuan29@baidu.com>
Date: Wed, 2 Aug 2023 15:10:38 +0800
Subject: [PATCH] fix security bug (#55866)

* fix security bug
---
 paddle/phi/infermeta/unary.cc            | 14 +++++++++++++-
 test/legacy_test/test_zero_dim_tensor.py | 15 ---------------
 2 files changed, 13 insertions(+), 16 deletions(-)

diff --git a/paddle/phi/infermeta/unary.cc b/paddle/phi/infermeta/unary.cc
index 417b6f358fc..fb5cb870cad 100644
--- a/paddle/phi/infermeta/unary.cc
+++ b/paddle/phi/infermeta/unary.cc
@@ -2272,7 +2272,19 @@ void ModeInferMeta(const MetaTensor& x,
       input_dims.size(),
       0,
       errors::InvalidArgument("input of ModeOp must have >= 0d shape"));
-  if (axis < 0) axis += dim_size;
+  if (axis < 0) {
+    axis += dim_size;
+    if (dim_size != 0) {
+      PADDLE_ENFORCE_GE(axis,
+                        0,
+                        phi::errors::InvalidArgument(
+                            "the axis must be [-%d, %d), but received %d .",
+                            dim_size,
+                            dim_size,
+                            axis - dim_size));
+    }
+  }
+
   std::vector<int64_t> dimvec;
   for (int64_t i = 0; i < axis; i++) {
     dimvec.emplace_back(input_dims[i]);
diff --git a/test/legacy_test/test_zero_dim_tensor.py b/test/legacy_test/test_zero_dim_tensor.py
index 7c814a83927..2f02a00209f 100644
--- a/test/legacy_test/test_zero_dim_tensor.py
+++ b/test/legacy_test/test_zero_dim_tensor.py
@@ -1245,21 +1245,6 @@ class TestSundryAPI(unittest.TestCase):
         self.assertEqual(x1.grad.shape, [5])
 
     def test_mode(self):
-        # 1) x is 0D
-        x = paddle.randn([])
-        x.stop_gradient = False
-        out, index = paddle.mode(x)
-        out.backward()
-
-        self.assertEqual(out.shape, [])
-        self.assertEqual(out, x)
-        self.assertEqual(index.shape, [])
-        self.assertEqual(index, 0)
-
-        self.assertEqual(x.grad.shape, [])
-        self.assertEqual(x.grad, 1.0)
-
-        # 2) x is 1D
         x1 = paddle.randn([5])
         x1.stop_gradient = False
         out1, index1 = paddle.mode(x1)
-- 
GitLab