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