feat(sdk): falls back to fs for large intermediate streams (#1341)

Co-authored-by: Dylan Ross <[email protected]>

Author: dyro

Cargo.lock

@@ -177,7 +177,7 @@ serde_json = { version = "1.0.117", features = ["preserve_order"] }
 serde_with = "3.11.0"
 serde-transcode = "1.1.1"
 sha1 = "0.10.6"
-sha2 = "0.10.6"
+sha2 = { version = "0.10.6", features = ["asm"] }
 spki = { version = "0.7.3", optional = true }
 static-iref = "3.0"
 tempfile = "=3.15.0"

sdk/Cargo.toml

@@ -173,6 +173,7 @@ pub(crate) struct Core {
     compress_manifests: bool,
     #[serde(skip_serializing_if = "Option::is_none")]
     max_memory_usage: Option<u64>,
+    backing_store_memory_threshold_in_mb: usize,
     // TODO: pending https://github.com/contentauth/c2pa-rs/pull/1180
     // prefer_update_manifests: bool,
 }
@@ -189,6 +190,7 @@ impl Default for Core {
             merkle_tree_max_proofs: 5,
             compress_manifests: true,
             max_memory_usage: None,
+            backing_store_memory_threshold_in_mb: 512,
             // prefer_update_manifests: true,
         }
     }

sdk/src/settings/mod.rs

