Improves odohconfig parsing

cloudflare · Dec 17, 2020 · f5e2948 · f5e2948
1 parent 58bb34d
commit f5e2948
Showing 1 changed file with 21 additions and 1 deletion.
diff --git a/src/dns_utils.rs b/src/dns_utils.rs
@@ -92,7 +92,7 @@ pub async fn fetch_odoh_config(target: &str) -> Result<Vec<u8>> {
 /// parse keys.
 fn odohconfig_from_https(data: &[u8]) -> Result<Vec<u8>> {
     let data_hex = hex::encode(data);
-    let vec = data_hex.split(ODOH_VERSION).collect::<Vec<_>>();
+    let vec = data_hex.splitn(2, ODOH_VERSION).collect::<Vec<_>>();
     let first_len = vec.first().unwrap().len();
     let odohconfig_len = vec.first().unwrap().to_string().split_off(first_len - 4);
     let odohconfig = format!(
@@ -136,3 +136,23 @@ fn get_qtype(qtype: &str) -> Result<RecordType> {
     };
     Ok(rtype)
 }
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[test]
+    fn test_odohconfig_parser() {
+        // if odohconfig string does not contain multiple ODOH_VERSION 
+        let mut data = "00010000010003026832000400086810f8f96810f9f9000600202606470000000000000000006810f8f92606470000000000000000006810f9f98001002e002cff0300280020000100010020f150567d5bfd7ffbb20f52b73cce923f0654e37265d7065dc5d6bfc8912b3e5e";
+        let mut odohconfig = odohconfig_from_https(&hex::decode(data).unwrap()).unwrap();
+        let mut expected_odohconfig = "002cff0300280020000100010020f150567d5bfd7ffbb20f52b73cce923f0654e37265d7065dc5d6bfc8912b3e5e";
+        assert_eq!(expected_odohconfig, hex::encode(odohconfig));
+
+        // if odohconfig string contains multiple ODOH_VERSION
+        data = "00010000010003026832000400086810f8f96810f9f9000600202606470000000000000000006810f8f92606470000000000000000006810f9f98001002e002cff0300280020000100010020128dfb25d235708df7031ff036bd82ec061c0877835186321e0c03e24f93213a";
+        odohconfig = odohconfig_from_https(&hex::decode(data).unwrap()).unwrap();
+        expected_odohconfig = "002cff0300280020000100010020128dfb25d235708df7031ff036bd82ec061c0877835186321e0c03e24f93213a";
+        assert_eq!(expected_odohconfig, hex::encode(odohconfig));
+    }
+}