1use crate::file::tempfile;
2use std::fs::File;
3use std::io::{self, Cursor, Read, Seek, SeekFrom, Write};
4
5/// A wrapper for the two states of a `SpooledTempFile`.
6#[derive(Debug)]
7pub 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)]
17pub 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]
62pub fn spooled_tempfile(max_size: usize) -> SpooledTempFile {
63 SpooledTempFile::new(max_size)
64}
65
66impl 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
118impl 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
155impl 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
196impl 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