🐛 fix byte-unit #184

mokeyish · Nov 29, 2023 · f19b323 · f19b323
1 parent fb2ae8c
commit f19b323
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 20 deletions.
diff --git a/src/config/parser/bytes.rs b/src/config/parser/bytes.rs
@@ -4,10 +4,9 @@ impl NomParser for Byte {
     fn parse(input: &str) -> IResult<&str, Self> {
         let num = recognize(pair(digit1, opt(pair(char('.'), digit1))));
         let unit = alpha1;
-        map_res(
-            recognize(tuple((num, space0, unit))),
-            <Byte as std::str::FromStr>::from_str,
-        )(input)
+        map_res(recognize(tuple((num, space0, unit))), |s| {
+            Byte::parse_str(s, true)
+        })(input)
     }
 }
 
@@ -20,19 +19,23 @@ mod tests {
         use byte_unit::Unit;
         assert_eq!(
             Byte::parse("12kb"),
-            Ok(("", Byte::from_i64_with_unit(12, Unit::Kbit).unwrap()))
+            Ok(("", Byte::from_u64_with_unit(12, Unit::KB).unwrap()))
         );
         assert_eq!(
             Byte::parse("123mb"),
-            Ok(("", Byte::from_i64_with_unit(123, Unit::Mbit).unwrap()))
+            Ok(("", Byte::from_u64_with_unit(123, Unit::MB).unwrap()))
         );
         assert_eq!(
             Byte::parse("80mb"),
-            Ok(("", Byte::from_i64_with_unit(80, Unit::Mbit).unwrap()))
+            Ok(("", Byte::from_u64_with_unit(80, Unit::MB).unwrap()))
         );
         assert_eq!(
             Byte::parse("30 gb"),
-            Ok(("", Byte::from_i64_with_unit(30, Unit::Gbit).unwrap()))
+            Ok(("", Byte::from_i64_with_unit(30, Unit::GB).unwrap()))
+        );
+        assert_eq!(
+            Byte::parse("30GB"),
+            Ok(("", Byte::from_i64_with_unit(30, Unit::GB).unwrap()))
         );
     }
 }
diff --git a/src/dns_conf.rs b/src/dns_conf.rs
@@ -379,11 +379,7 @@ impl RuntimeConfig {
 
     #[inline]
     pub fn log_size(&self) -> u64 {
-        use byte_unit::{Byte, Unit};
-        self.log
-            .size
-            .map(|n| n.as_u64())
-            .unwrap_or(Byte::from_i64_with_unit(128, Unit::KB).unwrap().as_u64())
+        self.log.size.map(|n| n.as_u64()).unwrap_or(128)
     }
     #[inline]
     pub fn log_num(&self) -> u64 {
@@ -422,11 +418,7 @@ impl RuntimeConfig {
 
     #[inline]
     pub fn audit_size(&self) -> u64 {
-        use byte_unit::{Byte, Unit};
-        self.audit
-            .size
-            .map(|n| n.as_u64())
-            .unwrap_or(Byte::from_i64_with_unit(128, Unit::KB).unwrap().as_u64())
+        self.audit.size.map(|n| n.as_u64()).unwrap_or(128)
     }
 
     #[inline]
@@ -1341,15 +1333,15 @@ mod parse {
             use byte_unit::Unit;
             let mut cfg = RuntimeConfig::builder();
             cfg.config("audit-size 80mb");
-            assert_eq!(cfg.audit.size, Byte::from_i64_with_unit(80, Unit::Mbit));
+            assert_eq!(cfg.audit.size, Byte::from_i64_with_unit(80, Unit::MB));
         }
 
         #[test]
         fn test_parse_config_audit_size_2() {
             use byte_unit::Unit;
             let mut cfg = RuntimeConfig::builder();
             cfg.config("audit-size 30 gb");
-            assert_eq!(cfg.audit.size, Byte::from_i64_with_unit(30, Unit::Gbit));
+            assert_eq!(cfg.audit.size, Byte::from_i64_with_unit(30, Unit::GB));
         }
 
         #[test]