🐛 fix byte-unit (mokeyish#185)

FlyLoongZ · Nov 29, 2023 · dce4ab2 · dce4ab2
1 parent df08ad1
commit dce4ab2
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 14 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
@@ -382,8 +382,8 @@ impl RuntimeConfig {
         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())
+            .unwrap_or_else(|| Byte::from_u64_with_unit(128, Unit::KB).unwrap())
+            .as_u64()
     }
     #[inline]
     pub fn log_num(&self) -> u64 {
@@ -425,8 +425,8 @@ impl RuntimeConfig {
         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())
+            .unwrap_or_else(|| Byte::from_u64_with_unit(128, Unit::KB).unwrap())
+            .as_u64()
     }
 
     #[inline]
@@ -1336,20 +1336,30 @@ mod parse {
             );
         }
 
+        #[test]
+        fn test_default_audit_size_1() {
+            use byte_unit::Unit;
+            let cfg = RuntimeConfig::builder().build();
+            assert_eq!(
+                cfg.audit_size(),
+                Byte::from_i64_with_unit(128, Unit::KB).unwrap().as_u64()
+            );
+        }
+
         #[test]
         fn test_parse_config_audit_size_1() {
             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]