zran: fix an assert when generating ZRAN images

Fix an assert when generating ZRAN images.

Signed-off-by: Jiang Liu <gerry@linux.alibaba.com>

src/bin/nydus-image/builder/tarball.rs

@@ -34,7 +34,7 @@ use nydus_storage::device::BlobFeatures;
 use nydus_storage::meta::ZranContextGenerator;
 use nydus_storage::RAFS_MAX_CHUNKS_PER_BLOB;
 use nydus_utils::compact::makedev;
-use nydus_utils::compress::zlib_random::ZranReader;
+use nydus_utils::compress::zlib_random::{ZranReader, ZRAN_READER_BUF_SIZE};
 use nydus_utils::compress::ZlibDecoder;
 use nydus_utils::digest::RafsDigest;
 use nydus_utils::{div_round_up, BufReaderInfo, ByteSize};
@@ -108,7 +108,7 @@ impl<'a> TarballTreeBuilder<'a> {
             }
             ConversionType::EStargzToRef | ConversionType::TargzToRef => {
                 // Use 64K buffer to keep consistence with zlib-random.
-                let mut buf_reader = BufReader::with_capacity(0x10000, file);
+                let mut buf_reader = BufReader::with_capacity(ZRAN_READER_BUF_SIZE, file);
                 let mut buf = [0u8; 3];
                 if buf_reader.read_exact(&mut buf).is_ok()
                     && buf[0] == 0x1f

utils/src/compress/zlib_random.rs

@@ -22,8 +22,9 @@ use sha2::{Digest, Sha256};
 pub const ZRAN_DICT_WIN_SIZE: usize = 1 << 15;
 /// Maximum number of random access slices per compression object.
 pub const ZRAN_MAX_CI_ENTRIES: usize = 1 << 24;
+/// Buffer size for ZRAN reader.
+pub const ZRAN_READER_BUF_SIZE: usize = 64 * 1024;
 
-const ZRAN_READER_BUF_SIZE: usize = 64 * 1024;
 const ZRAN_MIN_COMP_SIZE: u64 = 768 * 1024;
 const ZRAN_MAX_COMP_SIZE: u64 = 2048 * 1024;
 const ZRAN_MAX_UNCOMP_SIZE: u64 = 2048 * 1024;
@@ -117,7 +118,6 @@ impl ZranDecoder {
         self.stream.set_dict(dict)?;
 
         self.stream.set_next_in(input);
-        self.stream.set_avail_in(ctx.in_len as uInt);
         self.stream.set_next_out(output);
         self.stream.set_avail_out(ctx.out_len as uInt);
         let ret = self.stream.inflate(true);
@@ -384,8 +384,7 @@ impl<R> ZranReaderState<R> {
     fn new(reader: R) -> Result<Self> {
         let mut stream = ZranStream::new(false)?;
         let input = vec![0u8; ZRAN_READER_BUF_SIZE];
-        stream.set_next_in(&input);
-        stream.set_avail_in(0);
+        stream.set_next_in(&input[0..0]);
 
         Ok(ZranReaderState {
             stream,
@@ -422,14 +421,21 @@ impl<R: Read> Read for ZranReaderState<R> {
         loop {
             // Reload the input buffer when needed.
             if self.stream.avail_in() == 0 {
+                error!("last byte ptr {:x}", self.stream.stream.next_in as usize);
+                if self.stream.stream.next_in > self.input.as_mut_ptr() {
+                    self.stream.last_byte = unsafe { *self.stream.stream.next_in.sub(1) };
+                }
                 let sz = self.reader.read(self.input.as_mut_slice())?;
                 if sz == 0 {
                     return Ok(0);
                 }
                 self.reader_hash.update(&self.input[0..sz]);
                 self.reader_size += sz as u64;
-                self.stream.set_next_in(&self.input);
-                self.stream.set_avail_in(sz as u32);
+                self.stream.set_next_in(&self.input[..sz]);
+                error!(
+                    "zran last byte ptr {:x}",
+                    unsafe { self.stream.stream.next_in.add(sz) } as usize
+                );
             }
 
             match self.stream.inflate(false) {
@@ -475,6 +481,7 @@ struct ZranStream {
     stream: Box<z_stream>,
     total_in: u64,
     total_out: u64,
+    last_byte: u8,
 }
 
 impl ZranStream {
@@ -516,6 +523,7 @@ impl ZranStream {
             stream,
             total_in: 0,
             total_out: 0,
+            last_byte: 0,
         })
     }
 
@@ -543,9 +551,13 @@ impl ZranStream {
     }
 
     fn get_compression_info(&mut self, buf: &[u8], stream_switched: u8) -> ZranCompInfo {
-        let previous_byte = if self.stream.data_type & 0x3f != 0 {
-            assert!(self.stream.next_in as usize > buf.as_ptr() as usize);
-            unsafe { *self.stream.next_in.sub(1) }
+        let previous_byte = if self.stream.data_type & 0x7 != 0 {
+            assert!(self.stream.next_in as usize >= buf.as_ptr() as usize);
+            if self.stream.next_in as usize == buf.as_ptr() as usize {
+                self.last_byte
+            } else {
+                unsafe { *self.stream.next_in.sub(1) }
+            }
         } else {
             0
         };
@@ -554,7 +566,7 @@ impl ZranStream {
             out_pos: self.total_out,
             flags: self.stream.data_type as u32,
             previous_byte,
-            pending_bits: self.stream.data_type as u8 & 0x3f,
+            pending_bits: self.stream.data_type as u8 & 0x7,
             stream_switched,
         }
     }
@@ -604,6 +616,7 @@ impl ZranStream {
 
     fn set_next_in(&mut self, buf: &[u8]) {
         self.stream.next_in = buf.as_ptr() as *mut u8;
+        self.set_avail_in(buf.len() as u32);
     }
 
     fn avail_in(&self) -> u32 {