]>
Commit | Line | Data |
---|---|---|
f12f8ff1 DM |
1 | use failure::*; |
2 | use nix::unistd; | |
3 | use nix::sys::stat; | |
4 | ||
365bb90f | 5 | use std::fs::{File, OpenOptions}; |
f12f8ff1 DM |
6 | use std::io::Write; |
7 | use std::path::Path; | |
43eeef28 DM |
8 | use std::io::Read; |
9 | use std::io::ErrorKind; | |
1628a4c7 | 10 | use std::time::Duration; |
f12f8ff1 | 11 | |
eae8aa3a | 12 | use std::os::unix::io::AsRawFd; |
365bb90f | 13 | |
8cf6e764 WB |
14 | pub mod timer; |
15 | ||
50ea4396 | 16 | pub fn map_struct<T>(buffer: &[u8]) -> Result<&T, Error> { |
dc3de618 DM |
17 | if buffer.len() < ::std::mem::size_of::<T>() { |
18 | bail!("unable to map struct - buffer too small"); | |
19 | } | |
95bd5dfe | 20 | Ok(unsafe { & * (buffer.as_ptr() as *const T) }) |
dc3de618 DM |
21 | } |
22 | ||
50ea4396 | 23 | pub fn map_struct_mut<T>(buffer: &mut [u8]) -> Result<&mut T, Error> { |
dc3de618 DM |
24 | if buffer.len() < ::std::mem::size_of::<T>() { |
25 | bail!("unable to map struct - buffer too small"); | |
26 | } | |
95bd5dfe | 27 | Ok(unsafe { &mut * (buffer.as_ptr() as *mut T) }) |
dc3de618 DM |
28 | } |
29 | ||
30 | ||
f12f8ff1 DM |
31 | pub fn file_set_contents<P: AsRef<Path>>( |
32 | path: P, | |
33 | data: &[u8], | |
34 | perm: Option<stat::Mode>, | |
35 | ) -> Result<(), Error> { | |
36 | ||
37 | let path = path.as_ref(); | |
38 | ||
d64d80d2 DM |
39 | // Note: we use mkstemp heŕe, because this worka with different |
40 | // processes, threads, and even tokio tasks. | |
f12f8ff1 DM |
41 | let mut template = path.to_owned(); |
42 | template.set_extension("tmp_XXXXXX"); | |
43 | let (fd, tmp_path) = match unistd::mkstemp(&template) { | |
44 | Ok((fd, path)) => (fd, path), | |
45 | Err(err) => bail!("mkstemp {:?} failed: {}", template, err), | |
46 | }; | |
47 | ||
48 | let tmp_path = tmp_path.as_path(); | |
49 | ||
1a7bc3dd | 50 | let mode : stat::Mode = perm.unwrap_or(stat::Mode::from( |
f12f8ff1 DM |
51 | stat::Mode::S_IRUSR | stat::Mode::S_IWUSR | |
52 | stat::Mode::S_IRGRP | stat::Mode::S_IROTH | |
1a7bc3dd | 53 | )); |
f12f8ff1 DM |
54 | |
55 | if let Err(err) = stat::fchmod(fd, mode) { | |
56 | let _ = unistd::unlink(tmp_path); | |
57 | bail!("fchmod {:?} failed: {}", tmp_path, err); | |
58 | } | |
59 | ||
60 | use std::os::unix::io::FromRawFd; | |
61 | let mut file = unsafe { File::from_raw_fd(fd) }; | |
62 | ||
63 | if let Err(err) = file.write_all(data) { | |
64 | let _ = unistd::unlink(tmp_path); | |
65 | bail!("write failed: {}", err); | |
66 | } | |
67 | ||
68 | if let Err(err) = std::fs::rename(tmp_path, path) { | |
69 | let _ = unistd::unlink(tmp_path); | |
70 | bail!("Atomic rename failed for file {:?} - {}", path, err); | |
71 | } | |
72 | ||
73 | Ok(()) | |
74 | } | |
43eeef28 | 75 | |
1628a4c7 WB |
76 | pub fn lock_file<F: AsRawFd>( |
77 | file: &mut F, | |
78 | exclusive: bool, | |
79 | timeout: Option<Duration>, | |
80 | ) -> Result<(), Error> | |
81 | { | |
82 | let lockarg = | |
83 | if exclusive { | |
84 | nix::fcntl::FlockArg::LockExclusive | |
85 | } else { | |
86 | nix::fcntl::FlockArg::LockShared | |
87 | }; | |
88 | ||
89 | let timeout = match timeout { | |
90 | None => { | |
91 | nix::fcntl::flock(file.as_raw_fd(), lockarg)?; | |
92 | return Ok(()); | |
93 | } | |
94 | Some(t) => t, | |
95 | }; | |
96 | ||
97 | // unblock the timeout signal temporarily | |
98 | let _sigblock_guard = timer::unblock_timeout_signal(); | |
99 | ||
100 | // setup a timeout timer | |
101 | let mut timer = timer::Timer::create( | |
102 | timer::Clock::Realtime, | |
103 | timer::TimerEvent::ThisThreadSignal(timer::SIGTIMEOUT))?; | |
104 | ||
105 | timer.arm(timer::TimerSpec::new() | |
106 | .value(Some(timeout)) | |
107 | .interval(Some(Duration::from_millis(10))))?; | |
108 | ||
109 | nix::fcntl::flock(file.as_raw_fd(), lockarg)?; | |
110 | Ok(()) | |
111 | } | |
365bb90f | 112 | |
1628a4c7 WB |
113 | pub fn open_file_locked<P: AsRef<Path>>(path: P, timeout: Duration) |
114 | -> Result<File, Error> | |
115 | { | |
116 | let path = path.as_ref(); | |
117 | let mut file = | |
118 | match OpenOptions::new() | |
119 | .create(true) | |
120 | .append(true) | |
121 | .open(path) | |
122 | { | |
365bb90f DM |
123 | Ok(file) => file, |
124 | Err(err) => bail!("Unable to open lock {:?} - {}", | |
125 | path, err), | |
126 | }; | |
28b96b56 DM |
127 | match lock_file(&mut file, true, Some(timeout)) { |
128 | Ok(_) => Ok(file), | |
129 | Err(err) => bail!("Unable to aquire lock {:?} - {}", | |
130 | path, err), | |
131 | } | |
365bb90f DM |
132 | } |
133 | ||
43eeef28 DM |
134 | // Note: We cannot implement an Iterator, because Iterators cannot |
135 | // return a borrowed buffer ref (we want zero-copy) | |
136 | pub fn file_chunker<C, R>( | |
137 | mut file: R, | |
138 | chunk_size: usize, | |
606ce64b | 139 | mut chunk_cb: C |
43eeef28 | 140 | ) -> Result<(), Error> |
606ce64b | 141 | where C: FnMut(usize, &[u8]) -> Result<bool, Error>, |
43eeef28 DM |
142 | R: Read, |
143 | { | |
144 | ||
145 | const READ_BUFFER_SIZE: usize = 4*1024*1024; // 4M | |
146 | ||
147 | if chunk_size > READ_BUFFER_SIZE { bail!("chunk size too large!"); } | |
148 | ||
149 | let mut buf = vec![0u8; READ_BUFFER_SIZE]; | |
150 | ||
151 | let mut pos = 0; | |
152 | let mut file_pos = 0; | |
153 | loop { | |
154 | let mut eof = false; | |
155 | let mut tmp = &mut buf[..]; | |
156 | // try to read large portions, at least chunk_size | |
157 | while pos < chunk_size { | |
158 | match file.read(tmp) { | |
159 | Ok(0) => { eof = true; break; }, | |
160 | Ok(n) => { | |
161 | pos += n; | |
162 | if pos > chunk_size { break; } | |
163 | tmp = &mut tmp[n..]; | |
164 | } | |
165 | Err(ref e) if e.kind() == ErrorKind::Interrupted => { /* try again */ } | |
5f0c2d56 | 166 | Err(e) => bail!("read chunk failed - {}", e.to_string()), |
43eeef28 DM |
167 | } |
168 | } | |
43eeef28 DM |
169 | let mut start = 0; |
170 | while start + chunk_size <= pos { | |
171 | if !(chunk_cb)(file_pos, &buf[start..start+chunk_size])? { break; } | |
172 | file_pos += chunk_size; | |
173 | start += chunk_size; | |
174 | } | |
175 | if eof { | |
176 | if start < pos { | |
177 | (chunk_cb)(file_pos, &buf[start..pos])?; | |
178 | //file_pos += pos - start; | |
179 | } | |
180 | break; | |
181 | } else { | |
182 | let rest = pos - start; | |
183 | if rest > 0 { | |
184 | let ptr = buf.as_mut_ptr(); | |
185 | unsafe { std::ptr::copy_nonoverlapping(ptr.add(start), ptr, rest); } | |
186 | pos = rest; | |
187 | } else { | |
188 | pos = 0; | |
189 | } | |
190 | } | |
191 | } | |
192 | ||
193 | Ok(()) | |
43eeef28 | 194 | } |