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