1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122

//! # Interact with PK files in the file system

use std::{
    fs::{File, OpenOptions},
    io::{self, BufReader, BufWriter, Seek, Write},
    path::Path,
};

use crate::{
    common::{CRCTree, FileMetaPair},
    crc::CRC,
};

use super::{
    file::{PKEntryData, PKTrailer, MAGIC_SEP, MAGIC_START},
    reader::PackFile,
    writer::{write_pk_directory_tree, write_pk_trailer},
};

/// Handle to a PK file
///
/// This is a handle to an open PK file. Open means that the dictionary is in memory, but it
/// holds a handle to the underlying file and can add files as needed.
pub struct PKHandle {
    /// The file handle
    file: File,
    /// The last write position & num compressed
    trailer: PKTrailer,
    /// The directory
    directory: CRCTree<PKEntryData>,
}

/// Inversion of control to put bytes into PK
pub trait PKWriter {
    /// Write the bytes into the file
    fn write<W: Write>(&mut self, writer: &mut W) -> io::Result<()>;
}

impl PKHandle {
    /// Open a PK file
    pub fn open(path: &Path) -> io::Result<PKHandle> {
        let mut file = OpenOptions::new()
            .create(true)
            .truncate(false)
            .write(true)
            .read(true)
            .open(path)?;
        let meta = file.metadata()?;
        let new = meta.len() == 0;

        let (directory, trailer) = if new {
            file.write_all(&MAGIC_START)?;

            let file_list_base_addr = MAGIC_START.len() as u32;
            (
                CRCTree::new(),
                PKTrailer {
                    file_list_base_addr,
                    num_compressed: 0,
                },
            )
        } else {
            let buf = BufReader::new(&mut file);
            let mut pk = PackFile::open(buf);

            pk.check_magic()?;

            let trailer = pk.get_header()?;
            let mut acc = pk.get_entry_accessor(trailer.file_list_base_addr)?;

            (acc.read_all()?, trailer)
        };

        Ok(PKHandle {
            file,
            directory,
            trailer,
        })
    }

    /// Put a file into the PK
    pub fn put_file<W: PKWriter>(
        &mut self,
        crc: CRC,
        writer: &mut W,
        meta: FileMetaPair,
        is_compressed: bool,
    ) -> io::Result<()> {
        let mut buf = BufWriter::new(&mut self.file);
        let start = buf.stream_position()?;
        assert!(start <= u32::MAX.into());

        writer.write(&mut buf)?;
        buf.write_all(&MAGIC_SEP)?;
        let end = buf.stream_position()?;
        assert!(end <= u32::MAX.into());

        let is_compressed = u32::from(is_compressed);
        self.directory.insert(
            crc,
            PKEntryData {
                meta,
                file_data_addr: start as u32,
                is_compressed,
            },
        );

        self.trailer.file_list_base_addr = end as u32;
        self.trailer.num_compressed += is_compressed;

        Ok(())
    }

    /// Finish the file by writing the directory
    pub fn finish(&mut self) -> io::Result<()> {
        let mut buf = BufWriter::new(&mut self.file);
        write_pk_directory_tree(&mut buf, &self.directory)?;
        write_pk_trailer(&mut buf, &self.trailer)?;
        // FIXME: truncate file if necessary
        Ok(())
    }
}