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]
spooled_tempfile(max_size: usize) -> SpooledTempFile62 pub fn spooled_tempfile(max_size: usize) -> SpooledTempFile {
63 SpooledTempFile::new(max_size)
64 }
65
66 impl SpooledTempFile {
new(max_size: usize) -> SpooledTempFile67 pub fn new(max_size: usize) -> SpooledTempFile {
68 SpooledTempFile {
69 max_size: max_size,
70 inner: SpooledData::InMemory(Cursor::new(Vec::new())),
71 }
72 }
73
74 /// Returns true if the file has been rolled over to disk.
is_rolled(&self) -> bool75 pub fn is_rolled(&self) -> bool {
76 match self.inner {
77 SpooledData::InMemory(_) => false,
78 SpooledData::OnDisk(_) => true,
79 }
80 }
81
82 /// Rolls over to a file on disk, regardless of current size. Does nothing
83 /// if already rolled over.
roll(&mut self) -> io::Result<()>84 pub fn roll(&mut self) -> io::Result<()> {
85 if !self.is_rolled() {
86 let mut file = tempfile()?;
87 if let SpooledData::InMemory(ref mut cursor) = self.inner {
88 file.write_all(cursor.get_ref())?;
89 file.seek(SeekFrom::Start(cursor.position()))?;
90 }
91 self.inner = SpooledData::OnDisk(file);
92 }
93 Ok(())
94 }
95
set_len(&mut self, size: u64) -> Result<(), io::Error>96 pub fn set_len(&mut self, size: u64) -> Result<(), io::Error> {
97 if size as usize > self.max_size {
98 self.roll()?; // does nothing if already rolled over
99 }
100 match self.inner {
101 SpooledData::InMemory(ref mut cursor) => {
102 cursor.get_mut().resize(size as usize, 0);
103 Ok(())
104 }
105 SpooledData::OnDisk(ref mut file) => file.set_len(size),
106 }
107 }
108
109 /// Consumes and returns the inner `SpooledData` type.
into_inner(self) -> SpooledData110 pub fn into_inner(self) -> SpooledData {
111 self.inner
112 }
113 }
114
115 impl Read for SpooledTempFile {
read(&mut self, buf: &mut [u8]) -> io::Result<usize>116 fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
117 match self.inner {
118 SpooledData::InMemory(ref mut cursor) => cursor.read(buf),
119 SpooledData::OnDisk(ref mut file) => file.read(buf),
120 }
121 }
122 }
123
124 impl Write for SpooledTempFile {
write(&mut self, buf: &[u8]) -> io::Result<usize>125 fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
126 // roll over to file if necessary
127 let mut rolling = false;
128 if let SpooledData::InMemory(ref mut cursor) = self.inner {
129 rolling = cursor.position() as usize + buf.len() > self.max_size;
130 }
131 if rolling {
132 self.roll()?;
133 }
134
135 // write the bytes
136 match self.inner {
137 SpooledData::InMemory(ref mut cursor) => cursor.write(buf),
138 SpooledData::OnDisk(ref mut file) => file.write(buf),
139 }
140 }
141
142 #[inline]
flush(&mut self) -> io::Result<()>143 fn flush(&mut self) -> io::Result<()> {
144 match self.inner {
145 SpooledData::InMemory(ref mut cursor) => cursor.flush(),
146 SpooledData::OnDisk(ref mut file) => file.flush(),
147 }
148 }
149 }
150
151 impl Seek for SpooledTempFile {
seek(&mut self, pos: SeekFrom) -> io::Result<u64>152 fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
153 match self.inner {
154 SpooledData::InMemory(ref mut cursor) => cursor.seek(pos),
155 SpooledData::OnDisk(ref mut file) => file.seek(pos),
156 }
157 }
158 }
159