zip.rs

//! ZIP を読む
//!
//! docx、xlsx、pptx はどれも中身が ZIP である。ここはその外側を扱う。
//!
//! # 大きなパートを順に読むこと
//!
//! [`Archive::reader`] は展開を別のスレッドに任せ、出来たぶんから順に渡す。
//! 展開しきったものをメモリに置かないので、数 GB のシートでも読める。
//! 読む側と展開する側が同時に進むので、待ち時間も減る。
//!
//! ただし小さなパートは、スレッドを立てずにその場で展開する。
//! 展開がすぐ終わるので、スレッドを立てて受け渡す手間のほうが大きくなるためである。
//!
//! # 小さなファイルは丸ごと読む
//!
//! 決算の表のような小さなファイルを数多く読むときは、パートごとにファイルの位置を合わせて読む
//! 手間が目立つ。4MB 以下のファイルは開いたときに 1 回で全部読み、パートはそこから取り出す。

use crate::crc::{Crc, checksum};
use crate::error::Error;
use crate::inflate;
use std::collections::HashMap;
use std::fs::File;
use std::io::{Read, Seek, SeekFrom};
use std::path::Path;
use std::sync::mpsc::{Receiver, sync_channel};


/// 縮めていない
pub const STORED: u16 = 0;

/// deflate で縮めてある
pub const DEFLATED: u16 = 8;

/// これを超えたら 32 ビットに収まらない
const LIMIT: u64 = 0xffffffff;

/// 末尾から中央目録の在り処を探す範囲。追記コメントの上限 65535 に少し足したもの
const TAIL: u64 = 66 * 1024;

/// 一度に受け渡すバイト数
const PIECE: usize = 256 * 1024;

/// 展開したあとの長さがこれ以下のパートは、スレッドを立てずにその場で展開する
const SMALL: u64 = 256 * 1024;

/// これ以下の大きさのファイルは、開いたときに丸ごと読む
const WHOLE: u64 = 4 * 1024 * 1024;


/// 中央目録に載っている 1 エントリ
#[derive(Debug, Clone)]
pub struct Entry {
    /// パートの名前
    pub name: String,
    /// STORED か DEFLATED
    pub method: u16,
    /// 中身の CRC-32
    pub crc: u32,
    /// 縮めたあとの長さ
    pub compressed: u64,
    /// 縮める前の長さ
    pub uncompressed: u64,
    /// ローカルヘッダの位置
    pub offset: u64,
}


/// 読むために開いた ZIP
pub struct Archive {
    file: File,
    /// ファイルの大きさ
    size: u64,
    /// 小さなファイルの中身。丸ごと読んであれば、パートはここから取り出す
    whole: Option<Vec<u8>>,
    entry: Vec<Entry>,
    place: HashMap<String, usize>,
}


impl Archive {
    /// ファイルを開いて中央目録を読む
    pub fn open(path: &Path) -> Result<Archive, Error> {
        let mut file = File::open(path)?;
        let size = file.seek(SeekFrom::End(0))?;

        let whole = if size <= WHOLE {
            let mut data = Vec::with_capacity(size as usize);
            file.seek(SeekFrom::Start(0))?;
            file.read_to_end(&mut data)?;
            Some(data)
        } else {
            None
        };

        let mut archive = Archive {file, size, whole, entry: Vec::new(), place: HashMap::new()};
        let tail = TAIL.min(size);
        let region = archive.fetch(size - tail, tail as usize)?;

        let end = find_end(&region)?;
        let mut count = u16(&region, end + 10)? as u64;
        let mut length = u32(&region, end + 12)? as u64;
        let mut start = u32(&region, end + 16)? as u64;

        // ZIP64 なら、目録の位置と数は別の記録に入っている
        if count == 0xffff || length == LIMIT || start == LIMIT {
            let locator = find_locator(&region, end)?;
            let position = u64_at(&region, locator + 8)?;
            let record = archive.fetch(position, 56)?;

            if u32(&record, 0)? != 0x06064b50 {
                return Err(Error::new("ZIP64 の中央目録の印が見つかりません"));
            }

            count = u64_at(&record, 32)?;
            length = u64_at(&record, 40)?;
            start = u64_at(&record, 48)?;
        }

        let region = archive.fetch(start, length as usize)?;
        let mut entry = Vec::with_capacity((count as usize).min(region.len() / 46));
        let mut place = HashMap::with_capacity(entry.capacity());
        let mut cursor = 0;

        while cursor + 46 <= region.len() {
            if u32(&region, cursor)? != 0x02014b50 {
                break;
            }

            let one = read_entry(&region, cursor)?;
            let name_length = u16(&region, cursor + 28)? as usize;
            let extra_length = u16(&region, cursor + 30)? as usize;
            let comment_length = u16(&region, cursor + 32)? as usize;
            place.insert(one.name.clone(), entry.len());
            entry.push(one);
            cursor += 46 + name_length + extra_length + comment_length;
        }

        archive.entry = entry;
        archive.place = place;
        Ok(archive)
    }

