Skip to content

Commit 292bc2a

Browse files
committed
Fix possible overflows in _usmarray.pyx
When a bad typenum is passed to functions which make usm_ndarrays from pointers, there was a possible underflow, as itemsize was declared size_t Instead declare itemsize as int and cast to size_t when needed (as itemsize should always fit in int)
1 parent a34686f commit 292bc2a

File tree

1 file changed

+8
-4
lines changed

1 file changed

+8
-4
lines changed

dpctl/tensor/_usmarray.pyx

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1777,8 +1777,10 @@ cdef api object UsmNDArray_MakeSimpleFromPtr(
17771777
Returns:
17781778
Created usm_ndarray instance
17791779
"""
1780-
cdef size_t itemsize = type_bytesize(typenum)
1781-
cdef size_t nbytes = itemsize * nelems
1780+
cdef int itemsize = type_bytesize(typenum)
1781+
if (itemsize < 1):
1782+
raise ValueError("dtype with typenum=" + str(typenum) + " is not supported.")
1783+
cdef size_t nbytes = (<size_t> itemsize) * nelems
17821784
cdef c_dpmem._Memory mobj = c_dpmem._Memory.create_from_usm_pointer_size_qref(
17831785
ptr, nbytes, QRef, memory_owner=owner
17841786
)
@@ -1817,7 +1819,7 @@ cdef api object UsmNDArray_MakeFromPtr(
18171819
Returns:
18181820
Created usm_ndarray instance
18191821
"""
1820-
cdef size_t itemsize = type_bytesize(typenum)
1822+
cdef int itemsize = type_bytesize(typenum)
18211823
cdef int err = 0
18221824
cdef size_t nelems = 1
18231825
cdef Py_ssize_t min_disp = 0
@@ -1830,6 +1832,8 @@ cdef api object UsmNDArray_MakeFromPtr(
18301832
cdef object obj_shape
18311833
cdef object obj_strides
18321834

1835+
if (itemsize < 1):
1836+
raise ValueError("dtype with typenum=" + str(typenum) + " is not supported.")
18331837
if (nd < 0):
18341838
raise ValueError("Dimensionality must be non-negative")
18351839
if (ptr is NULL or QRef is NULL):
@@ -1881,7 +1885,7 @@ cdef api object UsmNDArray_MakeFromPtr(
18811885
raise ValueError(
18821886
"Given shape, strides and offset reference out-of-bound memory"
18831887
)
1884-
nbytes = itemsize * (offset + max_disp + 1)
1888+
nbytes = (<size_t> itemsize) * (offset + max_disp + 1)
18851889
mobj = c_dpmem._Memory.create_from_usm_pointer_size_qref(
18861890
ptr, nbytes, QRef, memory_owner=owner
18871891
)

0 commit comments

Comments
 (0)