diff --git a/src/dns_client/src/lib.rs b/src/dns_client/src/lib.rs index 9d131d6..6e68107 100644 --- 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::() % 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,7 +192,18 @@ fn _transform(raw: TransformArgs) -> Result { .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" { @@ -194,6 +211,12 @@ fn _transform(raw: TransformArgs) -> Result { } } + 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 index d579f92..bf45433 100644 --- 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 index e871e43..868012c 100644 --- 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 { } } // 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()); }