1 files changed, 49 insertions, 57 deletions

diff --git a/src/zip/file/read.rs b/src/zip/file/read.rs
index 7d683db..f5a54f3 100644
--- a/src/zip/file/read.rs
+++ b/src/zip/file/read.rs
@@ -1,27 +1,26 @@
 use crate::driver::FileDriver;
-use crate::utils::ReadUtils;
+use crate::utils::{IoCursor, ReadUtils};
 use crate::zip::{CompressionMethod, ZipError, ZipFileInfo, ZipResult};
 use bzip2::read::BzDecoder;
 use flate2::read::DeflateDecoder;
 use liblzma::read::XzDecoder;
 use liblzma::stream::{Filters, LzmaOptions, Stream};
 use std::io::{
-    Error as IoError, ErrorKind as IoErrorKind, Read, Result as IoResult, Seek, SeekFrom,
+    BufReader, Error as IoError, ErrorKind as IoErrorKind, Read, Result as IoResult, Seek, SeekFrom,
 };
+use zstd::stream::Decoder as ZstdDecoder;
 
-enum IoProxy<Io: Read> {
+enum Compression<Io: Read> {
     Store(Io),
     Deflate(DeflateDecoder<Io>),
     BZip2(BzDecoder<Io>),
+    Zstd(ZstdDecoder<'static, BufReader<Io>>),
     Xz(XzDecoder<Io>),
 }
 
 pub struct ZipFileReader<'d, Io: Read> {
-    io: IoProxy<&'d mut Io>,
+    io: Compression<IoCursor<&'d mut Io>>,
     info: &'d ZipFileInfo,
-
-    bounds: (u64, u64),
-    cursor: u64,
 }
 
 impl<'d, Io: Read> FileDriver for ZipFileReader<'d, Io> {
@@ -36,8 +35,8 @@ impl<'d, Io: Read> FileDriver for ZipFileReader<'d, Io> {
 impl<'d, Io: Read + Seek> ZipFileReader<'d, Io> {
     pub fn new(io: &'d mut Io, info: &'d ZipFileInfo) -> ZipResult<Self> {
         io.seek(SeekFrom::Start(info.header_pointer))?;
-        let buf = io.read_arr::<30>()?;
 
+        let buf = io.read_arr::<30>()?;
         if u32::from_le_bytes(buf[..4].try_into().unwrap()) != 0x04034b50 {
             return Err(ZipError::InvalidFileHeaderSignature.into());
         }
@@ -45,18 +44,28 @@ impl<'d, Io: Read + Seek> ZipFileReader<'d, Io> {
             + 30
             + u16::from_le_bytes(buf[26..28].try_into().unwrap()) as u64
             + u16::from_le_bytes(buf[28..30].try_into().unwrap()) as u64;
-        let mut cursor = io.seek(SeekFrom::Start(data_pointer))?;
 
         Ok(Self {
             io: match info.compression_method {
-                CompressionMethod::Store => IoProxy::Store(io),
-                CompressionMethod::Deflate => IoProxy::Deflate(DeflateDecoder::new(io)),
-                CompressionMethod::BZip2 => IoProxy::BZip2(BzDecoder::new(io)),
+                CompressionMethod::Store => Compression::Store(IoCursor::new(
+                    io,
+                    data_pointer,
+                    data_pointer + info.compressed_size,
+                )?),
+                CompressionMethod::Deflate => Compression::Deflate(DeflateDecoder::new(
+                    IoCursor::new(io, data_pointer, data_pointer + info.compressed_size)?,
+                )),
+                CompressionMethod::BZip2 => Compression::BZip2(BzDecoder::new(IoCursor::new(
+                    io,
+                    data_pointer,
+                    data_pointer + info.compressed_size,
+                )?)),
                 CompressionMethod::Lzma => {
+                    io.seek(SeekFrom::Start(data_pointer))?;
                     let buf = io.read_arr::<9>()?;
-                    cursor += 9;
-                    IoProxy::Xz(XzDecoder::new_stream(
-                        io,
+
+                    Compression::Xz(XzDecoder::new_stream(
+                        IoCursor::new(io, data_pointer + 9, data_pointer + info.compressed_size)?,
                         Stream::new_raw_decoder(
                             Filters::new().lzma1(
                                 LzmaOptions::new()
@@ -71,18 +80,30 @@ impl<'d, Io: Read + Seek> ZipFileReader<'d, Io> {
                         .unwrap(),
                     ))
                 }
-                CompressionMethod::Xz => IoProxy::Xz(XzDecoder::new(io)),
+                CompressionMethod::Zstd => Compression::Zstd(
+                    ZstdDecoder::new(IoCursor::new(
+                        io,
+                        data_pointer,
+                        data_pointer + info.compressed_size,
+                    )?)
+                    .unwrap(),
+                ),
+                CompressionMethod::Xz => Compression::Xz(XzDecoder::new(IoCursor::new(
+                    io,
+                    data_pointer,
+                    data_pointer + info.compressed_size,
+                )?)),
+                CompressionMethod::Unsupported => {
+                    return Err(ZipError::UnsupportedCompressionMethod.into())
+                }
             },
             info,
-
-            bounds: (cursor, data_pointer + info.compressed_size),
-            cursor,
         })
     }
 
     pub fn seekable(&self) -> bool {
         match self.io {
-            IoProxy::Store(..) => true,
+            Compression::Store(..) => true,
             _ => false,
         }
     }
@@ -90,49 +111,20 @@ impl<'d, Io: Read + Seek> ZipFileReader<'d, Io> {
 
 impl<'d, Io: Read> Read for ZipFileReader<'d, Io> {
     fn read(&mut self, buf: &mut [u8]) -> IoResult<usize> {
-        let upper = buf.len().min((self.bounds.1 - self.cursor) as usize);
-        let bytes = match &mut self.io {
-            IoProxy::Store(io) => io.read(&mut buf[..upper]),
-            IoProxy::Deflate(io) => io.read(&mut buf[..upper]),
-            IoProxy::BZip2(io) => io.read(&mut buf[..upper]),
-            IoProxy::Xz(io) => io.read(&mut buf[..upper]),
-        }?;
-        self.cursor += upper as u64;
-        Ok(bytes)
+        match &mut self.io {
+            Compression::Store(io) => io.read(buf),
+            Compression::Deflate(io) => io.read(buf),
+            Compression::BZip2(io) => io.read(buf),
+            Compression::Zstd(io) => io.read(buf),
+            Compression::Xz(io) => io.read(buf),
+        }
     }
 }
 
 impl<'d, Io: Read + Seek> Seek for ZipFileReader<'d, Io> {
     fn seek(&mut self, pos: SeekFrom) -> IoResult<u64> {
         match &mut self.io {
-            IoProxy::Store(io) => {
-                self.cursor = match pos {
-                    SeekFrom::Start(offset) => self.bounds.0 + offset,
-                    SeekFrom::End(offset) => {
-                        let cursor = self.bounds.1.saturating_add_signed(offset);
-                        if cursor < self.bounds.0 {
-                            return Err(IoError::new(
-                                IoErrorKind::InvalidInput,
-                                ZipError::NegativeFileOffset,
-                            ));
-                        }
-                        cursor
-                    }
-                    SeekFrom::Current(offset) => {
-                        let cursor = self.cursor.saturating_add_signed(offset);
-                        if cursor < self.bounds.0 {
-                            return Err(IoError::new(
-                                IoErrorKind::InvalidInput,
-                                ZipError::NegativeFileOffset,
-                            ));
-                        }
-                        cursor
-                    }
-                }
-                .min(self.bounds.1);
-
-                Ok(io.seek(SeekFrom::Start(self.cursor))? - self.bounds.0)
-            }
+            Compression::Store(io) => io.seek(pos),
             _ => Err(IoError::new(
                 IoErrorKind::Unsupported,
                 ZipError::CompressedDataIsUnseekable,