expose memlimit settings for xz2 decoders (#252)

Nullus157 · Nov 15, 2023 · 37d106f · 37d106f
1 parent d75f8c2
commit 37d106f
Show file tree

Hide file tree

Showing 4 changed files with 45 additions and 7 deletions.
diff --git a/src/codec/lzma/decoder.rs b/src/codec/lzma/decoder.rs
@@ -10,7 +10,13 @@ pub struct LzmaDecoder {
 impl LzmaDecoder {
     pub fn new() -> Self {
         Self {
-            inner: crate::codec::Xz2Decoder::new(),
+            inner: crate::codec::Xz2Decoder::new(u64::max_value()),
+        }
+    }
+
+    pub fn with_memlimit(memlimit: u64) -> Self {
+        Self {
+            inner: crate::codec::Xz2Decoder::new(memlimit),
         }
     }
 }

diff --git a/src/codec/xz/decoder.rs b/src/codec/xz/decoder.rs
@@ -11,7 +11,14 @@ pub struct XzDecoder {
 impl XzDecoder {
     pub fn new() -> Self {
         Self {
-            inner: crate::codec::Xz2Decoder::new(),
+            inner: crate::codec::Xz2Decoder::new(u64::max_value()),
+            skip_padding: None,
+        }
+    }
+
+    pub fn with_memlimit(memlimit: u64) -> Self {
+        Self {
+            inner: crate::codec::Xz2Decoder::new(memlimit),
             skip_padding: None,
         }
     }

diff --git a/src/codec/xz2/decoder.rs b/src/codec/xz2/decoder.rs
@@ -15,16 +15,16 @@ impl Debug for Xz2Decoder {
 }
 
 impl Xz2Decoder {
-    pub fn new() -> Self {
+    pub fn new(mem_limit: u64) -> Self {
         Self {
-            stream: Stream::new_auto_decoder(u64::max_value(), 0).unwrap(),
+            stream: Stream::new_auto_decoder(mem_limit, 0).unwrap(),
         }
     }
 }
 
 impl Decode for Xz2Decoder {
     fn reinit(&mut self) -> Result<()> {
-        *self = Self::new();
+        *self = Self::new(self.stream.memlimit());
         Ok(())
     }
 

diff --git a/src/macros.rs b/src/macros.rs
@@ -222,7 +222,19 @@ macro_rules! algos {
                 }
             }
         }
-        { @dec }
+        { @dec
+            /// Creates a new decoder, using the specified limit of memory. `Err(Custom { kind:
+            /// Other, error: MemLimit })` can be returned during decoding if the specified limit
+            /// is too small.
+            pub fn with_memlimit(read: $inner, memlimit: u64) -> Self {
+                Self {
+                    inner: crate::$($mod::)+generic::Decoder::new(
+                        read,
+                        crate::codec::XzDecoder::with_memlimit(memlimit),
+                    ),
+                }
+            }
+        }
         );
 
         algos!(@algo lzma ["lzma"] LzmaDecoder LzmaEncoder <$inner>
@@ -237,7 +249,20 @@ macro_rules! algos {
                 }
             }
         }
-        { @dec }
+        { @dec
+            /// Creates a new decoder, using the specified limit of memory. `Err(Custom { kind:
+            /// Other, error: MemLimit })` can be returned during decoding if the specified limit
+            /// is too small.
+            pub fn with_memlimit(read: $inner, memlimit: u64) -> Self {
+                Self {
+                    inner: crate::$($mod::)+generic::Decoder::new(
+                        read,
+                        crate::codec::LzmaDecoder::with_memlimit(memlimit),
+                    ),
+                }
+            }
+
+        }
         );
     }
 }