charge only required cycles, validate http response more, and charge …

…cycles in two stages
zkemail · Dec 19, 2024 · 124304c · 124304c
1 parent db64f39
commit 124304c
Show file tree

Hide file tree

Showing 3 changed files with 65 additions and 10 deletions.
diff --git a/src/dns_client/src/lib.rs b/src/dns_client/src/lib.rs
@@ -1,7 +1,7 @@
 use std::collections::HashMap;
 
 use base64::{engine::general_purpose, Engine as _};
-use candid::Nat;
+use candid::{Nat, Principal};
 use ic_cdk::api::management_canister::http_request::{
     http_request, CanisterHttpRequestArgument, HttpHeader, HttpMethod, HttpResponse, TransformArgs,
     TransformContext,
@@ -44,7 +44,7 @@ pub async fn get_dkim_public_key(
             return Err("Insufficient cycles".to_string());
         }
     }
-    let accepted_cycles = ic_cdk::api::call::msg_cycles_accept128(available_cycles);
+    let accepted_cycles = ic_cdk::api::call::msg_cycles_accept128(CHARGED_CYCLE);
     if available_cycles != accepted_cycles {
         return Err("Fail to accept all available cycles".to_string());
     }
@@ -65,6 +65,10 @@ pub async fn get_dkim_public_key(
 
     // #[cfg(debug_assertions)]
     // let prefixes = vec!["https://dns.google/resolve"];
+    // let (seed_raw,): ([u8; 32],) = ic_cdk::call(Principal::management_canister(), "raw_rand", ())
+    //     .await
+    //     .expect("Failed to call the management canister");
+    // let seed = (seed_raw[0..16].iter().map(|&b| b as u128).sum::<u128>() % 3) as usize;
 
     let seed = ic_cdk::api::time() as usize % prefixes.len();
     let mut shuffled_prefixes = vec![];
@@ -154,7 +158,9 @@ fn _construct_request(prefix: &str, selector: &str, domain: &str) -> CanisterHtt
         value: "application/dns-json".to_string(),
     }];
 