    /// offset から length バイトを取り出す。丸ごと読んであればそこから、なければファイルから読む
    fn fetch(&mut self, offset: u64, length: usize) -> Result<Vec<u8>, Error> {
        // 壊れた記録が大きな長さを指していても、ファイルの外は読みに行かない
        let Some(end) = offset.checked_add(length as u64).filter(|end| *end <= self.size) else {
            return Err(Error::new("ZIP の記録がファイルの外を指しています"));
        };

        if let Some(data) = &self.whole {
            return Ok(data[offset as usize..end as usize].to_vec());
        }

        let mut data = vec![0u8; length];
        self.file.seek(SeekFrom::Start(offset))?;
        self.file.read_exact(&mut data)?;
        Ok(data)
    }

    /// 入っているエントリ
    pub fn entry(&self) -> &[Entry] {
        &self.entry
    }

    /// その名前のパートがあるか
    pub fn has(&self, name: &str) -> bool {
        self.place.contains_key(name)
    }

    /// 縮めたままのバイト列を取り出す
    pub fn raw(&mut self, name: &str) -> Result<Vec<u8>, Error> {
        let index = self.find(name)?;
        let start = self.data_offset(index)?;
        let length = self.entry[index].compressed as usize;
        self.fetch(start, length).map_err(|error| error.within(name))
    }

    /// 中身を展開して全部返す
    pub fn read(&mut self, name: &str) -> Result<Vec<u8>, Error> {
        let index = self.find(name)?;
        let method = self.entry[index].method;
        let size = self.entry[index].uncompressed as usize;
        let expect = self.entry[index].crc;
        let data = self.raw(name)?;

        let plain = match method {
            STORED => data,
            DEFLATED => inflate::inflate_owned(data, size).map_err(|error| error.within(name))?,
            _ => return Err(Error::at("知らない圧縮の方式です", name)),
        };

        if checksum(&plain) != expect {
            return Err(Error::at("中身が壊れています", name));
        }

        Ok(plain)
    }

    /// 中身を順に受け取る。展開は別のスレッドが進める
    ///
    /// 縮めたバイトもまとめては読まない。ファイルから少しずつ渡すので、
    /// 縮めて数百 MB あるパートでも、手元に置くのは読みかけのぶんだけになる。
    pub fn reader(&mut self, name: &str) -> Result<Reader, Error> {
        let index = self.find(name)?;
        let method = self.entry[index].method;

        if method == DEFLATED && self.entry[index].uncompressed <= SMALL {
            return Ok(Reader::plain(self.read(name)?));
        }

        // 丸ごと読んであれば、縮めていないパートはそこから渡す
        if method == STORED && self.whole.is_some() {
            return Ok(Reader::plain(self.read(name)?));
        }

        let crc = self.entry[index].crc;
        let length = self.entry[index].compressed;
        let start = self.data_offset(index)?;
        let mut file = self.file.try_clone()?;
        file.seek(SeekFrom::Start(start))?;

        match method {
            STORED => Ok(Reader {
                origin: Origin::File(file, length, Crc::new(), crc, name.to_string()),
            }),

            DEFLATED => Ok(Reader {
                origin: Origin::Thread(spawn(file, length, crc, name.to_string())),
            }),

            _ => Err(Error::at("知らない圧縮の方式です", name)),
        }
    }

