xp3/
writer.rs

1/*
2 * Created on Mon Dec 14 2020
3 *
4 * Copyright (c) storycraft. Licensed under the Apache Licence 2.0.
5 */
6
7use std::{collections::HashMap, io::{self, Read, Seek, SeekFrom, Write}};
8
9use adler32::RollingAdler32;
10use byteorder::{WriteBytesExt, LittleEndian};
11use flate2::{Compression, read::ZlibEncoder};
12
13use super::{VirtualXP3, XP3Error, XP3_MAGIC, header::{XP3Header, XP3HeaderVersion}, index::{XP3Index, file::{IndexInfoFlag, IndexSegmentFlag, XP3FileIndex, XP3FileIndexAdler, XP3FileIndexInfo, XP3FileIndexSegment, XP3FileIndexTime}}, index_set::{XP3IndexCompression, XP3IndexSet}};
14
15pub struct XP3Writer<T: Write + Seek> {
16
17    stream: T,
18
19    // Position of start
20    start_pos: u64,
21
22    // Position of index offset writing
23    index_pos: u64,
24
25    header: XP3Header,
26    index_set: XP3IndexSet
27
28}
29
30impl<T: Write + Seek> XP3Writer<T> {
31
32    /// Start new XP3 file writing
33    pub fn start(
34        mut stream: T,
35        version: XP3HeaderVersion,
36        index_compression: XP3IndexCompression
37    ) -> Result<Self, XP3Error> {
38        let header = XP3Header::new(version);
39        let start_pos = stream.seek(SeekFrom::Current(0))?;
40
41        // Write magic
42        stream.write_all(&XP3_MAGIC)?;
43        // Fixed version
44        stream.write_u8(1)?;
45        // Write header
46        header.write_bytes(&mut stream)?;
47        // Save position
48        let index_pos = stream.seek(SeekFrom::Current(0))?;
49        // index offset will be overridden after data stream written.
50        stream.write_u64::<LittleEndian>(0)?;
51
52        Ok(Self {
53            stream,
54            start_pos,
55            index_pos,
56            header,
57            index_set: XP3IndexSet::new(index_compression, Vec::new(), HashMap::new())
58        })
59    }
60
61    pub fn append_extra_index(&mut self, index: XP3Index) {
62        self.index_set.extra_mut().push(index);
63    }
64
65    pub fn current_pos(&mut self) -> u64 {
66        self.stream.seek(SeekFrom::Current(0)).unwrap()
67    }
68
69    /// Enter file entry start
70    pub fn enter_file(
71        &mut self,
72        protection: IndexInfoFlag,
73        name: String,
74        timestamp: Option<u64>
75    ) -> EntryWriter<'_, T> {
76        EntryWriter::new(self, protection, name, timestamp)
77    }
78
79    /// Append indexes and finish file
80    pub fn finish(mut self) -> Result<(VirtualXP3, T), XP3Error> {
81        let current = self.stream.seek(SeekFrom::Current(0))?;
82
83        self.stream.seek(SeekFrom::Start(self.index_pos))?;
84        // Write index offset
85        self.stream.write_u64::<LittleEndian>(current - self.start_pos)?;
86
87        self.stream.seek(SeekFrom::Start(current))?;
88
89        // Write indexes
90        self.index_set.write_bytes(&mut self.stream)?;
91
92        Ok((VirtualXP3::new(self.header, self.index_set), self.stream))
93    }
94
95}
96
97pub struct EntryWriter<'a, T: Write + Seek> {
98
99    writer: &'a mut XP3Writer<T>,
100
101    protection: IndexInfoFlag,
102    name: String,
103    timestamp: Option<u64>,
104
105    saved_size: u64,
106    original_size: u64,
107    adler: RollingAdler32,
108
109    default_flag: IndexSegmentFlag,
110
111    written_segments: Vec<XP3FileIndexSegment>
112
113}
114
115impl<'a, T: Write + Seek> EntryWriter<'a, T> {
116
117    pub fn new(
118        writer: &'a mut XP3Writer<T>,
119        protection: IndexInfoFlag,
120        name: String,
121        timestamp: Option<u64>
122    ) -> Self {
123        Self {
124            writer,
125            protection,
126            name,
127            timestamp,
128            saved_size: 0,
129            original_size: 0,
130            adler: RollingAdler32::new(),
131            default_flag: IndexSegmentFlag::UnCompressed,
132            written_segments: Vec::new()
133        }
134    }
135
136    pub fn default_flag(&self) -> IndexSegmentFlag {
137        self.default_flag
138    }
139
140    pub fn set_default_flag(&mut self, flag: IndexSegmentFlag) {
141        self.default_flag = flag;
142    }
143
144    /// Write one segment from stream.
145    /// Returns write size.
146    pub fn write_segment(&mut self, flag: IndexSegmentFlag, data: &[u8]) -> Result<u64, io::Error> {
147        self.adler.update_buffer(&data);
148
149        let original_size: u64 = data.len() as u64;
150        let saved_size: u64;
151        match flag {
152            IndexSegmentFlag::UnCompressed => {
153                self.writer.stream.write_all(data)?;
154
155                saved_size = data.len() as u64;
156                self.saved_size += saved_size;
157            },
158
159            IndexSegmentFlag::Compressed => {
160                let mut compressed = Vec::<u8>::new();
161                let mut encoder = ZlibEncoder::new(&data[..], Compression::fast());
162
163                encoder.read_to_end(&mut compressed)?;
164
165                self.writer.stream.write_all(&mut compressed)?;
166                saved_size = compressed.len() as u64;
167                self.saved_size += saved_size;
168            }
169        }
170        self.original_size += original_size;
171
172        self.written_segments.push(XP3FileIndexSegment::new(
173            flag,
174            self.writer.current_pos() - saved_size - self.writer.start_pos,
175            original_size,
176            saved_size
177        ));
178
179        Ok(saved_size)
180    }
181
182    /// Write with default flag
183    pub fn write_default_segment(&mut self, data: &[u8]) -> Result<u64, io::Error> {
184        self.write_segment(self.default_flag, data)
185    }
186
187    /// Finish file entry
188    pub fn finish(self) {
189        let time = if self.timestamp.is_some() {
190            Some(XP3FileIndexTime::new(self.timestamp.unwrap()))
191        } else {
192            None
193        };
194
195        let index = XP3FileIndex::new(
196            XP3FileIndexInfo::new(
197                self.protection,
198                self.original_size,
199                self.saved_size,
200                self.name.clone()
201            ),
202            self.written_segments.clone(),
203            XP3FileIndexAdler::new(self.adler.hash()),
204            time
205        );
206
207        self.writer.index_set.file_map_mut().insert(index.info().name().clone(), index);
208    }
209
210}
211
212
213impl<'a, T: Write + Seek> Write for EntryWriter<'a, T> {
214
215    /// Write with default flag
216    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
217        Ok(self.write_segment(self.default_flag, buf)? as usize)
218    }
219
220    fn flush(&mut self) -> io::Result<()> {
221        Ok(())
222    }
223
224}
225
226/// Preventing creating more entry before previous entry finish.
227impl<'a, T: Write + Seek> Drop for EntryWriter<'a, T> {
228
229    fn drop(&mut self) { }
230
231}
xp3/writer.rs

xp3/
writer.rs
