src/tools.rs

   1 //! Tools and utilities
   2 //!
   3 //! This is a collection of small and useful tools.
   4 use failure::*;
   5 use nix::unistd;
   6 use nix::sys::stat;
   7
   8 use lazy_static::lazy_static;
   9
  10 use std::fs::{File, OpenOptions};
  11 use std::io::Write;
  12 use std::path::Path;
  13 use std::io::Read;
  14 use std::io::ErrorKind;
  15 use std::time::Duration;
  16
  17 use std::os::unix::io::RawFd;
  18 use std::os::unix::io::AsRawFd;
  19
  20 use serde_json::Value;
  21
  22 pub mod timer;
  23 pub mod wrapped_reader_stream;
  24 #[macro_use]
  25 pub mod common_regex;
  26
  27 /// The `BufferedReader` trait provides a single function
  28 /// `buffered_read`. It returns a reference to an internal buffer. The
  29 /// purpose of this traid is to avoid unnecessary data copies.
  30 pub trait BufferedReader {
  31     /// This functions tries to fill the internal buffers, then
  32     /// returns a reference to the available data. It returns an empty
  33     /// buffer if `offset` points to the end of the file.
  34     fn buffered_read(&mut self, offset: u64) -> Result<&[u8], Error>;
  35 }
  36
  37 /// Directly map a type into a binary buffer. This is mostly useful
  38 /// for reading structured data from a byte stream (file). You need to
  39 /// make sure that the buffer location does not change, so please
  40 /// avoid vec resize while you use such map.
  41 ///
  42 /// This function panics if the buffer is not large enough.
  43 pub fn map_struct<T>(buffer: &[u8]) -> Result<&T, Error> {
  44     if buffer.len() < ::std::mem::size_of::<T>() {
  45         bail!("unable to map struct - buffer too small");
  46     }
  47     Ok(unsafe { & * (buffer.as_ptr() as *const T) })
  48 }
  49
  50 /// Directly map a type into a mutable binary buffer. This is mostly
  51 /// useful for writing structured data into a byte stream (file). You
  52 /// need to make sure that the buffer location does not change, so
  53 /// please avoid vec resize while you use such map.
  54 ///
  55 /// This function panics if the buffer is not large enough.
  56 pub fn map_struct_mut<T>(buffer: &mut [u8]) -> Result<&mut T, Error> {
  57     if buffer.len() < ::std::mem::size_of::<T>() {
  58         bail!("unable to map struct - buffer too small");
  59     }
  60     Ok(unsafe { &mut * (buffer.as_ptr() as *mut T) })
  61 }
  62
  63 pub fn file_read_firstline<P: AsRef<Path>>(path: P) -> Result<String, std::io::Error> {
  64
  65     let path = path.as_ref();
  66
  67     let file = std::fs::File::open(path)?;
  68
  69     use std::io::{BufRead, BufReader};
  70
  71     let mut reader = BufReader::new(file);
  72
  73     let mut line = String::new();
  74
  75     let _ = reader.read_line(&mut line)?;
  76
  77     Ok(line)
  78 }
  79
  80 pub fn file_get_contents<P: AsRef<Path>>(path: P) -> Result<Vec<u8>, std::io::Error> {
  81
  82     let mut file = std::fs::File::open(path)?;
  83
  84     let mut buffer = Vec::new();
  85
  86     file.read_to_end(&mut buffer)?;
  87
  88     Ok(buffer)
  89 }
  90
  91 /// Atomically write a file. We first create a temporary file, which
  92 /// is then renamed.
  93 pub fn file_set_contents<P: AsRef<Path>>(
  94     path: P,
  95     data: &[u8],
  96     perm: Option<stat::Mode>,
  97 ) -> Result<(), Error> {
  98
  99     let path = path.as_ref();
 100
 101     // Note: we use mkstemp heŕe, because this worka with different
 102     // processes, threads, and even tokio tasks.
 103     let mut template = path.to_owned();
 104     template.set_extension("tmp_XXXXXX");
 105     let (fd, tmp_path) = match unistd::mkstemp(&template) {
 106         Ok((fd, path)) => (fd, path),
 107         Err(err) => bail!("mkstemp {:?} failed: {}", template, err),
 108     };
 109
 110     let tmp_path = tmp_path.as_path();
 111
 112     let mode : stat::Mode = perm.unwrap_or(stat::Mode::from(
 113         stat::Mode::S_IRUSR | stat::Mode::S_IWUSR |
 114         stat::Mode::S_IRGRP | stat::Mode::S_IROTH
 115     ));
 116
 117     if let Err(err) = stat::fchmod(fd, mode) {
 118         let _ = unistd::unlink(tmp_path);
 119         bail!("fchmod {:?} failed: {}", tmp_path, err);
 120     }
 121
 122     use std::os::unix::io::FromRawFd;
 123     let mut file = unsafe { File::from_raw_fd(fd) };
 124
 125     if let Err(err) = file.write_all(data) {
 126         let _ = unistd::unlink(tmp_path);
 127         bail!("write failed: {}", err);
 128     }
 129
 130     if let Err(err) = std::fs::rename(tmp_path, path) {
 131         let _ = unistd::unlink(tmp_path);
 132         bail!("Atomic rename failed for file {:?} - {}", path, err);
 133     }
 134
 135     Ok(())
 136 }
 137
 138 /// Create a file lock using fntl. This function allows you to specify
 139 /// a timeout if you want to avoid infinite blocking.
 140 pub fn lock_file<F: AsRawFd>(
 141     file: &mut F,
 142     exclusive: bool,
 143     timeout: Option<Duration>,
 144     ) -> Result<(), Error>
 145 {
 146     let lockarg =
 147         if exclusive {
 148             nix::fcntl::FlockArg::LockExclusive
 149         } else {
 150             nix::fcntl::FlockArg::LockShared
 151         };
 152
 153     let timeout = match timeout {
 154         None => {
 155             nix::fcntl::flock(file.as_raw_fd(), lockarg)?;
 156             return Ok(());
 157         }
 158         Some(t) => t,
 159     };
 160
 161     // unblock the timeout signal temporarily
 162     let _sigblock_guard = timer::unblock_timeout_signal();
 163
 164     // setup a timeout timer
 165     let mut timer = timer::Timer::create(
 166         timer::Clock::Realtime,
 167         timer::TimerEvent::ThisThreadSignal(timer::SIGTIMEOUT))?;
 168
 169     timer.arm(timer::TimerSpec::new()
 170         .value(Some(timeout))
 171         .interval(Some(Duration::from_millis(10))))?;
 172
 173     nix::fcntl::flock(file.as_raw_fd(), lockarg)?;
 174     Ok(())
 175 }
 176
 177 /// Open or create a lock file (append mode). Then try to
 178 /// aquire a lock using `lock_file()`.
 179 pub fn open_file_locked<P: AsRef<Path>>(path: P, timeout: Duration)
 180     -> Result<File, Error>
 181 {
 182     let path = path.as_ref();
 183     let mut file =
 184         match OpenOptions::new()
 185             .create(true)
 186             .append(true)
 187             .open(path)
 188         {
 189             Ok(file) => file,
 190             Err(err) => bail!("Unable to open lock {:?} - {}",
 191                               path, err),
 192         };
 193     match lock_file(&mut file, true, Some(timeout)) {
 194         Ok(_) => Ok(file),
 195         Err(err) => bail!("Unable to aquire lock {:?} - {}",
 196                           path, err),
 197     }
 198 }
 199
 200 /// Split a file into equal sized chunks. The last chunk may be
 201 /// smaller. Note: We cannot implement an `Iterator`, because iterators
 202 /// cannot return a borrowed buffer ref (we want zero-copy)
 203 pub fn file_chunker<C, R>(
 204     mut file: R,
 205     chunk_size: usize,
 206     mut chunk_cb: C
 207 ) -> Result<(), Error>
 208     where C: FnMut(usize, &[u8]) -> Result<bool, Error>,
 209           R: Read,
 210 {
 211
 212     const READ_BUFFER_SIZE: usize = 4*1024*1024; // 4M
 213
 214     if chunk_size > READ_BUFFER_SIZE { bail!("chunk size too large!"); }
 215
 216     let mut buf = vec![0u8; READ_BUFFER_SIZE];
 217
 218     let mut pos = 0;
 219     let mut file_pos = 0;
 220     loop {
 221         let mut eof = false;
 222         let mut tmp = &mut buf[..];
 223        // try to read large portions, at least chunk_size
 224         while pos < chunk_size {
 225             match file.read(tmp) {
 226                 Ok(0) => { eof = true; break; },
 227                 Ok(n) => {
 228                     pos += n;
 229                     if pos > chunk_size { break; }
 230                     tmp = &mut tmp[n..];
 231                 }
 232                 Err(ref e) if e.kind() == ErrorKind::Interrupted => { /* try again */ }
 233                 Err(e) => bail!("read chunk failed - {}", e.to_string()),
 234             }
 235         }
 236         let mut start = 0;
 237         while start + chunk_size <= pos {
 238             if !(chunk_cb)(file_pos, &buf[start..start+chunk_size])? { break; }
 239             file_pos += chunk_size;
 240             start += chunk_size;
 241         }
 242         if eof {
 243             if start < pos {
 244                 (chunk_cb)(file_pos, &buf[start..pos])?;
 245                 //file_pos += pos - start;
 246             }
 247             break;
 248         } else {
 249             let rest = pos - start;
 250             if rest > 0 {
 251                 let ptr = buf.as_mut_ptr();
 252                 unsafe { std::ptr::copy_nonoverlapping(ptr.add(start), ptr, rest); }
 253                 pos = rest;
 254             } else {
 255                 pos = 0;
 256             }
 257         }
 258     }
 259
 260     Ok(())
 261 }
 262
 263 /// Returns the hosts node name (UTS node name)
 264 pub fn nodename() -> &'static str {
 265
 266     lazy_static!{
 267         static ref NODENAME: String = {
 268
 269             nix::sys::utsname::uname()
 270                 .nodename()
 271                 .split('.')
 272                 .next()
 273                 .unwrap()
 274                 .to_owned()
 275         };
 276     }
 277
 278     &NODENAME
 279 }
 280
 281 pub fn required_string_param<'a>(param: &'a Value, name: &str) -> Result<&'a str, Error> {
 282     match param[name].as_str()   {
 283         Some(s) => Ok(s),
 284         None => bail!("missing parameter '{}'", name),
 285     }
 286 }
 287
 288 pub fn required_integer_param<'a>(param: &'a Value, name: &str) -> Result<i64, Error> {
 289     match param[name].as_i64()   {
 290         Some(s) => Ok(s),
 291         None => bail!("missing parameter '{}'", name),
 292     }
 293 }
 294
 295 pub fn complete_file_name(arg: &str) -> Vec<String> {
 296
 297     let mut result = vec![];
 298
 299     use nix::fcntl::OFlag;
 300     use nix::sys::stat::Mode;
 301     use nix::fcntl::AtFlags;
 302
 303     let mut dirname = std::path::PathBuf::from(arg);
 304
 305     let is_dir = match nix::sys::stat::fstatat(libc::AT_FDCWD, &dirname, AtFlags::empty()) {
 306         Ok(stat) => (stat.st_mode & libc::S_IFMT) == libc::S_IFDIR,
 307         Err(_) => false,
 308     };
 309
 310     if !is_dir {
 311         if let Some(parent) = dirname.parent() {
 312             dirname = parent.to_owned();
 313         }
 314     }
 315
 316     let mut dir = match nix::dir::Dir::openat(libc::AT_FDCWD, &dirname, OFlag::O_DIRECTORY, Mode::empty()) {
 317         Ok(d) => d,
 318         Err(_) => return result,
 319     };
 320
 321     for item in dir.iter() {
 322         if let Ok(entry) = item {
 323             if let Ok(name) = entry.file_name().to_str() {
 324                 if name == "." || name == ".." { continue; }
 325                 let mut newpath = dirname.clone();
 326                 newpath.push(name);
 327
 328                 if let Ok(stat) = nix::sys::stat::fstatat(libc::AT_FDCWD, &newpath, AtFlags::empty()) {
 329                     if (stat.st_mode & libc::S_IFMT) == libc::S_IFDIR {
 330                         newpath.push("");
 331                         if let Some(newpath) = newpath.to_str() {
 332                             result.push(newpath.to_owned());
 333                         }
 334                         continue;
 335                      }
 336                 }
 337                 if let Some(newpath) = newpath.to_str() {
 338                     result.push(newpath.to_owned());
 339                 }
 340
 341              }
 342         }
 343     }
 344
 345     result
 346 }
 347
 348 /// Scan directory for matching file names.
 349 ///
 350 /// Scan through all directory entries and call `callback()` function
 351 /// if the entry name matches the regular expression. This function
 352 /// used unix `openat()`, so you can pass absolute or relative file
 353 /// names. This function simply skips non-UTF8 encoded names.
 354 pub fn scandir<P, F>(
 355     dirfd: RawFd,
 356     path: P,
 357     regex: &regex::Regex,
 358     mut callback: F
 359 ) -> Result<(), Error>
 360     where F: FnMut(RawFd, &str, nix::dir::Type) -> Result<(), Error>,
 361           P: AsRef<Path>
 362 {
 363     use nix::fcntl::OFlag;
 364     use nix::sys::stat::Mode;
 365
 366     let mut subdir = nix::dir::Dir::openat(dirfd, path.as_ref(), OFlag::O_RDONLY, Mode::empty())?;
 367     let subdir_fd = subdir.as_raw_fd();
 368
 369     for entry in subdir.iter() {
 370         let entry = entry?;
 371         let file_type = match entry.file_type() {
 372             Some(file_type) => file_type,
 373             None => bail!("unable to detect file type"),
 374         };
 375         let filename = entry.file_name();
 376         let filename_str = match filename.to_str() {
 377             Ok(name) => name,
 378             Err(_) => continue /* ignore non utf8 entries*/,
 379         };
 380
 381         if !regex.is_match(filename_str) { continue; }
 382
 383         (callback)(subdir_fd, filename_str, file_type)?;
 384     }
 385     Ok(())
 386 }
 387
 388 pub fn get_hardware_address() -> Result<String, Error> {
 389
 390     static FILENAME: &str = "/etc/ssh/ssh_host_rsa_key.pub";
 391
 392     let contents = file_get_contents(FILENAME)?;
 393     let digest = md5::compute(contents);
 394
 395     Ok(format!("{:0x}", digest))
 396 }
 397
 398 pub fn digest_to_hex(digest: &[u8]) -> String {
 399
 400     const HEX_CHARS: &'static [u8; 16] = b"0123456789abcdef";
 401
 402     let mut buf = Vec::<u8>::with_capacity(digest.len()*2);
 403
 404     for i in 0..digest.len() {
 405         buf.push(HEX_CHARS[(digest[i] >> 4) as usize]);
 406         buf.push(HEX_CHARS[(digest[i] & 0xf) as usize]);
 407     }
 408
 409     unsafe { String::from_utf8_unchecked(buf) }
 410 }
 411
 412 pub fn assert_if_modified(digest1: &str, digest2: &str) -> Result<(), Error> {
 413     if digest1 != digest2 {
 414         bail!("detected modified configuration - file changed by other user? Try again.");
 415     }
 416     Ok(())
 417 }