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::Write;
40 ///
41 /// let mut file = spooled_tempfile(15);
42 ///
43 /// writeln!(file, "short line")?;
44 /// assert!(!file.is_rolled());
45 ///
46 /// // as a result of this write call, the size of the data will exceed
47 /// // `max_size` (15), so it will be written to a temporary file on disk,
48 /// // and the in-memory buffer will be dropped
49 /// writeln!(file, "marvin gardens")?;
50 /// assert!(file.is_rolled());
51 /// # Ok::<(), std::io::Error>(())
52 /// ```
53 #[inline]
spooled_tempfile(max_size: usize) -> SpooledTempFile54 pub fn spooled_tempfile(max_size: usize) -> SpooledTempFile {
55 SpooledTempFile::new(max_size)
56 }
57
58 impl SpooledTempFile {
59 #[must_use]
new(max_size: usize) -> SpooledTempFile60 pub fn new(max_size: usize) -> SpooledTempFile {
61 SpooledTempFile {
62 max_size,
63 inner: SpooledData::InMemory(Cursor::new(Vec::new())),
64 }
65 }
66
67 /// Returns true if the file has been rolled over to disk.
68 #[must_use]
is_rolled(&self) -> bool69 pub fn is_rolled(&self) -> bool {
70 match self.inner {
71 SpooledData::InMemory(_) => false,
72 SpooledData::OnDisk(_) => true,
73 }
74 }
75
76 /// Rolls over to a file on disk, regardless of current size. Does nothing
77 /// if already rolled over.
roll(&mut self) -> io::Result<()>78 pub fn roll(&mut self) -> io::Result<()> {
79 if !self.is_rolled() {
80 let mut file = tempfile()?;
81 if let SpooledData::InMemory(cursor) = &mut self.inner {
82 file.write_all(cursor.get_ref())?;
83 file.seek(SeekFrom::Start(cursor.position()))?;
84 }
85 self.inner = SpooledData::OnDisk(file);
86 }
87 Ok(())
88 }
89
set_len(&mut self, size: u64) -> Result<(), io::Error>90 pub fn set_len(&mut self, size: u64) -> Result<(), io::Error> {
91 if size > self.max_size as u64 {
92 self.roll()?; // does nothing if already rolled over
93 }
94 match &mut self.inner {
95 SpooledData::InMemory(cursor) => {
96 cursor.get_mut().resize(size as usize, 0);
97 Ok(())
98 }
99 SpooledData::OnDisk(file) => file.set_len(size),
100 }
101 }
102
103 /// Consumes and returns the inner `SpooledData` type.
104 #[must_use]
into_inner(self) -> SpooledData105 pub fn into_inner(self) -> SpooledData {
106 self.inner
107 }
108 }
109
110 impl Read for SpooledTempFile {
read(&mut self, buf: &mut [u8]) -> io::Result<usize>111 fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
112 match &mut self.inner {
113 SpooledData::InMemory(cursor) => cursor.read(buf),
114 SpooledData::OnDisk(file) => file.read(buf),
115 }
116 }
117
read_vectored(&mut self, bufs: &mut [io::IoSliceMut<'_>]) -> io::Result<usize>118 fn read_vectored(&mut self, bufs: &mut [io::IoSliceMut<'_>]) -> io::Result<usize> {
119 match &mut self.inner {
120 SpooledData::InMemory(cursor) => cursor.read_vectored(bufs),
121 SpooledData::OnDisk(file) => file.read_vectored(bufs),
122 }
123 }
124
read_to_end(&mut self, buf: &mut Vec<u8>) -> io::Result<usize>125 fn read_to_end(&mut self, buf: &mut Vec<u8>) -> io::Result<usize> {
126 match &mut self.inner {
127 SpooledData::InMemory(cursor) => cursor.read_to_end(buf),
128 SpooledData::OnDisk(file) => file.read_to_end(buf),
129 }
130 }
131
read_to_string(&mut self, buf: &mut String) -> io::Result<usize>132 fn read_to_string(&mut self, buf: &mut String) -> io::Result<usize> {
133 match &mut self.inner {
134 SpooledData::InMemory(cursor) => cursor.read_to_string(buf),
135 SpooledData::OnDisk(file) => file.read_to_string(buf),
136 }
137 }
138
read_exact(&mut self, buf: &mut [u8]) -> io::Result<()>139 fn read_exact(&mut self, buf: &mut [u8]) -> io::Result<()> {
140 match &mut self.inner {
141 SpooledData::InMemory(cursor) => cursor.read_exact(buf),
142 SpooledData::OnDisk(file) => file.read_exact(buf),
143 }
144 }
145 }
146
147 impl Write for SpooledTempFile {
write(&mut self, buf: &[u8]) -> io::Result<usize>148 fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
149 // roll over to file if necessary
150 if matches! {
151 &self.inner, SpooledData::InMemory(cursor)
152 if cursor.position().saturating_add(buf.len() as u64) > self.max_size as u64
153 } {
154 self.roll()?;
155 }
156
157 // write the bytes
158 match &mut self.inner {
159 SpooledData::InMemory(cursor) => cursor.write(buf),
160 SpooledData::OnDisk(file) => file.write(buf),
161 }
162 }
163
write_vectored(&mut self, bufs: &[io::IoSlice<'_>]) -> io::Result<usize>164 fn write_vectored(&mut self, bufs: &[io::IoSlice<'_>]) -> io::Result<usize> {
165 if matches! {
166 &self.inner, SpooledData::InMemory(cursor)
167 // Borrowed from the rust standard library.
168 if bufs
169 .iter()
170 .fold(cursor.position(), |a, b| a.saturating_add(b.len() as u64))
171 > self.max_size as u64
172 } {
173 self.roll()?;
174 }
175 match &mut self.inner {
176 SpooledData::InMemory(cursor) => cursor.write_vectored(bufs),
177 SpooledData::OnDisk(file) => file.write_vectored(bufs),
178 }
179 }
180
181 #[inline]
flush(&mut self) -> io::Result<()>182 fn flush(&mut self) -> io::Result<()> {
183 match &mut self.inner {
184 SpooledData::InMemory(cursor) => cursor.flush(),
185 SpooledData::OnDisk(file) => file.flush(),
186 }
187 }
188 }
189
190 impl Seek for SpooledTempFile {
seek(&mut self, pos: SeekFrom) -> io::Result<u64>191 fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
192 match &mut self.inner {
193 SpooledData::InMemory(cursor) => cursor.seek(pos),
194 SpooledData::OnDisk(file) => file.seek(pos),
195 }
196 }
197 }
198