-    let transform = TransformContext::from_name("transform".to_string(), vec![]);
+    let expected_name = format!("\"{}._domainkey.{}.\"", selector, domain);
+    let transform =
+        TransformContext::from_name("transform".to_string(), expected_name.as_bytes().to_vec());
     CanisterHttpRequestArgument {
         url: url.to_string(),
         method: HttpMethod::GET,
@@ -186,14 +192,31 @@ fn _transform(raw: TransformArgs) -> Result<HttpResponse, String> {
         .as_array()
         .ok_or_else(|| "No array of Answer")?
         .to_vec();
+    let expected_name = String::from_utf8(raw.context).expect("context is not a valid utf8 string");
     for i in 0..answers.len() {
+        if let Some(name) = answers[i].get("name") {
+            if name.to_string() != expected_name {
+                continue;
+            }
+        }
+        if let Some(dns_type) = answers[i].get("type") {
+            if dns_type.to_string() != "16" {
+                continue;
+            }
+        }
         let data = answers[i]["data"].to_string();
         if let Some(k_caps) = Regex::new("k=([a-z]+)").unwrap().captures(&data) {
             if &k_caps[1] != "rsa" {
                 continue;
             }
         }
 
+        if let Some(v_caps) = Regex::new("v=([A-Z0-9]+)").unwrap().captures(&data) {
+            if &v_caps[1] != "DKIM1" {
+                continue;
+            }
+        }
+
         if let Some(p_caps) = Regex::new(r#"p=([A-Za-z0-9\\+/" ]+);?"#)
             .unwrap()
             .captures(&data)

diff --git a/src/ic_dns_oracle_backend/src/lib.rs b/src/ic_dns_oracle_backend/src/lib.rs
@@ -11,11 +11,19 @@ use std::cell::RefCell;
 
 // consumed cycle for sign_dkim_public_key: 26_164_599_060 cycles
 // the consumed cycle * 1.5 is charged cycle = 39_246_898_590 cycles
-pub const SIGN_CHARGED_CYCLE: u128 = 39_246_898_590;
+pub const SIGN_CHARGED_CYCLE1: u128 = 39_246_898_590;
+
+// consumed cycle for sign_dkim_public_key: 26_164_599_060 cycles
+// the consumed cycle * 1.5 is charged cycle = 39_246_898_590 cycles
+pub const SIGN_CHARGED_CYCLE2: u128 = 39_246_898_590;
 
 // consumed cycle for revoke_dkim_public_key: 26_254_964_417 cycles
 // the consumed cycle * 1.5 is charged cycle = 39_382_446_626 cycles
-pub const REVOKE_CHARGED_CYCLE: u128 = 39_382_446_626;
+pub const REVOKE_CHARGED_CYCLE1: u128 = 39_382_446_626;
+
+// consumed cycle for revoke_dkim_public_key: 26_254_964_417 cycles
+// the consumed cycle * 1.5 is charged cycle = 39_382_446_626 cycles
+pub const REVOKE_CHARGED_CYCLE2: u128 = 39_382_446_626;
 
 // consumed cycle for get_dkim_public_key: 735_324_690 cycles
 // the consumed cycle * 1.5 is charged cycle = 1_102_987_035 cycles
@@ -259,12 +267,12 @@ pub async fn sign_dkim_public_key(
     let available_cycles = ic_cdk::api::call::msg_cycles_available128();
     #[cfg(not(debug_assertions))]
     {
-        if available_cycles < SIGN_CHARGED_CYCLE {
+        if available_cycles < SIGN_CHARGED_CYCLE1 {
             return Err("Insufficient cycles".to_string());
         }
     }
     // Accept all available cycles.
-    let accepted_cycles = ic_cdk::api::call::msg_cycles_accept128(available_cycles);
+    let accepted_cycles = ic_cdk::api::call::msg_cycles_accept128(SIGN_CHARGED_CYCLE1);
     if available_cycles != accepted_cycles {
         return Err("Fail to accept all available cycles".to_string());
     }
@@ -288,6 +296,18 @@ pub async fn sign_dkim_public_key(
         }
         Err(e) => e,
     };
+    let available_cycles = ic_cdk::api::call::msg_cycles_available128();
+    #[cfg(not(debug_assertions))]
+    {
+        if available_cycles < SIGN_CHARGED_CYCLE2 {
+            return Err("Insufficient cycles".to_string());
+        }
+    }
+    // Accept all available cycles.
+    let accepted_cycles = ic_cdk::api::call::msg_cycles_accept128(SIGN_CHARGED_CYCLE2);
+    if available_cycles != accepted_cycles {
+        return Err("Fail to accept all available cycles".to_string());
+    }
     let error1 = match _sign_dkim_public_key(selector, domain_with_gappssmtp, domain.clone()).await
     {
         Ok(res) => {
@@ -453,12 +473,12 @@ pub async fn revoke_dkim_public_key(
     let available_cycles = ic_cdk::api::call::msg_cycles_available128();
     #[cfg(not(debug_assertions))]
     {
-        if available_cycles < REVOKE_CHARGED_CYCLE {
+        if available_cycles < REVOKE_CHARGED_CYCLE1 {
             return Err("Insufficient cycles".to_string());
         }
     }
     // Accept all available cycles.
-    let accepted_cycles = ic_cdk::api::call::msg_cycles_accept128(available_cycles);
+    let accepted_cycles = ic_cdk::api::call::msg_cycles_accept128(REVOKE_CHARGED_CYCLE1);
     if available_cycles != accepted_cycles {
         return Err("Fail to accept all available cycles".to_string());
     }
@@ -486,6 +506,18 @@ pub async fn revoke_dkim_public_key(
         }
         Err(e) => e,
     };
+    let available_cycles = ic_cdk::api::call::msg_cycles_available128();
+    #[cfg(not(debug_assertions))]
+    {
+        if available_cycles < REVOKE_CHARGED_CYCLE2 {
+            return Err("Insufficient cycles".to_string());
+        }
+    }
+    // Accept all available cycles.
+    let accepted_cycles = ic_cdk::api::call::msg_cycles_accept128(REVOKE_CHARGED_CYCLE2);
+    if available_cycles != accepted_cycles {
+        return Err("Fail to accept all available cycles".to_string());
+    }
     let domain_with_gappssmtp =
         format!("{}.{}.gappssmtp.com", &domain.replace(".", "-"), &selector);
     let error1 = match _revoke_dkim_public_key(

diff --git a/src/poseidon/src/lib.rs b/src/poseidon/src/lib.rs
@@ -25,7 +25,7 @@ pub fn public_key_hash(public_key_hex: String) -> Result<String, String> {
         }
     }
     // Accept all available cycles.
-    let accepted_cycles = ic_cdk::api::call::msg_cycles_accept128(available_cycles);
+    let accepted_cycles = ic_cdk::api::call::msg_cycles_accept128(CHARGED_CYCLE);
     if available_cycles != accepted_cycles {
         return Err("Fail to accept all available cycles".to_string());
     }