    /// 名前からエントリの番号を引く
    fn find(&self, name: &str) -> Result<usize, Error> {
        match self.place.get(name) {
            Some(index) => Ok(*index),
            None => Err(Error::at("パートが見つかりません", name)),
        }
    }

    /// ローカルヘッダを読んで、中身が始まる位置を求める
    fn data_offset(&mut self, index: usize) -> Result<u64, Error> {
        let offset = self.entry[index].offset;
        let header = self.fetch(offset, 30)?;

        if u32(&header, 0)? != 0x04034b50 {
            return Err(Error::at("ローカルヘッダの印が見つかりません", &self.entry[index].name));
        }

        let name_length = u16(&header, 26)? as u64;
        let extra_length = u16(&header, 28)? as u64;
        Ok(offset + 30 + name_length + extra_length)
    }
}


/// 中身を順に受け取るもの
pub struct Reader {
    origin: Origin,
}


/// どこから受け取るか
enum Origin {
    /// すでに手元にあるバイト列を、そのまま切って渡す
    Plain(Vec<u8>, usize),
    /// 縮めていないので、ファイルから読んでそのまま渡す
    File(File, u64, Crc, u32, String),
    /// 別のスレッドが展開しながら送ってくる
    Thread(Receiver<Result<Vec<u8>, Error>>),
}


impl Reader {
    /// そろっているバイト列を、順に受け取る形にして返す
    ///
    /// 小さなパートは展開しきってからこちらに載せる。読む側は大きなパートと同じ書き方でよい。
    pub fn plain(data: Vec<u8>) -> Reader {
        Reader {origin: Origin::Plain(data, 0)}
    }

    /// 続きを受け取る。終わったら None
    pub fn next(&mut self) -> Result<Option<Vec<u8>>, Error> {
        match &mut self.origin {
            Origin::Plain(data, sent) => {
                if *sent >= data.len() {
                    return Ok(None);
                }

                // まだ何も渡していなければ、写さずに丸ごと渡す
                if *sent == 0 {
                    *sent = data.len();
                    return Ok(Some(std::mem::take(data)));
                }

                let stop = (*sent + PIECE).min(data.len());
                let piece = data[*sent..stop].to_vec();
                *sent = stop;
                Ok(Some(piece))
            }

            Origin::File(file, rest, crc, expect, name) => {
                if *rest == 0 {
                    if crc.finish() != *expect {
                        return Err(Error::at("中身が壊れています", name.as_str()));
                    }

                    return Ok(None);
                }

                let take = (*rest as usize).min(PIECE);
                let mut piece = vec![0u8; take];
                file.read_exact(&mut piece)?;
                crc.update(&piece);
                *rest -= take as u64;
                Ok(Some(piece))
            }

            Origin::Thread(receiver) => match receiver.recv() {
                Ok(Ok(piece)) => Ok(Some(piece)),
                Ok(Err(error)) => Err(error),
                Err(_) => Ok(None),
            },
        }
    }
}


/// 展開する側のスレッドを立てる
///
/// 展開は読む側と同時に進む。待ち行列は少しだけ持ち、読む側が遅ければそこで止まる。
fn spawn(mut file: File, mut rest: u64, expect: u32, name: String) -> Receiver<Result<Vec<u8>, Error>> {
    let (sender, receiver) = sync_channel(4);

    std::thread::spawn(move || {
        let carry = sender.clone();
        let mut crc = Crc::new();

        let mut supply = move || -> Result<Option<Vec<u8>>, Error> {
            if rest == 0 {
                return Ok(None);
            }

            let take = (rest as usize).min(PIECE);
            let mut piece = vec![0u8; take];
            file.read_exact(&mut piece)?;
            rest -= take as u64;
            Ok(Some(piece))
        };

        let mut result = inflate::inflate_stream(&mut supply, &mut |part| {
            crc.update(part);

            match sender.send(Ok(part.to_vec())) {
                Ok(()) => Ok(()),
                Err(_) => Err(Error::new("読む側がいなくなりました")),
            }
        });

        if result.is_ok() && crc.finish() != expect {
            result = Err(Error::new("中身が壊れています"));
        }

        if let Err(error) = result {
            let _ = carry.send(Err(error.within(&name)));
        }
    });

    receiver
}


