[MXNET-953] Fix oob memory read (apache#12631)

azai91 · Dec 1, 2018 · 91125ff · 91125ff
1 parent 30ad9f6
commit 91125ff
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/src/operator/tensor/elemwise_unary_op_basic.cc b/src/operator/tensor/elemwise_unary_op_basic.cc
@@ -492,7 +492,8 @@ void ShapeComputeCPU(const nnvm::NodeAttrs& attrs,
   CHECK_EQ(req.size(), 1U);
   const TBlob& in_data = inputs[0];
   const TBlob& out_data = outputs[0];
-  memcpy(out_data.dptr_, in_data.shape_.data(), in_data.ndim() * sizeof(int64_t));
+  size_t type_size = mshadow::mshadow_sizeof(out_data.type_flag_);
+  memcpy(out_data.dptr_, in_data.shape_.data(), in_data.ndim() * type_size);
 }
 
 NNVM_REGISTER_OP(shape_array)
@@ -542,8 +543,9 @@ void SizeComputeCPU(const nnvm::NodeAttrs& attrs,
   CHECK_EQ(req.size(), 1U);
   const TBlob& in_data = inputs[0];
   const TBlob& out_data = outputs[0];
+  size_t type_size = mshadow::mshadow_sizeof(out_data.type_flag_);
   const index_t size_var = in_data.Size();
-  memcpy(out_data.dptr_, &size_var, 1U * sizeof(int64_t));
+  memcpy(out_data.dptr_, &size_var, type_size);
 }
 
 NNVM_REGISTER_OP(size_array)