From a63b88ef3c8f58c7c0ac57d361d06ff0bb3ed385 Mon Sep 17 00:00:00 2001
From: kieran <kieran@harkin.me>
Date: Mon, 27 Jan 2025 22:50:43 +0000
Subject: [PATCH] fix: leaky demuxer file handles

---
 Cargo.lock   | 12 +++++++++++-
 Cargo.toml   |  1 +
 src/demux.rs | 32 +++++++++++++++++++++++++-------
 3 files changed, 37 insertions(+), 8 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index cb79ad5..ef6cfa6 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1,6 +1,6 @@
 # This file is automatically @generated by Cargo.
 # It is not intended for manual editing.
-version = 3
+version = 4
 
 [[package]]
 name = "aho-corasick"
@@ -181,6 +181,7 @@ dependencies = [
  "ffmpeg-sys-the-third",
  "libc",
  "log",
+ "rlimit",
  "slimbox",
 ]
 
@@ -333,6 +334,15 @@ version = "0.8.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "2b15c43186be67a4fd63bee50d0303afffcef381492ebe2c5d87f324e1b8815c"
 
+[[package]]
+name = "rlimit"
+version = "0.10.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7043b63bd0cd1aaa628e476b80e6d4023a3b50eb32789f2728908107bd0c793a"
+dependencies = [
+ "libc",
+]
+
 [[package]]
 name = "rustc-hash"
 version = "1.1.0"
diff --git a/Cargo.toml b/Cargo.toml
index 8c599d8..bbe96f5 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -17,3 +17,4 @@ log = "0.4.22"
 
 [dev-dependencies]
 env_logger = "0.11.5"
+rlimit = "0.10.2"
diff --git a/src/demux.rs b/src/demux.rs
index 2165045..b23dce8 100644
--- a/src/demux.rs
+++ b/src/demux.rs
@@ -250,14 +250,15 @@ impl Drop for Demuxer {
     fn drop(&mut self) {
         unsafe {
             if !self.ctx.is_null() {
-                if let DemuxerInput::Reader(_, _) = self.input {
-                    av_free((*(*self.ctx).pb).buffer as *mut _);
-                    drop(SlimBox::<dyn Read>::from_raw((*(*self.ctx).pb).opaque));
+                match self.input {
+                    DemuxerInput::Reader(_, _) => {
+                        av_free((*(*self.ctx).pb).buffer as *mut _);
+                        drop(SlimBox::<dyn Read>::from_raw((*(*self.ctx).pb).opaque));
+                        avio_context_free(&mut (*self.ctx).pb);
+                    }
+                    _ => {}
                 }
-                if !(*self.ctx).pb.is_null() {
-                    avio_context_free(&mut (*self.ctx).pb);
-                }
-                avformat_free_context(self.ctx);
+                avformat_close_input(&mut self.ctx);
             }
         }
     }
@@ -284,4 +285,21 @@ mod tests {
         }
         Ok(())
     }
+
+    /// Test for leaking file handles
+    #[test]
+    fn probe_lots() -> Result<()> {
+        rlimit::setrlimit(rlimit::Resource::NOFILE, 64, 128)?;
+
+        let nof_limit = rlimit::Resource::NOFILE.get_hard()?;
+        for n in 0..nof_limit {
+            let mut demux = Demuxer::new("./test_output/test.png")?;
+            unsafe {
+                if let Err(e) = demux.probe_input() {
+                    bail!("Failed on {}: {}", n, e);
+                }
+            }
+        }
+        Ok(())
+    }
 }