Skip to main content

linux_info/
util.rs

1use std::convert::TryInto;
2use std::ffi::CString;
3use std::fs::File;
4use std::io::Read;
5use std::os::unix::ffi::OsStrExt;
6use std::os::unix::prelude::AsRawFd;
7use std::path::Path;
8use std::{fmt, io, mem};
9
10use byte_parser::{ParseIterator, StrParser};
11
12use libc::c_int;
13
14const DEF_PRECISION: usize = 2;
15
16/// Represents a size, for example `1024 kB`.
17#[derive(Debug, Clone, PartialEq, Eq)]
18pub struct DataSize {
19	bytes: u128,
20}
21
22impl DataSize {
23	// not implemeting FromStr because this is private.
24	pub(crate) fn from_str(s: &str) -> Option<Self> {
25		let mut iter = StrParser::new(s);
26		let float = parse_f64(&mut iter)?;
27		// now we need to parse the unit
28		let unit = iter.record().consume_to_str().trim();
29
30		let unit = DataSizeUnit::from_str(unit)?;
31		Some(Self {
32			bytes: unit.to_byte(float),
33		})
34	}
35
36	pub(crate) fn from_size_bytes(bytes: impl TryInto<u128>) -> Option<Self> {
37		bytes.try_into().ok().map(|bytes| Self { bytes })
38	}
39
40	/// Convert the data unit into a specific unit.
41	pub fn to(self, unit: &DataSizeUnit) -> f64 {
42		DataSizeUnit::convert(self.bytes, unit)
43	}
44}
45
46#[derive(Debug, Clone, Copy, PartialEq, Eq)]
47pub enum DataSizeUnit {
48	/// Byte
49	B,
50	/// Kilobyte
51	Kb, // 1_000
52	/// Megabyte
53	Mb, // 1_000_000
54	/// Gigabyte
55	Gb, // 1_000_000_000
56	/// Terabyte
57	Tb, // 1_000_000_000_000
58}
59
60impl DataSizeUnit {
61	const fn val(&self) -> u128 {
62		match self {
63			Self::B => 1,
64			Self::Kb => 1_024,
65			Self::Mb => 1_024 * 1_024,
66			Self::Gb => 1_024 * 1_024 * 1_024,
67			Self::Tb => 1_024 * 1_024 * 1_024 * 1_024,
68		}
69	}
70
71	fn from_str(s: &str) -> Option<Self> {
72		Some(match s {
73			"" => Self::B,
74			s if eqs(s, "b") => Self::B,
75			s if eqs(s, "kb") => Self::Kb,
76			s if eqs(s, "mb") => Self::Mb,
77			s if eqs(s, "gb") => Self::Gb,
78			s if eqs(s, "tb") => Self::Tb,
79			_ => return None,
80		})
81	}
82
83	fn to_byte(&self, val: f64) -> u128 {
84		// TODO probably need fix this overflowing
85		(val * self.val() as f64) as u128
86	}
87
88	fn adjust_to(byte: u128) -> Self {
89		match byte {
90			b if b < Self::Kb.val() => Self::B,
91			b if b < Self::Mb.val() => Self::Kb,
92			b if b < Self::Gb.val() => Self::Mb,
93			b if b < Self::Tb.val() => Self::Gb,
94			_ => Self::Tb,
95		}
96	}
97
98	fn convert(byte: u128, to: &Self) -> f64 {
99		byte as f64 / to.val() as f64
100	}
101
102	const fn as_str(&self) -> &'static str {
103		match self {
104			Self::B => "b",
105			Self::Kb => "kb",
106			Self::Mb => "mb",
107			Self::Gb => "gb",
108			Self::Tb => "tb",
109		}
110	}
111
112	// val needs to be already adjusted to self
113	fn fmt_val(&self, val: f64, f: &mut fmt::Formatter<'_>) -> fmt::Result {
114		if val == 0f64 {
115			return write!(f, "{}", val);
116		}
117
118		// calculate the precision we wan't to use
119		let fract = val.fract();
120		let precision = if fract == 0f64 {
121			0
122		} else {
123			let max = f.precision().unwrap_or(DEF_PRECISION);
124
125			if max <= 1 {
126				max
127			} else {
128				calculate_precision(fract, max)
129			}
130		};
131
132		write!(f, "{:.*} {}", precision, val, self.as_str())
133	}
134}
135
136fn calculate_precision(mut fract: f64, max: usize) -> usize {
137	let max_number = 10usize.pow(max as u32) as f64;
138
139	// round to the max position
140	fract *= max_number;
141	fract = fract.round();
142
143	// then go slowly back until we hit a fraction
144	for m in (1..=max).rev() {
145		fract /= 10f64;
146		if fract.fract() != 0f64 {
147			return m;
148		}
149	}
150
151	0
152}
153
154#[inline(always)]
155fn eqs(a: &str, b: &str) -> bool {
156	a.eq_ignore_ascii_case(b)
157}
158
159impl fmt::Display for DataSize {
160	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
161		let unit = DataSizeUnit::adjust_to(self.bytes);
162		let val = DataSizeUnit::convert(self.bytes, &unit);
163		unit.fmt_val(val, f)
164	}
165}
166
167// parses a part of a slice
168// Panics if Iterator contains not valid utf8
169fn parse_f64<'s, I>(iter: &mut I) -> Option<f64>
170where
171	I: ParseIterator<'s>,
172{
173	let mut iter = iter.record();
174
175	// consume first digits
176	iter.while_byte_fn(u8::is_ascii_digit)
177		.consume_at_least(1)
178		.ok()?;
179
180	// dot
181	let has_dot = iter.next_if(|&b| b == b'.').is_some();
182
183	if has_dot {
184		// consume next digits
185		iter.consume_while_byte_fn(u8::is_ascii_digit);
186	}
187
188	iter.to_str().parse().ok()
189}
190
191/// Clears the string the writes the entire file to the string.  
192/// Does not allocate in advance like std::fs::read_to_string.
193pub fn read_to_string_mut(
194	path: impl AsRef<Path>,
195	s: &mut String,
196) -> io::Result<()> {
197	s.clear();
198	let mut file = File::open(path)?;
199	file.read_to_string(s).map(|_| ())
200}
201
202fn cstr(path: impl AsRef<Path>) -> io::Result<CString> {
203	CString::new(path.as_ref().as_os_str().as_bytes()).map_err(From::from)
204}
205
206// see https://man7.org/linux/man-pages/man2/fstatfs.2.html
207pub fn statfs(path: impl AsRef<Path>) -> io::Result<libc::statfs> {
208	unsafe {
209		let mut stat = mem::MaybeUninit::<libc::statfs>::uninit();
210		let c = cstr(path)?;
211		let r = libc::statfs(c.as_ptr(), stat.as_mut_ptr());
212		match r {
213			0 => Ok(stat.assume_init()),
214			-1 => Err(io::Error::last_os_error()),
215			_ => panic!("unexpected return value from statfs {:?}", r),
216		}
217	}
218}
219
220// BLKSSZGET
221
222pub fn blkdev_sector_size(fd: impl AsRawFd) -> io::Result<u64> {
223	let s = unsafe {
224		let mut size: c_int = 0; // todo set default to 512
225		match blksszget(fd.as_raw_fd(), &mut size) {
226			-1 => return Err(io::Error::last_os_error()),
227			_ => size,
228		}
229	};
230
231	s.try_into()
232		.map_err(|e| io::Error::new(io::ErrorKind::Other, e))
233}
234
235#[cfg(any(
236	target_arch = "x86",
237	target_arch = "arm",
238	target_arch = "x86_64",
239	target_arch = "aarch64"
240))]
241unsafe fn blksszget(fd: c_int, data: *mut c_int) -> c_int {
242	let nr = (0x12 << 8) | (104 << 0);
243	libc::ioctl(fd, nr, data)
244}
245
246#[cfg(test)]
247mod tests {
248	use super::*;
249
250	#[test]
251	fn test_size() {
252		let size = DataSize::from_str("24576 kB").unwrap();
253		assert_eq!(size.to(&DataSizeUnit::Kb), 24576.0);
254	}
255
256	#[test]
257	fn size_str() {
258		// TODO update the formatter
259		let s = DataSize::from_str("1024").unwrap();
260		assert_eq!(s.to_string(), "1 kb");
261		let s = DataSize::from_str("10 kb").unwrap();
262		assert_eq!(s.to_string(), "10 kb");
263		let s = DataSize::from_str("42.1 mB").unwrap();
264		assert_eq!(s.to_string(), "42.1 mb");
265		let s = DataSize::from_str("4.22 Gb").unwrap();
266		assert_eq!(s.to_string(), "4.22 gb");
267		let s = DataSize::from_str("2000 Tb").unwrap();
268		assert_eq!(s.to_string(), "2000 tb");
269		// and precision
270		assert_eq!(
271			format!("{:.0}", DataSize::from_str("1.2 kb").unwrap()),
272			"1 kb"
273		);
274	}
275
276	#[test]
277	fn test_precision() {
278		assert_eq!(calculate_precision(0.00005, 4), 4);
279		assert_eq!(calculate_precision(0.00001, 4), 0);
280		assert_eq!(calculate_precision(0.0001, 4), 4);
281		assert_eq!(calculate_precision(0.001, 4), 3);
282		assert_eq!(calculate_precision(0.01, 4), 2);
283		assert_eq!(calculate_precision(0.1, 4), 1);
284		assert_eq!(calculate_precision(0.0, 4), 0);
285	}
286
287	#[test]
288	fn run_statfs() {
289		statfs("/").unwrap();
290	}
291}