1use crate::checksum::{make_checksum, parse_checksum, write_checksum, CHECKSUM_LEN};
2use crate::reader::Reader;
3use crate::version::{Version, VERSION_LEN};
4use crate::writer::Writer;
5use std::io::{Read, Seek, SeekFrom, Write};
6
7#[cfg(feature = "async")]
8use crate::reader::AsyncReader;
9#[cfg(feature = "async")]
10use crate::writer::AsyncWriter;
11#[cfg(feature = "async")]
12use tokio::io::{AsyncRead, AsyncReadExt, AsyncSeek, AsyncSeekExt, AsyncWrite, AsyncWriteExt};
13
14#[derive(Debug, Clone, Copy, PartialEq)]
15pub struct Header {
16 version: Version,
17 index_size: u32,
18}
19
20impl Header {
21 pub const MAGIC_LEN: usize = 8;
22 pub const MAGIC: [u8; Header::MAGIC_LEN] = [0xf0, 0x9f, 0x8c, 0x90, 0xf0, 0x9f, 0x8e, 0x81];
23 pub const MAGIC_OFFSET: u64 = 0;
24 pub const VERSION_OFFSET: u64 = Header::MAGIC_LEN as u64;
25 pub const INDEX_SIZE_OFFSET: u64 = Self::VERSION_OFFSET + VERSION_LEN as u64;
26 pub const INDEX_SIZE_BYTES_LEN: usize = 4;
27 pub const CHECKSUM_OFFSET: u64 = Self::INDEX_SIZE_OFFSET + Self::INDEX_SIZE_BYTES_LEN as u64;
28 pub const END_OFFSET: u64 = Self::CHECKSUM_OFFSET + CHECKSUM_LEN as u64;
29
30 pub fn index_end_offset(&self) -> u64 {
31 Self::END_OFFSET + self.index_size as u64 + CHECKSUM_LEN as u64
32 }
33
34 pub fn new(version: Version, index_size: u32) -> Self {
35 Self {
36 version,
37 index_size,
38 }
39 }
40
41 pub fn version(&self) -> Version {
42 self.version
43 }
44
45 pub fn index_size(&self) -> u32 {
46 self.index_size
47 }
48}
49
50fn write_magic() -> Vec<u8> {
51 Header::MAGIC.to_vec()
52}
53
54fn write_version(version: Version) -> Vec<u8> {
55 version.bytes().to_vec()
56}
57
58fn write_index_size(index_size: u32) -> Vec<u8> {
59 index_size.to_be_bytes().to_vec()
60}
61
62#[derive(Debug, Default, Clone, Copy, PartialEq)]
63pub struct HeaderWriterOptions {
64 pub(crate) checksum_seed: u32,
65}
66
67impl HeaderWriterOptions {
68 pub fn new() -> Self {
69 Self::default()
70 }
71
72 pub fn checksum_seed(&mut self, seed: u32) -> &mut Self {
73 self.checksum_seed = seed;
74 self
75 }
76}
77
78pub struct HeaderWriter<W: Write> {
79 w: W,
80 options: HeaderWriterOptions,
81}
82
83impl<W: Write> HeaderWriter<W> {
84 pub fn new(w: W) -> Self {
85 Self {
86 w,
87 options: Default::default(),
88 }
89 }
90
91 pub fn new_with_options(w: W, options: HeaderWriterOptions) -> Self {
92 Self { w, options }
93 }
94
95 pub fn set_options(&mut self, options: HeaderWriterOptions) -> &mut Self {
96 self.options = options;
97 self
98 }
99
100 pub fn write_magic(&mut self) -> crate::Result<Vec<u8>> {
101 let bytes = write_magic();
102 self.w.write_all(&bytes)?;
103 Ok(bytes)
104 }
105
106 pub fn write_version(&mut self, version: Version) -> crate::Result<Vec<u8>> {
107 let bytes = write_version(version);
108 self.w.write_all(&bytes)?;
109 Ok(bytes)
110 }
111
112 pub fn write_index_size(&mut self, index_size: u32) -> crate::Result<Vec<u8>> {
113 let bytes = write_index_size(index_size);
114 self.w.write_all(&bytes)?;
115 Ok(bytes)
116 }
117
118 pub fn write_checksum(&mut self, checksum: u32) -> crate::Result<Vec<u8>> {
119 let bytes = write_checksum(checksum);
120 self.w.write_all(&bytes)?;
121 Ok(bytes)
122 }
123}
124
125impl<W: Write> Writer<Header> for HeaderWriter<W> {
126 fn write(&mut self, header: &Header) -> crate::Result<usize> {
127 let mut bytes = vec![];
128 bytes.extend(self.write_magic()?);
129 bytes.extend(self.write_version(header.version)?);
130 bytes.extend(self.write_index_size(header.index_size)?);
131
132 let checksum = make_checksum(self.options.checksum_seed, &bytes);
133 bytes.extend(self.write_checksum(checksum)?);
134 Ok(bytes.len())
135 }
136}
137
138#[cfg(feature = "async")]
139pub struct AsyncHeaderWriter<W: AsyncWrite + Unpin> {
140 w: W,
141 options: HeaderWriterOptions,
142}
143
144#[cfg(feature = "async")]
145impl<W: AsyncWrite + Unpin> AsyncHeaderWriter<W> {
146 pub fn new(w: W) -> Self {
147 Self {
148 w,
149 options: Default::default(),
150 }
151 }
152
153 pub fn new_with_options(w: W, options: HeaderWriterOptions) -> Self {
154 Self { w, options }
155 }
156
157 pub fn set_options(&mut self, options: HeaderWriterOptions) -> &mut Self {
158 self.options = options;
159 self
160 }
161
162 pub async fn write_magic(&mut self) -> crate::Result<Vec<u8>> {
163 let bytes = write_magic();
164 self.w.write_all(&bytes).await?;
165 Ok(bytes)
166 }
167
168 pub async fn write_version(&mut self, version: Version) -> crate::Result<Vec<u8>> {
169 let bytes = write_version(version);
170 self.w.write_all(&bytes).await?;
171 Ok(bytes)
172 }
173
174 pub async fn write_index_size(&mut self, index_size: u32) -> crate::Result<Vec<u8>> {
175 let bytes = write_index_size(index_size);
176 self.w.write_all(&bytes).await?;
177 Ok(bytes)
178 }
179
180 pub async fn write_checksum(&mut self, checksum: u32) -> crate::Result<Vec<u8>> {
181 let bytes = write_checksum(checksum);
182 self.w.write_all(&bytes).await?;
183 Ok(bytes)
184 }
185}
186
187#[cfg(feature = "async")]
188impl<W: AsyncWrite + Unpin> AsyncWriter<Header> for AsyncHeaderWriter<W> {
189 async fn write(&mut self, header: &Header) -> crate::Result<usize> {
190 let mut bytes = vec![];
191 bytes.extend(self.write_magic().await?);
192 bytes.extend(self.write_version(header.version).await?);
193 bytes.extend(self.write_index_size(header.index_size).await?);
194
195 let checksum = make_checksum(self.options.checksum_seed, &bytes);
196 bytes.extend(self.write_checksum(checksum).await?);
197 Ok(bytes.len())
198 }
199}
200
201fn read_magic() -> (u64, [u8; Header::MAGIC_LEN]) {
202 (Header::MAGIC_OFFSET, [0u8; Header::MAGIC_LEN])
203}
204
205fn parse_magic(buf: &[u8; Header::MAGIC_LEN]) -> crate::Result<()> {
206 if buf != Header::MAGIC.as_ref() {
207 return Err(crate::Error::InvalidMagicNum);
208 }
209 Ok(())
210}
211
212fn read_version() -> (u64, [u8; VERSION_LEN]) {
213 (Header::VERSION_OFFSET, [0u8; VERSION_LEN])
214}
215
216fn parse_version(buf: &[u8; VERSION_LEN]) -> crate::Result<Version> {
217 if buf == Version::V1.bytes().as_ref() {
218 return Ok(Version::V1);
219 }
220 Err(crate::Error::InvalidVersion)
221}
222
223fn read_index_size() -> (u64, [u8; Header::INDEX_SIZE_BYTES_LEN]) {
224 (
225 Header::INDEX_SIZE_OFFSET,
226 [0u8; Header::INDEX_SIZE_BYTES_LEN],
227 )
228}
229
230fn parse_index_size(buf: &[u8; Header::INDEX_SIZE_BYTES_LEN]) -> u32 {
231 u32::from_be_bytes(AsRef::<[u8]>::as_ref(&buf).try_into().unwrap())
232}
233
234fn read_checksum() -> (u64, [u8; CHECKSUM_LEN]) {
235 (Header::CHECKSUM_OFFSET, [0u8; CHECKSUM_LEN])
236}
237
238fn read_total() -> (u64, [u8; Header::CHECKSUM_OFFSET as usize]) {
239 (
240 Header::MAGIC_OFFSET,
241 [0u8; Header::CHECKSUM_OFFSET as usize],
242 )
243}
244
245#[derive(Debug, Default, Clone, Copy, PartialEq)]
246pub struct HeaderReaderOptions {
247 pub checksum_seed: u32,
248 pub verify_checksum: bool,
249}
250
251impl HeaderReaderOptions {
252 pub fn new() -> Self {
253 Self::default()
254 }
255
256 pub fn checksum_seed(mut self, seed: u32) -> Self {
257 self.checksum_seed = seed;
258 self
259 }
260
261 pub fn verify_checksum(mut self, verify: bool) -> Self {
262 self.verify_checksum = verify;
263 self
264 }
265}
266
267pub struct HeaderReader<R: Read + Seek> {
268 r: R,
269 options: HeaderReaderOptions,
270}
271
272impl<R: Read + Seek> HeaderReader<R> {
273 pub fn new(r: R) -> Self {
274 Self::new_with_options(r, Default::default())
275 }
276
277 pub fn new_with_options(r: R, options: HeaderReaderOptions) -> Self {
278 Self { r, options }
279 }
280
281 pub fn set_options(&mut self, options: HeaderReaderOptions) -> &mut Self {
282 self.options = options;
283 self
284 }
285
286 pub fn read_magic(&mut self) -> crate::Result<[u8; Header::MAGIC_LEN]> {
287 let (offset, mut buf) = read_magic();
288 self.r.seek(SeekFrom::Start(offset))?;
289 self.r.read_exact(&mut buf)?;
290 parse_magic(&buf)?;
291 Ok(buf)
292 }
293
294 pub fn read_version(&mut self) -> crate::Result<Version> {
295 let (offset, mut buf) = read_version();
296 self.r.seek(SeekFrom::Start(offset))?;
297 self.r.read_exact(&mut buf)?;
298 parse_version(&buf)
299 }
300
301 pub fn read_index_size(&mut self) -> crate::Result<u32> {
302 let (offset, mut buf) = read_index_size();
303 self.r.seek(SeekFrom::Start(offset))?;
304 self.r.read_exact(&mut buf)?;
305 Ok(parse_index_size(&buf))
306 }
307
308 pub fn read_checksum(&mut self) -> crate::Result<u32> {
309 let (offset, mut buf) = read_checksum();
310 self.r.seek(SeekFrom::Start(offset))?;
311 self.r.read_exact(&mut buf)?;
312 let checksum = parse_checksum(&buf);
313 Ok(checksum)
314 }
315
316 fn verify_checksum(&mut self, checksum: u32) -> crate::Result<()> {
317 let (offset, mut total) = read_total();
318 self.r.seek(SeekFrom::Start(offset))?;
319 self.r.read_exact(&mut total)?;
320
321 let expected_checksum = make_checksum(self.options.checksum_seed, &total);
322 if checksum != expected_checksum {
323 return Err(crate::Error::InvalidHeaderChecksum);
324 }
325 Ok(())
326 }
327}
328
329impl<R: Read + Seek> Reader<Header> for HeaderReader<R> {
330 fn read(&mut self) -> crate::Result<Header> {
331 self.read_magic()?;
332 let version = self.read_version()?;
333 let index_size = self.read_index_size()?;
334 let checksum = self.read_checksum()?;
335 if self.options.verify_checksum {
336 self.verify_checksum(checksum)?;
337 }
338 Ok(Header::new(version, index_size))
339 }
340}
341
342#[cfg(feature = "async")]
343pub struct AsyncHeaderReader<R: AsyncRead + AsyncSeek + Unpin> {
344 r: R,
345 options: HeaderReaderOptions,
346}
347
348#[cfg(feature = "async")]
349impl<R: AsyncRead + AsyncSeek + Unpin> AsyncHeaderReader<R> {
350 pub fn new(r: R) -> Self {
351 Self::new_with_options(r, Default::default())
352 }
353
354 pub fn new_with_options(r: R, options: HeaderReaderOptions) -> Self {
355 Self { r, options }
356 }
357
358 pub fn set_options(&mut self, options: HeaderReaderOptions) -> &mut Self {
359 self.options = options;
360 self
361 }
362
363 pub async fn read_magic(&mut self) -> crate::Result<[u8; Header::MAGIC_LEN]> {
364 let (offset, mut buf) = read_magic();
365 self.r.seek(SeekFrom::Start(offset)).await?;
366 self.r.read_exact(&mut buf).await?;
367 parse_magic(&buf)?;
368 Ok(buf)
369 }
370
371 pub async fn read_version(&mut self) -> crate::Result<Version> {
372 let (offset, mut buf) = read_version();
373 self.r.seek(SeekFrom::Start(offset)).await?;
374 self.r.read_exact(&mut buf).await?;
375 parse_version(&buf)
376 }
377
378 pub async fn read_index_size(&mut self) -> crate::Result<u32> {
379 let (offset, mut buf) = read_index_size();
380 self.r.seek(SeekFrom::Start(offset)).await?;
381 self.r.read_exact(&mut buf).await?;
382 Ok(parse_index_size(&buf))
383 }
384
385 pub async fn read_checksum(&mut self) -> crate::Result<u32> {
386 let (offset, mut buf) = read_checksum();
387 self.r.seek(SeekFrom::Start(offset)).await?;
388 self.r.read_exact(&mut buf).await?;
389 let checksum = parse_checksum(&buf);
390 Ok(checksum)
391 }
392
393 async fn verify_checksum(&mut self, checksum: u32) -> crate::Result<()> {
394 let (offset, mut total) = read_total();
395 self.r.seek(SeekFrom::Start(offset)).await?;
396 self.r.read_exact(&mut total).await?;
397
398 let expected_checksum = make_checksum(self.options.checksum_seed, &total);
399 if checksum != expected_checksum {
400 return Err(crate::Error::InvalidHeaderChecksum);
401 }
402 Ok(())
403 }
404}
405
406#[cfg(feature = "async")]
407impl<R: AsyncRead + AsyncSeek + Unpin> AsyncReader<Header> for AsyncHeaderReader<R> {
408 async fn read(&mut self) -> crate::Result<Header> {
409 self.read_magic().await?;
410 let version = self.read_version().await?;
411 let index_size = self.read_index_size().await?;
412 let checksum = self.read_checksum().await?;
413 if self.options.verify_checksum {
414 self.verify_checksum(checksum).await?;
415 }
416 Ok(Header::new(version, index_size))
417 }
418}
419
420#[cfg(test)]
421mod tests {
422 use super::*;
423 use std::io::Cursor;
424
425 #[test]
426 fn read_and_write() {
427 let header = Header::new(Version::V1, 1234);
428 let mut buf = vec![];
429 let mut writer = HeaderWriter::new(Cursor::new(&mut buf));
430 writer.write(&header).unwrap();
431 assert_eq!(
432 buf,
433 [240, 159, 140, 144, 240, 159, 142, 129, 1, 0, 0, 4, 210, 49, 56, 3, 16]
434 );
435 let mut reader = HeaderReader::new(Cursor::new(&buf));
436 let read_header = reader.read().unwrap();
437 assert_eq!(header, read_header);
438 assert_eq!(read_header.version(), Version::V1);
439 assert_eq!(read_header.index_size(), 1234);
440 }
441
442 #[cfg(feature = "async")]
443 #[tokio::test]
444 async fn async_read_and_write() {
445 let header = Header::new(Version::V1, 1234);
446 let mut buf = vec![];
447 let mut writer = AsyncHeaderWriter::new(Cursor::new(&mut buf));
448 writer.write(&header).await.unwrap();
449 assert_eq!(
450 buf,
451 [240, 159, 140, 144, 240, 159, 142, 129, 1, 0, 0, 4, 210, 49, 56, 3, 16]
452 );
453 let mut reader = AsyncHeaderReader::new(Cursor::new(&buf));
454 let read_header = reader.read().await.unwrap();
455 assert_eq!(header, read_header);
456 assert_eq!(read_header.version(), Version::V1);
457 assert_eq!(read_header.index_size(), 1234);
458 }
459}