gemini-code-assist[bot] commented on code in PR #496:
URL: https://github.com/apache/tvm-ffi/pull/496#discussion_r2880815988
##########
rust/tvm-ffi/src/collections/tensor.rs:
##########
@@ -158,7 +158,11 @@ impl Tensor {
///
/// # Returns
/// * `Result<&mut [T]>` - The data as a mutable slice
- pub fn data_as_slice_mut<T: AsDLDataType>(&self) -> Result<&mut [T]> {
+ /// Returns the tensor data as a mutable slice.
+ ///
+ /// Note: if the `Tensor` has been cloned (via `ObjectArc`), the caller
+ /// must ensure no other clone is concurrently reading the data.
+ pub fn data_as_slice_mut<T: AsDLDataType>(&mut self) -> Result<&mut [T]> {
Review Comment:
The function `data_as_slice_mut` has two significant security issues:
1. **Unsoundness (Aliased Mutable References):** Although the signature is
changed to `&mut self`, the `Tensor` type implements `Clone` (via `ObjectArc`),
which allows multiple `Tensor` instances to point to the same underlying data.
This means safe code can still obtain multiple aliased `&mut [T]` references to
the same memory by cloning the `Tensor` and calling this function on each
clone. This is undefined behavior in Rust. The function should be marked
`unsafe fn` to require the caller to uphold the invariant that no other
references to the data exist.
2. **Null Pointer Dereference:** The function calls
`std::slice::from_raw_parts_mut` using the underlying data pointer without
verifying it is not null. If memory allocation fails, the pointer will be null,
leading to undefined behavior when `from_raw_parts_mut` is called with a
non-zero length. A check should be added to ensure `self.data.dltensor.data` is
not null when `self.numel() > 0`.
##########
rust/tvm-ffi/src/device.rs:
##########
@@ -26,8 +26,12 @@ use tvm_ffi_sys::{TVMFFIEnvGetStream, TVMFFIEnvSetStream,
TVMFFIStreamHandle};
pub fn current_stream(device: &DLDevice) -> TVMFFIStreamHandle {
unsafe { TVMFFIEnvGetStream(device.device_type as i32, device.device_id) }
}
-/// call f with the stream set to the given stream
-pub fn with_stream<T>(
+/// Call `f` with the device stream temporarily set to `stream`.
+///
+/// # Safety
+///
+/// `stream` must be a valid stream handle for the given device, or null.
+pub unsafe fn with_stream<T>(
Review Comment:
The function `with_stream` fails to restore the previous stream if the
provided closure `f` returns an error. The `?` operator on line 48 causes an
early return, skipping the restoration logic in the subsequent `unsafe` block.
This leaves the device in an inconsistent state with the temporary stream still
set, which can lead to race conditions or unexpected behavior in other parts of
the application that expect the original stream. It is recommended to use a
guard object (RAII pattern) to ensure the stream is restored even on error or
panic.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
