From d5c78e7ba45fcebfbafd55a82ba2601ee3ea9617 Mon Sep 17 00:00:00 2001
From: Qqwy / Marten <qqwy@gmx.com>
Date: Thu, 27 Jul 2023 15:10:00 +0200
Subject: [PATCH] ArrowArrayStreamReader::try_new(): Safeguard against released
 streams (#1501)

---
 src/ffi/stream.rs      |  9 ++++++++-
 tests/it/ffi/stream.rs | 13 ++++++++++++-
 2 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/src/ffi/stream.rs b/src/ffi/stream.rs
index 96119358213..08fcaf0f43a 100644
--- a/src/ffi/stream.rs
+++ b/src/ffi/stream.rs
@@ -54,7 +54,8 @@ pub struct ArrowArrayStreamReader<Iter: DerefMut<Target = ArrowArrayStream>> {
 impl<Iter: DerefMut<Target = ArrowArrayStream>> ArrowArrayStreamReader<Iter> {
     /// Returns a new [`ArrowArrayStreamReader`]
     /// # Error
-    /// Errors iff the [`ArrowArrayStream`] is out of specification
+    /// Errors iff the [`ArrowArrayStream`] is out of specification,
+    /// or was already released prior to calling this function.
     /// # Safety
     /// This method is intrinsically `unsafe` since it assumes that the `ArrowArrayStream`
     /// contains a valid Arrow C stream interface.
@@ -62,6 +63,12 @@ impl<Iter: DerefMut<Target = ArrowArrayStream>> ArrowArrayStreamReader<Iter> {
     /// * The `ArrowArrayStream` fulfills the invariants of the C stream interface
     /// * The schema `get_schema` produces fulfills the C data interface
     pub unsafe fn try_new(mut iter: Iter) -> Result<Self, Error> {
+        if iter.release.is_none() {
+            return Err(Error::InvalidArgumentError(
+                "The C stream was already released".to_string(),
+            ));
+        };
+
         if iter.get_next.is_none() {
             return Err(Error::OutOfSpec(
                 "The C stream MUST contain a non-null get_next".to_string(),
diff --git a/tests/it/ffi/stream.rs b/tests/it/ffi/stream.rs
index 44d0e1e7cc1..53887d4362f 100644
--- a/tests/it/ffi/stream.rs
+++ b/tests/it/ffi/stream.rs
@@ -1,6 +1,6 @@
 use arrow2::array::*;
 use arrow2::datatypes::Field;
-use arrow2::{error::Result, ffi};
+use arrow2::{error::Error, error::Result, ffi};
 
 fn _test_round_trip(arrays: Vec<Box<dyn Array>>) -> Result<()> {
     let field = Field::new("a", arrays[0].data_type().clone(), true);
@@ -30,3 +30,14 @@ fn round_trip() -> Result<()> {
 
     _test_round_trip(vec![array.clone(), array.clone(), array])
 }
+
+#[test]
+fn stream_reader_try_new_invalid_argument_error_on_released_stream() {
+    let released_stream = Box::new(ffi::ArrowArrayStream::empty());
+    let reader = unsafe { ffi::ArrowArrayStreamReader::try_new(released_stream) };
+    // poor man's assert_matches:
+    match reader {
+        Err(Error::InvalidArgumentError(_)) => {}
+        _ => panic!("ArrowArrayStreamReader::try_new did not return an InvalidArgumentError"),
+    }
+}