@@ -84,7 +84,8 @@ use crate::{
     status_tracker::{ErrorBehavior, StatusTracker},
     utils::{
         hash_utils::HashRange,
-        io_utils::{insert_data_at, safe_vec, stream_len},
+        io_utils,
+        io_utils::{insert_data_at, stream_len},
         is_zero,
         patch::patch_bytes,
     },
@@ -98,7 +99,6 @@ use crate::{
 use crate::{external_manifest::ManifestPatchCallback, RemoteSigner};
 
 const MANIFEST_STORE_EXT: &str = "c2pa"; // file extension for external manifests
-const MANIFEST_RESERVE_SIZE: usize = 10 * 1024 * 1024; // 10MB reserve size for manifest
 
 pub(crate) struct ManifestHashes {
     pub manifest_box_hash: Vec<u8>,
@@ -2788,11 +2788,7 @@ impl Store {
                 .await?
         };
 
-        let intermediate_output: Vec<u8> = safe_vec(
-            stream_len(input_stream)? + MANIFEST_RESERVE_SIZE as u64,
-            None,
-        )?;
-        let mut intermediate_stream = Cursor::new(intermediate_output);
+        let mut intermediate_stream = io_utils::stream_with_fs_fallback(None)?;
 
         #[allow(unused_mut)] // Not mutable in the non-async case.
         let mut jumbf_bytes = self.start_save_stream(
@@ -3314,12 +3310,7 @@ impl Store {
         output_stream: &mut dyn CAIReadWrite,
         reserve_size: usize,
     ) -> Result<Vec<u8>> {
-        // make sure we can hold the intermediate stream before attempting
-        let intermediate_output: Vec<u8> = safe_vec(
-            stream_len(input_stream)? + MANIFEST_RESERVE_SIZE as u64,
-            None,
-        )?;
-        let mut intermediate_stream = Cursor::new(intermediate_output);
+        let mut intermediate_stream = io_utils::stream_with_fs_fallback(None)?;
 
         let pc = self.provenance_claim_mut().ok_or(Error::ClaimEncoding)?;
 
@@ -3345,11 +3336,7 @@ impl Store {
                     .get_writer(format)
                     .ok_or(Error::UnsupportedType)?;
 
-                let tmp_output: Vec<u8> = safe_vec(
-                    stream_len(input_stream)? + MANIFEST_RESERVE_SIZE as u64,
-                    None,
-                )?;
-                let mut tmp_stream = Cursor::new(tmp_output);
+                let mut tmp_stream = io_utils::stream_with_fs_fallback(None)?;
                 manifest_writer.remove_cai_store_from_stream(input_stream, &mut tmp_stream)?;
 
                 // add external ref if possible
@@ -3393,11 +3380,7 @@ impl Store {
 
                 // insert UUID boxes at the correct location if required
                 if let Some(merkle_uuid_boxes) = &bmff_hash.merkle_uuid_boxes {
-                    let temp_data: Vec<u8> = safe_vec(
-                        stream_len(&mut intermediate_stream)? + MANIFEST_RESERVE_SIZE as u64,
-                        None,
-                    )?;
-                    let mut temp_stream = Cursor::new(temp_data);
+                    let mut temp_stream = io_utils::stream_with_fs_fallback(None)?;
 
                     insert_data_at(
                         &mut intermediate_stream,

sdk/src/store.rs

@@ -20,9 +20,9 @@ use std::{
 };
 
 #[allow(unused)] // different code path for WASI
-use tempfile::{tempdir, Builder, NamedTempFile, TempDir};
+use tempfile::{tempdir, Builder, NamedTempFile, SpooledTempFile, TempDir};
 
-use crate::{asset_io::rename_or_move, Error, Result};
+use crate::{asset_io::rename_or_move, settings::get_settings_value, Error, Result};
 // Replace data at arbitrary location and len in a file.
 // start_location is where the replacement data will start
 // replace_len is how many bytes from source to replaced starting a start_location
@@ -118,6 +118,44 @@ pub(crate) fn stream_len<R: Read + Seek + ?Sized>(reader: &mut R) -> Result<u64>
     Ok(len)
 }
 
+#[cfg(target_arch = "wasm32")]
+fn stream_with_fs_fallback_wasm(
+    _threshold_override: Option<usize>,
+) -> Result<std::io::Cursor<Vec<u8>>> {
+    Ok(std::io::Cursor::new(Vec::new()))
+}
+
+#[cfg(not(target_arch = "wasm32"))]
+fn stream_with_fs_fallback_file_io(threshold_override: Option<usize>) -> Result<SpooledTempFile> {
+    let threshold = threshold_override.unwrap_or(get_settings_value::<usize>(
+        "core.backing_store_memory_threshold_in_mb",
+    )?);
+
+    Ok(SpooledTempFile::new(threshold))
+}
+
+/// Will create a [Read], [Write], and [Seek] capable stream that will stay in memory
+/// as long as the threshold is not exceeded. The threshold is specified in MB in the
+/// settings under ""core.backing_store_memory_threshold_in_mb"
+///
+/// # Parameters
+/// - `threshold_override`: Optional override for the threshold value in MB. If provided, this
+///   value will be used instead of the one from settings.
+///
+/// # Errors
+/// - Returns an error if the threshold value from settings is not valid.
+///
+/// # Note
+/// This will return a an in-memory stream when the compilation target doesn't support file I/O.
+pub(crate) fn stream_with_fs_fallback(
+    threshold_override: Option<usize>,
+) -> Result<impl Read + Write + Seek> {
+    #[cfg(target_arch = "wasm32")]
+    return stream_with_fs_fallback_wasm(threshold_override);
+    #[cfg(not(target_arch = "wasm32"))]
+    return stream_with_fs_fallback_file_io(threshold_override);
+}
+
 // Returns a new Vec first making sure it can hold the desired capacity.  Fill
 // with default value if provided
 pub(crate) fn safe_vec<T: Clone>(item_cnt: u64, init_with: Option<T>) -> Result<Vec<T>> {
@@ -386,4 +424,41 @@ mod tests {
         )
         .is_err());
     }
+
+    #[cfg(not(target_arch = "wasm32"))]
+    #[test]
+    fn test_safe_stream_threshold_behavior() {
+        let mut stream = stream_with_fs_fallback_file_io(Some(10)).unwrap();
+
+        // Less data written than required to write to the FS.
+        let small_data = b"small"; // 5 bytes
+        stream.write_all(small_data).unwrap();
+        assert!(!stream.is_rolled(), "data still in memory");
+
+        // Adds more data to exceed the threshold.
+        let large_data = b"this is larger than 10 bytes total";
+        stream.write_all(large_data).unwrap();
+        assert!(stream.is_rolled(), "data moved to disk");
+    }
+
+    #[cfg(not(target_arch = "wasm32"))]
+    #[test]
+    fn test_safe_stream_no_threshold_behavior() {
+        let mut stream = stream_with_fs_fallback_file_io(None).unwrap();
+
+        // Less data written than required to write to the FS.
+        let small_data = b"small"; // 5 bytes
+        stream.write_all(small_data).unwrap();
+        assert!(!stream.is_rolled(), "data still in memory");
+
+        let large_data = vec![0; 1024 * 1024]; // 1MB.
+        let threshold =
+            get_settings_value::<usize>("core.backing_store_memory_threshold_in_mb").unwrap();
+
+        for _ in 0..threshold {
+            stream.write_all(&large_data).unwrap();
+        }
+
+        assert!(stream.is_rolled(), "data moved to disk");
+    }
 }