]>
Commit | Line | Data |
---|---|---|
51b499db DM |
1 | //! Tools and utilities |
2 | //! | |
3 | //! This is a collection of small and useful tools. | |
f12f8ff1 DM |
4 | use failure::*; |
5 | use nix::unistd; | |
6 | use nix::sys::stat; | |
7 | ||
74a69302 DM |
8 | use lazy_static::lazy_static; |
9 | ||
365bb90f | 10 | use std::fs::{File, OpenOptions}; |
f12f8ff1 DM |
11 | use std::io::Write; |
12 | use std::path::Path; | |
43eeef28 DM |
13 | use std::io::Read; |
14 | use std::io::ErrorKind; | |
1628a4c7 | 15 | use std::time::Duration; |
f12f8ff1 | 16 | |
443f3743 | 17 | use std::os::unix::io::RawFd; |
eae8aa3a | 18 | use std::os::unix::io::AsRawFd; |
365bb90f | 19 | |
0fe5d605 DM |
20 | use serde_json::Value; |
21 | ||
8cf6e764 | 22 | pub mod timer; |
7f0d67cf | 23 | pub mod wrapped_reader_stream; |
8f973f81 DM |
24 | #[macro_use] |
25 | pub mod common_regex; | |
8cf6e764 | 26 | |
0a72e267 DM |
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 { | |
318564ac DM |
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. | |
0a72e267 DM |
34 | fn buffered_read(&mut self, offset: u64) -> Result<&[u8], Error>; |
35 | } | |
36 | ||
51b499db DM |
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. | |
50ea4396 | 43 | pub fn map_struct<T>(buffer: &[u8]) -> Result<&T, Error> { |
dc3de618 DM |
44 | if buffer.len() < ::std::mem::size_of::<T>() { |
45 | bail!("unable to map struct - buffer too small"); | |
46 | } | |
95bd5dfe | 47 | Ok(unsafe { & * (buffer.as_ptr() as *const T) }) |
dc3de618 DM |
48 | } |
49 | ||
51b499db DM |
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. | |
50ea4396 | 56 | pub fn map_struct_mut<T>(buffer: &mut [u8]) -> Result<&mut T, Error> { |
dc3de618 DM |
57 | if buffer.len() < ::std::mem::size_of::<T>() { |
58 | bail!("unable to map struct - buffer too small"); | |
59 | } | |
95bd5dfe | 60 | Ok(unsafe { &mut * (buffer.as_ptr() as *mut T) }) |
dc3de618 DM |
61 | } |
62 | ||
447787ab DM |
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 | ||
53157ca6 DM |
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 | ||
51b499db DM |
91 | /// Atomically write a file. We first create a temporary file, which |
92 | /// is then renamed. | |
f12f8ff1 DM |
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 | ||
d64d80d2 DM |
101 | // Note: we use mkstemp heŕe, because this worka with different |
102 | // processes, threads, and even tokio tasks. | |
f12f8ff1 DM |
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 | ||
1a7bc3dd | 112 | let mode : stat::Mode = perm.unwrap_or(stat::Mode::from( |
f12f8ff1 DM |
113 | stat::Mode::S_IRUSR | stat::Mode::S_IWUSR | |
114 | stat::Mode::S_IRGRP | stat::Mode::S_IROTH | |
1a7bc3dd | 115 | )); |
f12f8ff1 DM |
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 | } | |
43eeef28 | 137 | |
51b499db DM |
138 | /// Create a file lock using fntl. This function allows you to specify |
139 | /// a timeout if you want to avoid infinite blocking. | |
1628a4c7 WB |
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 | } | |
365bb90f | 176 | |
51b499db DM |
177 | /// Open or create a lock file (append mode). Then try to |
178 | /// aquire a lock using `lock_file()`. | |
1628a4c7 WB |
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 | { | |
365bb90f DM |
189 | Ok(file) => file, |
190 | Err(err) => bail!("Unable to open lock {:?} - {}", | |
191 | path, err), | |
192 | }; | |
28b96b56 DM |
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 | } | |
365bb90f DM |
198 | } |
199 | ||
51b499db DM |
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) | |
43eeef28 DM |
203 | pub fn file_chunker<C, R>( |
204 | mut file: R, | |
205 | chunk_size: usize, | |
606ce64b | 206 | mut chunk_cb: C |
43eeef28 | 207 | ) -> Result<(), Error> |
606ce64b | 208 | where C: FnMut(usize, &[u8]) -> Result<bool, Error>, |
43eeef28 DM |
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 */ } | |
5f0c2d56 | 233 | Err(e) => bail!("read chunk failed - {}", e.to_string()), |
43eeef28 DM |
234 | } |
235 | } | |
43eeef28 DM |
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(()) | |
43eeef28 | 261 | } |
0fe5d605 | 262 | |
443f3743 | 263 | /// Returns the hosts node name (UTS node name) |
74a69302 DM |
264 | pub fn nodename() -> &'static str { |
265 | ||
266 | lazy_static!{ | |
267 | static ref NODENAME: String = { | |
268 | ||
0d38dcb4 DM |
269 | nix::sys::utsname::uname() |
270 | .nodename() | |
271 | .split('.') | |
272 | .next() | |
273 | .unwrap() | |
274 | .to_owned() | |
74a69302 DM |
275 | }; |
276 | } | |
277 | ||
278 | &NODENAME | |
279 | } | |
280 | ||
0fe5d605 DM |
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 | } | |
0d38dcb4 DM |
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 | } | |
383e8577 DM |
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 | } | |
443f3743 DM |
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: ®ex::Regex, | |
cce1676a | 358 | mut callback: F |
443f3743 | 359 | ) -> Result<(), Error> |
cce1676a | 360 | where F: FnMut(RawFd, &str, nix::dir::Type) -> Result<(), Error>, |
443f3743 DM |
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 | } | |
7e13b2d6 DM |
387 | |
388 | pub fn get_hardware_address() -> Result<String, Error> { | |
389 | ||
390 | static FILENAME: &str = "/etc/ssh/assh_host_rsa_key.pub"; | |
391 | ||
392 | let mut file = File::open(FILENAME)?; | |
393 | let mut contents = Vec::new(); | |
394 | file.read_to_end(&mut contents)?; | |
395 | ||
396 | let digest = md5::compute(contents); | |
397 | ||
398 | Ok(format!("{:0x}", digest)) | |
399 | } |