1 use crate::file::tempfile;
2 use std::fs::File;
3 use std::io::{self, Cursor, Read, Seek, SeekFrom, Write};
4
5 #[derive(Debug)]
6 enum SpooledInner {
7 InMemory(Cursor<Vec<u8>>),
8 OnDisk(File),
9 }
10
11 /// An object that behaves like a regular temporary file, but keeps data in
12 /// memory until it reaches a configured size, at which point the data is
13 /// written to a temporary file on disk, and further operations use the file
14 /// on disk.
15 #[derive(Debug)]
16 pub struct SpooledTempFile {
17 max_size: usize,
18 inner: SpooledInner,
19 }
20
21 /// Create a new spooled temporary file.
22 ///
23 /// # Security
24 ///
25 /// This variant is secure/reliable in the presence of a pathological temporary
26 /// file cleaner.
27 ///
28 /// # Resource Leaking
29 ///
30 /// The temporary file will be automatically removed by the OS when the last
31 /// handle to it is closed. This doesn't rely on Rust destructors being run, so
32 /// will (almost) never fail to clean up the temporary file.
33 ///
34 /// # Examples
35 ///
36 /// ```
37 /// use tempfile::spooled_tempfile;
38 /// use std::io::{self, Write};
39 ///
40 /// # fn main() {
41 /// # if let Err(_) = run() {
42 /// # ::std::process::exit(1);
43 /// # }
44 /// # }
45 /// # fn run() -> Result<(), io::Error> {
46 /// let mut file = spooled_tempfile(15);
47 ///
48 /// writeln!(file, "short line")?;
49 /// assert!(!file.is_rolled());
50 ///
51 /// // as a result of this write call, the size of the data will exceed
52 /// // `max_size` (15), so it will be written to a temporary file on disk,
53 /// // and the in-memory buffer will be dropped
54 /// writeln!(file, "marvin gardens")?;
55 /// assert!(file.is_rolled());
56 ///
57 /// # Ok(())
58 /// # }
59 /// ```
60 #[inline]
spooled_tempfile(max_size: usize) -> SpooledTempFile61 pub fn spooled_tempfile(max_size: usize) -> SpooledTempFile {
62 SpooledTempFile::new(max_size)
63 }
64
65 impl SpooledTempFile {
new(max_size: usize) -> SpooledTempFile66 pub fn new(max_size: usize) -> SpooledTempFile {
67 SpooledTempFile {
68 max_size: max_size,
69 inner: SpooledInner::InMemory(Cursor::new(Vec::new())),
70 }
71 }
72
73 /// Returns true if the file has been rolled over to disk.
is_rolled(&self) -> bool74 pub fn is_rolled(&self) -> bool {
75 match self.inner {
76 SpooledInner::InMemory(_) => false,
77 SpooledInner::OnDisk(_) => true,
78 }
79 }
80
81 /// Rolls over to a file on disk, regardless of current size. Does nothing
82 /// if already rolled over.
roll(&mut self) -> io::Result<()>83 pub fn roll(&mut self) -> io::Result<()> {
84 if !self.is_rolled() {
85 let mut file = tempfile()?;
86 if let SpooledInner::InMemory(ref mut cursor) = self.inner {
87 file.write_all(cursor.get_ref())?;
88 file.seek(SeekFrom::Start(cursor.position()))?;
89 }
90 self.inner = SpooledInner::OnDisk(file);
91 }
92 Ok(())
93 }
94
set_len(&mut self, size: u64) -> Result<(), io::Error>95 pub fn set_len(&mut self, size: u64) -> Result<(), io::Error> {
96 if size as usize > self.max_size {
97 self.roll()?; // does nothing if already rolled over
98 }
99 match self.inner {
100 SpooledInner::InMemory(ref mut cursor) => {
101 cursor.get_mut().resize(size as usize, 0);
102 Ok(())
103 }
104 SpooledInner::OnDisk(ref mut file) => file.set_len(size),
105 }
106 }
107 }
108
109 impl Read for SpooledTempFile {
read(&mut self, buf: &mut [u8]) -> io::Result<usize>110 fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
111 match self.inner {
112 SpooledInner::InMemory(ref mut cursor) => cursor.read(buf),
113 SpooledInner::OnDisk(ref mut file) => file.read(buf),
114 }
115 }
116 }
117
118 impl Write for SpooledTempFile {
write(&mut self, buf: &[u8]) -> io::Result<usize>119 fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
120 // roll over to file if necessary
121 let mut rolling = false;
122 if let SpooledInner::InMemory(ref mut cursor) = self.inner {
123 rolling = cursor.position() as usize + buf.len() > self.max_size;
124 }
125 if rolling {
126 self.roll()?;
127 }
128
129 // write the bytes
130 match self.inner {
131 SpooledInner::InMemory(ref mut cursor) => cursor.write(buf),
132 SpooledInner::OnDisk(ref mut file) => file.write(buf),
133 }
134 }
135
136 #[inline]
flush(&mut self) -> io::Result<()>137 fn flush(&mut self) -> io::Result<()> {
138 match self.inner {
139 SpooledInner::InMemory(ref mut cursor) => cursor.flush(),
140 SpooledInner::OnDisk(ref mut file) => file.flush(),
141 }
142 }
143 }
144
145 impl Seek for SpooledTempFile {
seek(&mut self, pos: SeekFrom) -> io::Result<u64>146 fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
147 match self.inner {
148 SpooledInner::InMemory(ref mut cursor) => cursor.seek(pos),
149 SpooledInner::OnDisk(ref mut file) => file.seek(pos),
150 }
151 }
152 }
153