src/tools.rs

   1 //! Tools and utilities
   2 //!
   3 //! This is a collection of small and useful tools.
   4
   5 use failure::*;
   6 use nix::unistd;
   7 use nix::sys::stat;
   8
   9 use std::fs::{File, OpenOptions};
  10 use std::io::Write;
  11 use std::path::Path;
  12 use std::io::Read;
  13 use std::io::ErrorKind;
  14 use std::time::Duration;
  15
  16 use std::os::unix::io::AsRawFd;
  17
  18 use serde_json::Value;
  19
  20 pub mod timer;
  21
  22 /// The `BufferedReader` trait provides a single function
  23 /// `buffered_read`. It returns a reference to an internal buffer. The
  24 /// purpose of this traid is to avoid unnecessary data copies.
  25 pub trait BufferedReader {
  26     /// This functions tries to fill the internal buffers, then
  27     /// returns a reference to the available data. It returns an empty
  28     /// buffer if `offset` points to the end of the file.
  29     fn buffered_read(&mut self, offset: u64) -> Result<&[u8], Error>;
  30 }
  31
  32 /// Directly map a type into a binary buffer. This is mostly useful
  33 /// for reading structured data from a byte stream (file). You need to
  34 /// make sure that the buffer location does not change, so please
  35 /// avoid vec resize while you use such map.
  36 ///
  37 /// This function panics if the buffer is not large enough.
  38 pub fn map_struct<T>(buffer: &[u8]) -> Result<&T, Error> {
  39     if buffer.len() < ::std::mem::size_of::<T>() {
  40         bail!("unable to map struct - buffer too small");
  41     }
  42     Ok(unsafe { & * (buffer.as_ptr() as *const T) })
  43 }
  44
  45 /// Directly map a type into a mutable binary buffer. This is mostly
  46 /// useful for writing structured data into a byte stream (file). You
  47 /// need to make sure that the buffer location does not change, so
  48 /// please avoid vec resize while you use such map.
  49 ///
  50 /// This function panics if the buffer is not large enough.
  51 pub fn map_struct_mut<T>(buffer: &mut [u8]) -> Result<&mut T, Error> {
  52     if buffer.len() < ::std::mem::size_of::<T>() {
  53         bail!("unable to map struct - buffer too small");
  54     }
  55     Ok(unsafe { &mut * (buffer.as_ptr() as *mut T) })
  56 }
  57
  58 /// Atomically write a file. We first create a temporary file, which
  59 /// is then renamed.
  60 pub fn file_set_contents<P: AsRef<Path>>(
  61     path: P,
  62     data: &[u8],
  63     perm: Option<stat::Mode>,
  64 ) -> Result<(), Error> {
  65
  66     let path = path.as_ref();
  67
  68     // Note: we use mkstemp heŕe, because this worka with different
  69     // processes, threads, and even tokio tasks.
  70     let mut template = path.to_owned();
  71     template.set_extension("tmp_XXXXXX");
  72     let (fd, tmp_path) = match unistd::mkstemp(&template) {
  73         Ok((fd, path)) => (fd, path),
  74         Err(err) => bail!("mkstemp {:?} failed: {}", template, err),
  75     };
  76
  77     let tmp_path = tmp_path.as_path();
  78
  79     let mode : stat::Mode = perm.unwrap_or(stat::Mode::from(
  80         stat::Mode::S_IRUSR | stat::Mode::S_IWUSR |
  81         stat::Mode::S_IRGRP | stat::Mode::S_IROTH
  82     ));
  83
  84     if let Err(err) = stat::fchmod(fd, mode) {
  85         let _ = unistd::unlink(tmp_path);
  86         bail!("fchmod {:?} failed: {}", tmp_path, err);
  87     }
  88
  89     use std::os::unix::io::FromRawFd;
  90     let mut file = unsafe { File::from_raw_fd(fd) };
  91
  92     if let Err(err) = file.write_all(data) {
  93         let _ = unistd::unlink(tmp_path);
  94         bail!("write failed: {}", err);
  95     }
  96
  97     if let Err(err) = std::fs::rename(tmp_path, path) {
  98         let _ = unistd::unlink(tmp_path);
  99         bail!("Atomic rename failed for file {:?} - {}", path, err);
 100     }
 101
 102     Ok(())
 103 }
 104
 105 /// Create a file lock using fntl. This function allows you to specify
 106 /// a timeout if you want to avoid infinite blocking.
 107 pub fn lock_file<F: AsRawFd>(
 108     file: &mut F,
 109     exclusive: bool,
 110     timeout: Option<Duration>,
 111     ) -> Result<(), Error>
 112 {
 113     let lockarg =
 114         if exclusive {
 115             nix::fcntl::FlockArg::LockExclusive
 116         } else {
 117             nix::fcntl::FlockArg::LockShared
 118         };
 119
 120     let timeout = match timeout {
 121         None => {
 122             nix::fcntl::flock(file.as_raw_fd(), lockarg)?;
 123             return Ok(());
 124         }
 125         Some(t) => t,
 126     };
 127
 128     // unblock the timeout signal temporarily
 129     let _sigblock_guard = timer::unblock_timeout_signal();
 130
 131     // setup a timeout timer
 132     let mut timer = timer::Timer::create(
 133         timer::Clock::Realtime,
 134         timer::TimerEvent::ThisThreadSignal(timer::SIGTIMEOUT))?;
 135
 136     timer.arm(timer::TimerSpec::new()
 137         .value(Some(timeout))
 138         .interval(Some(Duration::from_millis(10))))?;
 139
 140     nix::fcntl::flock(file.as_raw_fd(), lockarg)?;
 141     Ok(())
 142 }
 143
 144 /// Open or create a lock file (append mode). Then try to
 145 /// aquire a lock using `lock_file()`.
 146 pub fn open_file_locked<P: AsRef<Path>>(path: P, timeout: Duration)
 147     -> Result<File, Error>
 148 {
 149     let path = path.as_ref();
 150     let mut file =
 151         match OpenOptions::new()
 152             .create(true)
 153             .append(true)
 154             .open(path)
 155         {
 156             Ok(file) => file,
 157             Err(err) => bail!("Unable to open lock {:?} - {}",
 158                               path, err),
 159         };
 160     match lock_file(&mut file, true, Some(timeout)) {
 161         Ok(_) => Ok(file),
 162         Err(err) => bail!("Unable to aquire lock {:?} - {}",
 163                           path, err),
 164     }
 165 }
 166
 167 /// Split a file into equal sized chunks. The last chunk may be
 168 /// smaller. Note: We cannot implement an `Iterator`, because iterators
 169 /// cannot return a borrowed buffer ref (we want zero-copy)
 170 pub fn file_chunker<C, R>(
 171     mut file: R,
 172     chunk_size: usize,
 173     mut chunk_cb: C
 174 ) -> Result<(), Error>
 175     where C: FnMut(usize, &[u8]) -> Result<bool, Error>,
 176           R: Read,
 177 {
 178
 179     const READ_BUFFER_SIZE: usize = 4*1024*1024; // 4M
 180
 181     if chunk_size > READ_BUFFER_SIZE { bail!("chunk size too large!"); }
 182
 183     let mut buf = vec![0u8; READ_BUFFER_SIZE];
 184
 185     let mut pos = 0;
 186     let mut file_pos = 0;
 187     loop {
 188         let mut eof = false;
 189         let mut tmp = &mut buf[..];
 190        // try to read large portions, at least chunk_size
 191         while pos < chunk_size {
 192             match file.read(tmp) {
 193                 Ok(0) => { eof = true; break; },
 194                 Ok(n) => {
 195                     pos += n;
 196                     if pos > chunk_size { break; }
 197                     tmp = &mut tmp[n..];
 198                 }
 199                 Err(ref e) if e.kind() == ErrorKind::Interrupted => { /* try again */ }
 200                 Err(e) => bail!("read chunk failed - {}", e.to_string()),
 201             }
 202         }
 203         let mut start = 0;
 204         while start + chunk_size <= pos {
 205             if !(chunk_cb)(file_pos, &buf[start..start+chunk_size])? { break; }
 206             file_pos += chunk_size;
 207             start += chunk_size;
 208         }
 209         if eof {
 210             if start < pos {
 211                 (chunk_cb)(file_pos, &buf[start..pos])?;
 212                 //file_pos += pos - start;
 213             }
 214             break;
 215         } else {
 216             let rest = pos - start;
 217             if rest > 0 {
 218                 let ptr = buf.as_mut_ptr();
 219                 unsafe { std::ptr::copy_nonoverlapping(ptr.add(start), ptr, rest); }
 220                 pos = rest;
 221             } else {
 222                 pos = 0;
 223             }
 224         }
 225     }
 226
 227     Ok(())
 228 }
 229
 230 pub fn required_string_param<'a>(param: &'a Value, name: &str) -> Result<&'a str, Error> {
 231     match param[name].as_str()   {
 232         Some(s) => Ok(s),
 233         None => bail!("missing parameter '{}'", name),
 234     }
 235 }