From a2582fdce81a853b2a7fe9f201149bb23aa3609f Mon Sep 17 00:00:00 2001
From: realaravinth <realaravinth@batsense.net>
Date: Thu, 19 May 2022 16:24:18 +0530
Subject: [PATCH] feat: use constant format challenge TXT key

---
 src/verify.rs | 55 +++++++++++++++++++++++++++++++--------------------
 1 file changed, 34 insertions(+), 21 deletions(-)

diff --git a/src/verify.rs b/src/verify.rs
index e788d2f..d5838b9 100644
--- a/src/verify.rs
+++ b/src/verify.rs
@@ -19,12 +19,14 @@ use trust_dns_resolver::{
     config::{ResolverConfig, ResolverOpts},
     AsyncResolver,
 };
+use lazy_static::lazy_static;
 
 use crate::utils::get_random;
+use crate::ArcCtx;
 
 pub struct TXTChallenge {
     key: String,
-    base_hostname: String,
+    _base_hostname: String,
     value: String,
 }
 
@@ -32,54 +34,65 @@ const KEY_LEN: usize = 30;
 const VALUES_LEN: usize = 30;
 
 impl TXTChallenge {
-    pub async fn new(hostname: &str) -> Self {
-        let key = get_random(KEY_LEN);
+    pub fn get_challenge_txt_key_prefix(ctx: &ArcCtx) -> String {
+        // starchart-{{ starchart instance's hostname}}.{{ forge instance's hostname }}
+        format!("starchart-{}", &ctx.settings.server.domain)
+    }
+
+    pub fn get_challenge_txt_key(ctx: &ArcCtx, hostname: &str) -> String {
+        format!("{}.{}", Self::get_challenge_txt_key_prefix(ctx), hostname)
+    }
+
+
+
+    pub async fn new(ctx: &ArcCtx, hostname: &str) -> Self {
+        let key = Self::get_challenge_txt_key(ctx, hostname);
         let value = get_random(VALUES_LEN);
         Self {
             key,
             value,
-            base_hostname: hostname.to_string(),
+            _base_hostname: hostname.to_string(),
         }
     }
 
-    pub fn get_txt_key(&self) -> String {
-        format!("{}.{}", self.key, self.base_hostname)
-    }
-
-    pub async fn verify_txt(&self) -> bool {
+    pub async fn verify_txt(&self) -> Result<bool, Box<dyn std::error::Error>> {
         let conf = ResolverConfig::cloudflare_tls();
         let opts = ResolverOpts::default();
-        let resolver = AsyncResolver::tokio(conf, opts).unwrap();
-        let res = resolver.txt_lookup(&self.get_txt_key()).await.unwrap();
-        res.iter().any(|r| r.to_string() == self.value)
+        let resolver = AsyncResolver::tokio(conf, opts)?;
+        let res = resolver.txt_lookup(&self.key).await?;
+        Ok(res.iter().any(|r| r.to_string() == self.value))
     }
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::tests::sqlx_sqlite;
 
     #[actix_rt::test]
     async fn verify_txt_works() {
         // please note that this DNS record is in prod
-        const KEY: &str = "test-starchart-foobardontmindme";
-        const BASE_DOMAIN: &str = "starchart.forgeflux.org";
+        const BASE_DOMAIN: &str = "forge.forgeflux.org";
         const VALUE: &str = "ifthisvalueisretrievedbyforgefluxstarchartthenthetestshouldpass";
+
+        let (_db, ctx, _federate, _tmp_dir) = sqlx_sqlite::get_ctx().await;
+
+        let key =  TXTChallenge::get_challenge_txt_key(&ctx, BASE_DOMAIN);
         let mut txt_challenge = TXTChallenge {
             value: VALUE.to_string(),
-            base_hostname: BASE_DOMAIN.to_string(),
-            key: KEY.to_string(),
+            _base_hostname: BASE_DOMAIN.to_string(),
+            key: key.clone(),
         };
         assert_eq!(
-            txt_challenge.get_txt_key(),
-            "test-starchart-foobardontmindme.starchart.forgeflux.org"
+            TXTChallenge::get_challenge_txt_key(&ctx, BASE_DOMAIN),
+            key,
         );
 
         assert!(
-            txt_challenge.verify_txt().await,
+            txt_challenge.verify_txt().await.unwrap(),
             "TXT Challenge verification test"
         );
-        txt_challenge.value = KEY.into();
-        assert!(!txt_challenge.verify_txt().await);
+        txt_challenge.value = key;
+        assert!(!txt_challenge.verify_txt().await.unwrap());
     }
 }