Star79
24
代码
代码
Issues3
Pull Requests1
项目讨论
Wiki
分析
Star79
24
已合并
[bugfix] 重明安全排查问题修改 #62
zhaishangzhao创建于 6月30日
[bugfix] 重明安全排查问题修改 #62
已合并
zhaishangzhao创建于 6月30日
从zsz_wlf/mstxbr_test合入到Ascend/mstx26.1.0
共 1 个文件变更+36-19
Mpython/mstx_api.cpp+36-19
@@ -20,20 +20,19 @@
20class GILCtrl {20class GILCtrl {
21public:21public:
22 GILCtrl() : tstate(PyEval_SaveThread()) {} // 释放GIL22 GILCtrl() : tstate(PyEval_SaveThread()) {} // 释放GIL
23- ~GILCtrl()23+ ~GILCtrl() {
24- {24+ PyEval_RestoreThread(tstate); // 恢复GIL
25- PyEval_RestoreThread(tstate); // 恢复GIL
26 }25 }
26+ 
27private:27private:
28- PyThreadState* tstate;28+ PyThreadState *tstate;
29};29};
30 30 
31-bool ParseArgs(PyObject *args, PyObject *kwds, const char*& message, PyObject*& py_stream)31+bool ParseArgs(PyObject *args, PyObject *kwds, const char *&message, PyObject *&py_stream) {
32-{
33 message = nullptr;32 message = nullptr;
34 py_stream = Py_None;33 py_stream = Py_None;
35 34 
36- static char* kwlist[] = { "message", "stream", nullptr };35+ static char *kwlist[] = {"message", "stream", nullptr};
37 36 
38 if (!PyArg_ParseTupleAndKeywords(args, kwds, "|sO", kwlist, &message, &py_stream)) {37 if (!PyArg_ParseTupleAndKeywords(args, kwds, "|sO", kwlist, &message, &py_stream)) {
39 return false;38 return false;
@@ -41,8 +40,7 @@ bool ParseArgs(PyObject *args, PyObject *kwds, const char*& message, PyObject*&
41 return true;40 return true;
42}41}
43 42 
44-PyObject *WrapMstxMarkA(PyObject *self, PyObject *args, PyObject *kwds)43+PyObject *WrapMstxMarkA(PyObject *self, PyObject *args, PyObject *kwds) {
45-{
46 const char *message;44 const char *message;
47 PyObject *py_stream;45 PyObject *py_stream;
48 46 
@@ -52,7 +50,18 @@ PyObject *WrapMstxMarkA(PyObject *self, PyObject *args, PyObject *kwds)
52 50 
53 aclrtStream stream = nullptr;51 aclrtStream stream = nullptr;
54 if (py_stream != Py_None) {52 if (py_stream != Py_None) {
55- void* ptr = reinterpret_cast<void*>(PyLong_AsVoidPtr(py_stream));53+ // 校验输入必须是整数
54+ if (!PyLong_Check(py_stream)) {
55+ PyErr_SetString(PyExc_TypeError, "stream must be an integer (aclrtStream handle) or None");
56+ return nullptr;
57+ }
58+ 
59+ unsigned long long handle_val = PyLong_AsUnsignedLongLong(py_stream);
60+ if (PyErr_Occurred()) {
61+ return nullptr;
62+ }
63+ 
64+ void *ptr = reinterpret_cast<void *>(handle_val);
56 stream = static_cast<aclrtStream>(ptr);65 stream = static_cast<aclrtStream>(ptr);
57 }66 }
58 if (!mstxMarkA) {67 if (!mstxMarkA) {
@@ -66,9 +75,8 @@ PyObject *WrapMstxMarkA(PyObject *self, PyObject *args, PyObject *kwds)
66 Py_RETURN_NONE;75 Py_RETURN_NONE;
67}76}
68 77 
69-PyObject *WrapMstxRangeStartA(PyObject *self, PyObject *args, PyObject *kwds)78+PyObject *WrapMstxRangeStartA(PyObject *self, PyObject *args, PyObject *kwds) {
70-{79+ const char *message = nullptr;
71- const char* message = nullptr;
72 PyObject *py_stream = Py_None;80 PyObject *py_stream = Py_None;
73 81 
74 if (!ParseArgs(args, kwds, message, py_stream)) {82 if (!ParseArgs(args, kwds, message, py_stream)) {
@@ -77,7 +85,18 @@ PyObject *WrapMstxRangeStartA(PyObject *self, PyObject *args, PyObject *kwds)
77 85 
78 aclrtStream stream = nullptr;86 aclrtStream stream = nullptr;
79 if (py_stream != Py_None) {87 if (py_stream != Py_None) {
80- void* ptr = reinterpret_cast<void*>(PyLong_AsVoidPtr(py_stream));88+ // 校验输入必须是整数
89+ if (!PyLong_Check(py_stream)) {
90+ PyErr_SetString(PyExc_TypeError, "stream must be an integer (aclrtStream handle) or None");
91+ return nullptr;
92+ }
93+ 
94+ unsigned long long handle_val = PyLong_AsUnsignedLongLong(py_stream);
95+ if (PyErr_Occurred()) {
96+ return nullptr;
97+ }
98+ 
99+ void *ptr = reinterpret_cast<void *>(handle_val);
81 stream = static_cast<aclrtStream>(ptr);100 stream = static_cast<aclrtStream>(ptr);
82 }101 }
83 if (!mstxRangeStartA) {102 if (!mstxRangeStartA) {
@@ -92,8 +111,7 @@ PyObject *WrapMstxRangeStartA(PyObject *self, PyObject *args, PyObject *kwds)
92 return Py_BuildValue("I", ret);111 return Py_BuildValue("I", ret);
93}112}
94 113 
95-PyObject *WrapMstxRangeEnd(PyObject *self, PyObject *args, PyObject *kwds)114+PyObject *WrapMstxRangeEnd(PyObject *self, PyObject *args, PyObject *kwds) {
96-{
97 uint32_t rangeId = 0;115 uint32_t rangeId = 0;
98 static char arg1[] = "rangeId";116 static char arg1[] = "rangeId";
99 static char *kwlist[] = {arg1, nullptr};117 static char *kwlist[] = {arg1, nullptr};
@@ -112,12 +130,11 @@ PyObject *WrapMstxRangeEnd(PyObject *self, PyObject *args, PyObject *kwds)
112 Py_RETURN_NONE;130 Py_RETURN_NONE;
113}131}
114 132 
115-PyObject *WrapMstxGetToolId(PyObject *self, PyObject *args, PyObject *kwds)133+PyObject *WrapMstxGetToolId(PyObject *self, PyObject *args, PyObject *kwds) {
116-{
117 uint64_t id = 0;134 uint64_t id = 0;
118 {135 {
119 GILCtrl gilCtrl;136 GILCtrl gilCtrl;
120 mstxGetToolId(&id);137 mstxGetToolId(&id);
121 }138 }
122 return Py_BuildValue("K", id);139 return Py_BuildValue("K", id);
123-}140+}