assetpack_core/file_transform/
spool.rs

1use std::{
2  fs,
3  io::{self, Cursor, Read, Seek, SeekFrom, Write},
4  path::Path,
5};
6
7use crate::error::{Error, Result};
8
9#[derive(Debug)]
10pub enum StoredStreamSpool {
11  Memory { buffer: Vec<u8>, max_bytes: u64 },
12  File { file: tempfile::NamedTempFile, len: u64 },
13}
14
15impl StoredStreamSpool {
16  pub fn from_bytes(buffer: Vec<u8>) -> Self {
17    let max_bytes = buffer.len().max(1) as u64;
18    Self::Memory { buffer, max_bytes }
19  }
20
21  pub fn new_memory(max_bytes: u64) -> Self {
22    Self::Memory {
23      buffer: Vec::new(),
24      max_bytes: max_bytes.max(1),
25    }
26  }
27
28  pub fn new_file(temp_dir: Option<&Path>) -> Result<Self> {
29    let file = if let Some(dir) = temp_dir {
30      fs::create_dir_all(dir)?;
31      tempfile::Builder::new()
32        .prefix("assetpack-transform-precomp2-")
33        .tempfile_in(dir)
34        .map_err(Error::Io)?
35    } else {
36      tempfile::Builder::new()
37        .prefix("assetpack-transform-precomp2-")
38        .tempfile()
39        .map_err(Error::Io)?
40    };
41    Ok(Self::File { file, len: 0 })
42  }
43
44  pub fn len(&self) -> u64 {
45    match self {
46      StoredStreamSpool::Memory { buffer, .. } => buffer.len() as u64,
47      StoredStreamSpool::File { len, .. } => *len,
48    }
49  }
50
51  pub fn is_empty(&self) -> bool {
52    self.len() == 0
53  }
54
55  pub fn as_slice(&self) -> Option<&[u8]> {
56    match self {
57      StoredStreamSpool::Memory { buffer, .. } => Some(buffer.as_slice()),
58      StoredStreamSpool::File { .. } => None,
59    }
60  }
61
62  pub fn reader(&self) -> Result<StoredStreamReader<'_>> {
63    match self {
64      StoredStreamSpool::Memory { buffer, .. } => Ok(StoredStreamReader::Memory(Cursor::new(buffer.as_slice()))),
65      StoredStreamSpool::File { file, .. } => Ok(StoredStreamReader::File(file.reopen()?)),
66    }
67  }
68
69  pub fn into_bytes(self) -> Result<Vec<u8>> {
70    match self {
71      StoredStreamSpool::Memory { buffer, .. } => Ok(buffer),
72      StoredStreamSpool::File { mut file, .. } => {
73        let mut buf = Vec::new();
74        file.rewind()?;
75        file.read_to_end(&mut buf)?;
76        Ok(buf)
77      }
78    }
79  }
80}
81
82impl Write for StoredStreamSpool {
83  fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
84    match self {
85      StoredStreamSpool::Memory { buffer, max_bytes } => {
86        let next_len = buffer.len().saturating_add(buf.len());
87        if next_len as u64 > *max_bytes {
88          return Err(io::Error::new(io::ErrorKind::OutOfMemory, "stored stream exceeded memory limit"));
89        }
90        buffer.extend_from_slice(buf);
91        Ok(buf.len())
92      }
93      StoredStreamSpool::File { file, len } => {
94        file.as_file_mut().write_all(buf)?;
95        *len += buf.len() as u64;
96        Ok(buf.len())
97      }
98    }
99  }
100
101  fn flush(&mut self) -> io::Result<()> {
102    match self {
103      StoredStreamSpool::Memory { .. } => Ok(()),
104      StoredStreamSpool::File { file, .. } => file.as_file_mut().flush(),
105    }
106  }
107}
108
109pub enum StoredStreamReader<'a> {
110  Memory(Cursor<&'a [u8]>),
111  File(std::fs::File),
112}
113
114impl Read for StoredStreamReader<'_> {
115  fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
116    match self {
117      StoredStreamReader::Memory(cursor) => cursor.read(buf),
118      StoredStreamReader::File(file) => file.read(buf),
119    }
120  }
121}
122
123impl Seek for StoredStreamReader<'_> {
124  fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
125    match self {
126      StoredStreamReader::Memory(cursor) => cursor.seek(pos),
127      StoredStreamReader::File(file) => file.seek(pos),
128    }
129  }
130}
assetpack_core/file_transform/spool.rs

assetpack_core/file_transform/
spool.rs
