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