/// 中央目録の 1 エントリを読む
fn read_entry(region: &[u8], cursor: usize) -> Result<Entry, Error> {
    let method = u16(region, cursor + 10)?;
    let crc = u32(region, cursor + 16)?;
    let mut compressed = u32(region, cursor + 20)? as u64;
    let mut uncompressed = u32(region, cursor + 24)? as u64;
    let name_length = u16(region, cursor + 28)? as usize;
    let extra_length = u16(region, cursor + 30)? as usize;
    let mut offset = u32(region, cursor + 42)? as u64;

    let start = cursor + 46;

    if start + name_length + extra_length > region.len() {
        return Err(Error::new("中央目録が途中で終わっています"));
    }

    let name = String::from_utf8_lossy(&region[start..start + name_length]).into_owned();
    let extra = &region[start + name_length..start + name_length + extra_length];
    let mut cursor = 0;

    while cursor + 4 <= extra.len() {
        let kind = u16(extra, cursor)?;
        let length = u16(extra, cursor + 2)? as usize;
        let body = cursor + 4;

        if body + length > extra.len() {
            break;
        }

        if kind == 1 {
            // 32 ビットに収まらなかったものだけが、この順に入っている
            let mut place = body;

            if uncompressed == LIMIT && place + 8 <= body + length {
                uncompressed = u64_at(extra, place)?;
                place += 8;
            }

            if compressed == LIMIT && place + 8 <= body + length {
                compressed = u64_at(extra, place)?;
                place += 8;
            }

            if offset == LIMIT && place + 8 <= body + length {
                offset = u64_at(extra, place)?;
            }
        }

        cursor = body + length;
    }

    Ok(Entry {name, method, crc, compressed, uncompressed, offset})
}


/// 末尾から中央目録の終わりの印を探す
fn find_end(region: &[u8]) -> Result<usize, Error> {
    if region.len() < 22 {
        return Err(Error::new("ZIP として短すぎます"));
    }

    let mut cursor = region.len() - 22;

    loop {
        if u32(region, cursor)? == 0x06054b50 {
            let comment = u16(region, cursor + 20)? as usize;

            if cursor + 22 + comment == region.len() {
                return Ok(cursor);
            }
        }

        if cursor == 0 {
            return Err(Error::new("ZIP の中央目録が見つかりません"));
        }

        cursor -= 1;
    }
}


/// ZIP64 の中央目録の在り処を指す印を探す
fn find_locator(region: &[u8], end: usize) -> Result<usize, Error> {
    if end < 20 {
        return Err(Error::new("ZIP64 の目印が見つかりません"));
    }

    let cursor = end - 20;

    if u32(region, cursor)? != 0x07064b50 {
        return Err(Error::new("ZIP64 の目印が見つかりません"));
    }

    Ok(cursor)
}


/// 2 バイトを読む
fn u16(region: &[u8], offset: usize) -> Result<u16, Error> {
    if offset + 2 > region.len() {
        return Err(Error::new("ZIP の記録が途中で終わっています"));
    }

    Ok(u16::from_le_bytes([region[offset], region[offset + 1]]))
}


/// 4 バイトを読む
fn u32(region: &[u8], offset: usize) -> Result<u32, Error> {
    if offset + 4 > region.len() {
        return Err(Error::new("ZIP の記録が途中で終わっています"));
    }

    Ok(u32::from_le_bytes([
        region[offset],
        region[offset + 1],
        region[offset + 2],
        region[offset + 3],
    ]))
}


/// 8 バイトを読む
fn u64_at(region: &[u8], offset: usize) -> Result<u64, Error> {
    if offset + 8 > region.len() {
        return Err(Error::new("ZIP の記録が途中で終わっています"));
    }

    let mut value = [0u8; 8];
    value.copy_from_slice(&region[offset..offset + 8]);
    Ok(u64::from_le_bytes(value))
}