1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153

use file::tempfile;
use std::fs::File;
use std::io::{self, Cursor, Read, Seek, SeekFrom, Write};

#[derive(Debug)]
enum SpooledInner {
    InMemory(Cursor<Vec<u8>>),
    OnDisk(File),
}

/// An object that behaves like a regular temporary file, but keeps data in
/// memory until it reaches a configured size, at which point the data is
/// written to a temporary file on disk, and further operations use the file
/// on disk.
#[derive(Debug)]
pub struct SpooledTempFile {
    max_size: usize,
    inner: SpooledInner,
}

/// Create a new spooled temporary file.
///
/// # Security
///
/// This variant is secure/reliable in the presence of a pathological temporary
/// file cleaner.
///
/// # Resource Leaking
///
/// The temporary file will be automatically removed by the OS when the last
/// handle to it is closed. This doesn't rely on Rust destructors being run, so
/// will (almost) never fail to clean up the temporary file.
///
/// # Examples
///
/// ```
/// # extern crate tempfile;
/// use tempfile::spooled_tempfile;
/// use std::io::{self, Write};
///
/// # fn main() {
/// #     if let Err(_) = run() {
/// #         ::std::process::exit(1);
/// #     }
/// # }
/// # fn run() -> Result<(), io::Error> {
/// let mut file = spooled_tempfile(15);
///
/// writeln!(file, "short line")?;
/// assert!(!file.is_rolled());
///
/// // as a result of this write call, the size of the data will exceed
/// // `max_size` (15), so it will be written to a temporary file on disk,
/// // and the in-memory buffer will be dropped
/// writeln!(file, "marvin gardens")?;
/// assert!(file.is_rolled());
///
/// # Ok(())
/// # }
/// ```
#[inline]
pub fn spooled_tempfile(max_size: usize) -> SpooledTempFile {
    SpooledTempFile::new(max_size)
}

impl SpooledTempFile {
    pub fn new(max_size: usize) -> SpooledTempFile {
        SpooledTempFile {
            max_size: max_size,
            inner: SpooledInner::InMemory(Cursor::new(Vec::new())),
        }
    }

    /// Returns true if the file has been rolled over to disk.
    pub fn is_rolled(&self) -> bool {
        match self.inner {
            SpooledInner::InMemory(_) => false,
            SpooledInner::OnDisk(_) => true,
        }
    }

    /// Rolls over to a file on disk, regardless of current size. Does nothing
    /// if already rolled over.
    pub fn roll(&mut self) -> io::Result<()> {
        if !self.is_rolled() {
            let mut file = tempfile()?;
            if let SpooledInner::InMemory(ref mut cursor) = self.inner {
                file.write_all(cursor.get_ref())?;
                file.seek(SeekFrom::Start(cursor.position()))?;
            }
            self.inner = SpooledInner::OnDisk(file);
        }
        Ok(())
    }

    pub fn set_len(&mut self, size: u64) -> Result<(), io::Error> {
        if size as usize > self.max_size {
            self.roll()?; // does nothing if already rolled over
        }
        match self.inner {
            SpooledInner::InMemory(ref mut cursor) => {
                cursor.get_mut().resize(size as usize, 0);
                Ok(())
            }
            SpooledInner::OnDisk(ref mut file) => file.set_len(size),
        }
    }
}

impl Read for SpooledTempFile {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        match self.inner {
            SpooledInner::InMemory(ref mut cursor) => cursor.read(buf),
            SpooledInner::OnDisk(ref mut file) => file.read(buf),
        }
    }
}

impl Write for SpooledTempFile {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        // roll over to file if necessary
        let mut rolling = false;
        if let SpooledInner::InMemory(ref mut cursor) = self.inner {
            rolling = cursor.position() as usize + buf.len() > self.max_size;
        }
        if rolling {
            self.roll()?;
        }

        // write the bytes
        match self.inner {
            SpooledInner::InMemory(ref mut cursor) => cursor.write(buf),
            SpooledInner::OnDisk(ref mut file) => file.write(buf),
        }
    }

    #[inline]
    fn flush(&mut self) -> io::Result<()> {
        match self.inner {
            SpooledInner::InMemory(ref mut cursor) => cursor.flush(),
            SpooledInner::OnDisk(ref mut file) => file.flush(),
        }
    }
}

impl Seek for SpooledTempFile {
    fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
        match self.inner {
            SpooledInner::InMemory(ref mut cursor) => cursor.seek(pos),
            SpooledInner::OnDisk(ref mut file) => file.seek(pos),
        }
    }
}