1 | use crate::file::tempfile; |
2 | use std::fs::File; |
3 | use std::io::{self, Cursor, Read, Seek, SeekFrom, Write}; |
4 | |
5 | /// A wrapper for the two states of a `SpooledTempFile`. |
6 | #[derive(Debug)] |
7 | pub enum SpooledData { |
8 | InMemory(Cursor<Vec<u8>>), |
9 | OnDisk(File), |
10 | } |
11 | |
12 | /// An object that behaves like a regular temporary file, but keeps data in |
13 | /// memory until it reaches a configured size, at which point the data is |
14 | /// written to a temporary file on disk, and further operations use the file |
15 | /// on disk. |
16 | #[derive(Debug)] |
17 | pub struct SpooledTempFile { |
18 | max_size: usize, |
19 | inner: SpooledData, |
20 | } |
21 | |
22 | /// Create a new spooled temporary file. |
23 | /// |
24 | /// # Security |
25 | /// |
26 | /// This variant is secure/reliable in the presence of a pathological temporary |
27 | /// file cleaner. |
28 | /// |
29 | /// # Resource Leaking |
30 | /// |
31 | /// The temporary file will be automatically removed by the OS when the last |
32 | /// handle to it is closed. This doesn't rely on Rust destructors being run, so |
33 | /// will (almost) never fail to clean up the temporary file. |
34 | /// |
35 | /// # Examples |
36 | /// |
37 | /// ``` |
38 | /// use tempfile::spooled_tempfile; |
39 | /// use std::io::{self, Write}; |
40 | /// |
41 | /// # fn main() { |
42 | /// # if let Err(_) = run() { |
43 | /// # ::std::process::exit(1); |
44 | /// # } |
45 | /// # } |
46 | /// # fn run() -> Result<(), io::Error> { |
47 | /// let mut file = spooled_tempfile(15); |
48 | /// |
49 | /// writeln!(file, "short line" )?; |
50 | /// assert!(!file.is_rolled()); |
51 | /// |
52 | /// // as a result of this write call, the size of the data will exceed |
53 | /// // `max_size` (15), so it will be written to a temporary file on disk, |
54 | /// // and the in-memory buffer will be dropped |
55 | /// writeln!(file, "marvin gardens" )?; |
56 | /// assert!(file.is_rolled()); |
57 | /// |
58 | /// # Ok(()) |
59 | /// # } |
60 | /// ``` |
61 | #[inline ] |
62 | pub fn spooled_tempfile(max_size: usize) -> SpooledTempFile { |
63 | SpooledTempFile::new(max_size) |
64 | } |
65 | |
66 | impl SpooledTempFile { |
67 | #[must_use ] |
68 | pub fn new(max_size: usize) -> SpooledTempFile { |
69 | SpooledTempFile { |
70 | max_size, |
71 | inner: SpooledData::InMemory(Cursor::new(Vec::new())), |
72 | } |
73 | } |
74 | |
75 | /// Returns true if the file has been rolled over to disk. |
76 | #[must_use ] |
77 | pub fn is_rolled(&self) -> bool { |
78 | match self.inner { |
79 | SpooledData::InMemory(_) => false, |
80 | SpooledData::OnDisk(_) => true, |
81 | } |
82 | } |
83 | |
84 | /// Rolls over to a file on disk, regardless of current size. Does nothing |
85 | /// if already rolled over. |
86 | pub fn roll(&mut self) -> io::Result<()> { |
87 | if !self.is_rolled() { |
88 | let mut file = tempfile()?; |
89 | if let SpooledData::InMemory(cursor) = &mut self.inner { |
90 | file.write_all(cursor.get_ref())?; |
91 | file.seek(SeekFrom::Start(cursor.position()))?; |
92 | } |
93 | self.inner = SpooledData::OnDisk(file); |
94 | } |
95 | Ok(()) |
96 | } |
97 | |
98 | pub fn set_len(&mut self, size: u64) -> Result<(), io::Error> { |
99 | if size as usize > self.max_size { |
100 | self.roll()?; // does nothing if already rolled over |
101 | } |
102 | match &mut self.inner { |
103 | SpooledData::InMemory(cursor) => { |
104 | cursor.get_mut().resize(size as usize, 0); |
105 | Ok(()) |
106 | } |
107 | SpooledData::OnDisk(file) => file.set_len(size), |
108 | } |
109 | } |
110 | |
111 | /// Consumes and returns the inner `SpooledData` type. |
112 | #[must_use ] |
113 | pub fn into_inner(self) -> SpooledData { |
114 | self.inner |
115 | } |
116 | } |
117 | |
118 | impl Read for SpooledTempFile { |
119 | fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> { |
120 | match &mut self.inner { |
121 | SpooledData::InMemory(cursor) => cursor.read(buf), |
122 | SpooledData::OnDisk(file) => file.read(buf), |
123 | } |
124 | } |
125 | |
126 | fn read_vectored(&mut self, bufs: &mut [io::IoSliceMut<'_>]) -> io::Result<usize> { |
127 | match &mut self.inner { |
128 | SpooledData::InMemory(cursor) => cursor.read_vectored(bufs), |
129 | SpooledData::OnDisk(file) => file.read_vectored(bufs), |
130 | } |
131 | } |
132 | |
133 | fn read_to_end(&mut self, buf: &mut Vec<u8>) -> io::Result<usize> { |
134 | match &mut self.inner { |
135 | SpooledData::InMemory(cursor) => cursor.read_to_end(buf), |
136 | SpooledData::OnDisk(file) => file.read_to_end(buf), |
137 | } |
138 | } |
139 | |
140 | fn read_to_string(&mut self, buf: &mut String) -> io::Result<usize> { |
141 | match &mut self.inner { |
142 | SpooledData::InMemory(cursor) => cursor.read_to_string(buf), |
143 | SpooledData::OnDisk(file) => file.read_to_string(buf), |
144 | } |
145 | } |
146 | |
147 | fn read_exact(&mut self, buf: &mut [u8]) -> io::Result<()> { |
148 | match &mut self.inner { |
149 | SpooledData::InMemory(cursor) => cursor.read_exact(buf), |
150 | SpooledData::OnDisk(file) => file.read_exact(buf), |
151 | } |
152 | } |
153 | } |
154 | |
155 | impl Write for SpooledTempFile { |
156 | fn write(&mut self, buf: &[u8]) -> io::Result<usize> { |
157 | // roll over to file if necessary |
158 | if matches! { |
159 | &self.inner, SpooledData::InMemory(cursor) |
160 | if cursor.position() as usize + buf.len() > self.max_size |
161 | } { |
162 | self.roll()?; |
163 | } |
164 | |
165 | // write the bytes |
166 | match &mut self.inner { |
167 | SpooledData::InMemory(cursor) => cursor.write(buf), |
168 | SpooledData::OnDisk(file) => file.write(buf), |
169 | } |
170 | } |
171 | |
172 | fn write_vectored(&mut self, bufs: &[io::IoSlice<'_>]) -> io::Result<usize> { |
173 | if matches! { |
174 | &self.inner, SpooledData::InMemory(cursor) |
175 | // Borrowed from the rust standard library. |
176 | if cursor.position() as usize + bufs.iter() |
177 | .fold(0usize, |a, b| a.saturating_add(b.len())) > self.max_size |
178 | } { |
179 | self.roll()?; |
180 | } |
181 | match &mut self.inner { |
182 | SpooledData::InMemory(cursor) => cursor.write_vectored(bufs), |
183 | SpooledData::OnDisk(file) => file.write_vectored(bufs), |
184 | } |
185 | } |
186 | |
187 | #[inline ] |
188 | fn flush(&mut self) -> io::Result<()> { |
189 | match &mut self.inner { |
190 | SpooledData::InMemory(cursor) => cursor.flush(), |
191 | SpooledData::OnDisk(file) => file.flush(), |
192 | } |
193 | } |
194 | } |
195 | |
196 | impl Seek for SpooledTempFile { |
197 | fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> { |
198 | match &mut self.inner { |
199 | SpooledData::InMemory(cursor) => cursor.seek(pos), |
200 | SpooledData::OnDisk(file) => file.seek(pos), |
201 | } |
202 | } |
203 | } |
204 | |