lzma_rust2 binding: keep the stream usable after LZMA_MEMLIMIT_ERROR

lzma/lzma_rust_api.h

@@ -546,10 +546,10 @@ lzma_check lzma_get_check(const lzma_stream* strm);
 // For encoders and raw decoders where options are explicitly configured at
 // initialization, returns the calculated memory usage in bytes.
 //
-// Note: For streaming decoders where dictionary size is determined dynamically
+// For streaming decoders where dictionary size is determined dynamically
 // from the stream header (e.g., `lzma_stream_decoder`, `lzma_alone_decoder`,
-// `lzma_auto_decoder`, `lzma_lzip_decoder`), this implementation returns 0
-// as the underlying Rust decoders do not expose runtime memory usage.
+// `lzma_auto_decoder`, `lzma_lzip_decoder`), returns the memory usage in bytes
+// once the header has been parsed, and 0 before that.
 uint64_t lzma_memusage(const lzma_stream* strm);
 
 #ifdef __cplusplus

lzma/rust/code.rs

@@ -8,8 +8,9 @@ use lzma_rust2::Status;
 use crate::helpers::{self, FfiAction};
 use crate::state::{ActionState, CoderInner, CoderResult, InternalState};
 use crate::types::{
-    lzma_ret, lzma_stream, LZMA_BUF_ERROR, LZMA_CHECK_NONE, LZMA_GET_CHECK, LZMA_NO_CHECK, LZMA_OK,
-    LZMA_OPTIONS_ERROR, LZMA_PROG_ERROR, LZMA_STREAM_END, LZMA_TELL_ANY_CHECK, LZMA_TELL_NO_CHECK,
+    lzma_ret, lzma_stream, LZMA_BUF_ERROR, LZMA_CHECK_NONE, LZMA_GET_CHECK, LZMA_MEMLIMIT_ERROR,
+    LZMA_NO_CHECK, LZMA_OK, LZMA_OPTIONS_ERROR, LZMA_PROG_ERROR, LZMA_STREAM_END,
+    LZMA_TELL_ANY_CHECK, LZMA_TELL_NO_CHECK,
 };
 
 pub(crate) struct StreamView<'a> {
@@ -339,10 +340,17 @@ unsafe fn lzma_code_inner(
         }
         Err(e) => {
             state.allow_buf_error = false;
-            // NOTE(b/563249135): liblzma keeps reporting LZMA_MEMLIMIT_ERROR
-            // instead of failing the stream.
-            state.action_state = ActionState::Error;
-            helpers::map_rust_error(e)
+            let ret = helpers::map_rust_error(e);
+            if ret == LZMA_MEMLIMIT_ERROR {
+                // Like liblzma, a memory limit error leaves the stream usable: the decoder
+                // checks the limit again on the next call, from the same point in the input.
+                if let ActionState::Finish { .. } = state.action_state {
+                    state.action_state = ActionState::Finish { saved_avail_in: strm_ref.avail_in };
+                }
+            } else {
+                state.action_state = ActionState::Error;
+            }
+            ret
         }
     }
 }

lzma/rust/state.rs

@@ -54,7 +54,9 @@ fn decoder_result(
             bytes_produced: r.bytes_produced,
             result: Ok(r.status),
         },
-        Err(e) if e.kind() != io::ErrorKind::UnexpectedEof => {
+        Err(e)
+            if !matches!(e.kind(), io::ErrorKind::UnexpectedEof | io::ErrorKind::OutOfMemory) =>
+        {
             CoderResult { bytes_consumed: 0, bytes_produced: 0, result: Err(e) }
         }
         Err(e) => {
@@ -201,15 +203,20 @@ impl InternalState {
     }
 
     pub fn memusage(&self) -> u64 {
-        match &self.coder {
+        // The auto decoder answers for the format it detected.
+        let coder = match &self.coder {
+            CoderInner::AutoDecoder(auto) => auto.inner().unwrap_or(&self.coder),
+            coder => coder,
+        };
+        match coder {
             CoderInner::RawLzma2Decoder(_) => {
                 if let Some(opts) = &self.lzma_options {
                     let dict_size = opts.dict_size.max(lzma_rust2::DICT_SIZE_MIN);
                     return (lzma_rust2::lzma2_get_memory_usage(dict_size) as u64) * 1024;
                 }
                 0
             }
-            CoderInner::LzmaDecoder(_) => {
+            CoderInner::LzmaDecoder(lzma) => {
                 if let Some(opts) = &self.lzma_options {
                     let dict_size = opts.dict_size.max(lzma_rust2::DICT_SIZE_MIN);
                     if let Ok(mem_kb) =
@@ -218,10 +225,12 @@ impl InternalState {
                         return (mem_kb as u64) * 1024;
                     }
                 }
-                0
+                lzma.memory_usage().map_or(0, |kb| u64::from(kb) * 1024)
+            }
+            CoderInner::XzDecoder(xz) => xz.memory_usage().map_or(0, |kb| u64::from(kb) * 1024),
+            CoderInner::LzipDecoder(lzip) => {
+                lzip.memory_usage().map_or(0, |kb| u64::from(kb) * 1024)
             }
-            // NOTE (b/562586054): Add support for lzma_stream_decoder,
-            // lzma_auto_decoder, lzma_lzip_decoder, and lzma_alone_decoder.
             _ => 0,
         }
     }

lzma/rust/util.rs

@@ -27,15 +27,15 @@ pub extern "C" fn lzma_cputhreads() -> u32 {
 ///
 /// Implements `lzma_memusage` by interrogating the rust state struct.
 /// Returns memory usage for encoders and raw decoders where options are known.
-/// For streaming decoders (XZ, Alone, Lzip, Auto), returns 0 because
-/// the underlying decoders do not expose dynamically parsed dictionary size.
+/// For streaming decoders (XZ, Alone, Lzip, Auto), returns the figure the crate
+/// checks against the memory limit, which is 0 until the header has been parsed.
 ///
 /// # Safety
 /// - `strm` must be null or point to a valid, properly aligned `lzma_stream`.
 /// - If `strm` is non-null and `strm.internal` is non-null, it must point to a valid
 ///   `InternalState` allocated by this crate.
-/// - The caller must guarantee that `*strm` and `*strm.internal` are not concurrently modified
-///   or deallocated during this call.
+/// - The caller must guarantee that `*strm` and `*strm.internal` are not concurrently modified or
+///   deallocated during this call.
 #[unsafe(export_name = crate::prefix!(lzma_memusage))]
 pub unsafe extern "C" fn lzma_memusage(strm: Option<&crate::types::lzma_stream>) -> u64 {
     let Some(strm) = strm else {