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

//! # Write PK files

use std::io::{self, Seek, Write};

use crate::common::{writer::write_crc_tree, CRCTree, FileMeta};

use super::file::{PKEntryData, PKTrailer};

/// Write the directory of a PK file.
///
/// This function takes a [Write] implementation and a `CRCTree<PKEntryData>`
/// and writes the tree part of the PK directory to disk
pub fn write_pk_directory_tree<W: Write>(
    writer: &mut W,
    tree: &CRCTree<PKEntryData>,
) -> io::Result<()> {
    write_crc_tree(writer, tree, write_pk_entry_data)
}

/// Write the trailer of a PK file
pub fn write_pk_trailer<W: Write>(writer: &mut W, trailer: &PKTrailer) -> io::Result<()> {
    writer.write_all(&trailer.file_list_base_addr.to_le_bytes())?;
    writer.write_all(&trailer.num_compressed.to_le_bytes())?;
    Ok(())
}

/// Write the full directory to disk
///
/// For a "complete" PK file, this function takes the dictionary as a sorted tree
/// and writes the PK directory as well as the trailer to disk.
pub fn write_pk_directory<W: Write + Seek>(
    writer: &mut W,
    tree: &CRCTree<PKEntryData>,
) -> io::Result<()> {
    let file_list_base_addr = writer.stream_position()? as u32;
    let num_compressed = tree
        .iter()
        .filter(|(_, &x)| x.is_compressed & 0xFF > 0)
        .count() as u32;
    let trailer = PKTrailer {
        file_list_base_addr,
        num_compressed,
    };
    write_pk_directory_tree(writer, tree)?;
    write_pk_trailer(writer, &trailer)?;
    Ok(())
}

fn write_file_meta<W: Write>(writer: &mut W, meta: &FileMeta) -> io::Result<()> {
    writer.write_all(&meta.size.to_le_bytes())?;
    write!(writer, "{}\0\0\0\0", &meta.hash)?;
    Ok(())
}

/// Write out a [PKEntryData]
fn write_pk_entry_data<W: Write>(writer: &mut W, entry: &PKEntryData) -> io::Result<()> {
    write_file_meta(writer, &entry.meta.raw)?;
    write_file_meta(writer, &entry.meta.compressed)?;
    writer.write_all(&entry.file_data_addr.to_le_bytes())?;
    writer.write_all(&entry.is_compressed.to_le_bytes())?;
    Ok(())
}

#[cfg(test)]
mod tests {
    use crate::{
        common::{FileMeta, FileMetaPair},
        md5::MD5Sum,
        pk::{file::PKEntryData, parser::parse_pk_entry_data, writer::write_pk_entry_data},
    };

    #[test]
    fn test_write_pk_entry() {
        let mut out: Vec<u8> = vec![];
        let pke = PKEntryData {
            meta: FileMetaPair {
                raw: FileMeta {
                    size: 100,
                    hash: MD5Sum([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]),
                },
                compressed: FileMeta {
                    size: 101,
                    hash: MD5Sum([
                        32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47,
                    ]),
                },
            },
            file_data_addr: 50,
            is_compressed: 256,
        };
        write_pk_entry_data(&mut out, &pke).unwrap();
        assert_eq!(
            out.as_slice(),
            &[
                100, 0, 0, 0, //
                b'0', b'0', b'0', b'1', b'0', b'2', b'0', b'3', b'0', b'4', b'0', b'5', b'0', b'6',
                b'0', b'7', b'0', b'8', b'0', b'9', b'0', b'a', b'0', b'b', b'0', b'c', b'0', b'd',
                b'0', b'e', b'0', b'f', 0, 0, 0, 0, //
                101, 0, 0, 0, //
                b'2', b'0', b'2', b'1', b'2', b'2', b'2', b'3', b'2', b'4', b'2', b'5', b'2', b'6',
                b'2', b'7', b'2', b'8', b'2', b'9', b'2', b'a', b'2', b'b', b'2', b'c', b'2', b'd',
                b'2', b'e', b'2', b'f', 0, 0, 0, 0, //
                50, 0, 0, 0, //
                0, 1, 0, 0
            ]
        );
        let (r, data) = parse_pk_entry_data(&out).unwrap();
        assert_eq!(r, &[] as &[u8]);
        assert_eq!(data, pke);
    }
}