From e901f63bb9627ba6ea5a11ae2cbe5cff8e99713b Mon Sep 17 00:00:00 2001
From: core <core@coredoes.dev>
Date: Thu, 17 Aug 2023 11:22:12 -0400
Subject: [PATCH] keystore v2 work

---
 trifid-api/src/routes/v1/dnclient.rs | 88 +++++++++++++++-------------
 1 file changed, 47 insertions(+), 41 deletions(-)

diff --git a/trifid-api/src/routes/v1/dnclient.rs b/trifid-api/src/routes/v1/dnclient.rs
index a6d2e80..75ce2ba 100644
--- a/trifid-api/src/routes/v1/dnclient.rs
+++ b/trifid-api/src/routes/v1/dnclient.rs
@@ -17,10 +17,11 @@ use std::time::{SystemTime, UNIX_EPOCH};
 use sea_orm::{ActiveModelTrait, EntityTrait};
 use trifid_pki::cert::{deserialize_ed25519_public, deserialize_x25519_public};
 use trifid_pki::x25519_dalek::PublicKey;
-use trifid_api_entities::entity::host;
+use trifid_api_entities::entity::{host, keystore_entry, keystore_host};
 use crate::error::APIErrorsResponse;
 use sea_orm::{ColumnTrait, QueryFilter, IntoActiveModel};
 use sea_orm::ActiveValue::Set;
+use trifid_api_entities::entity::prelude::KeystoreHost;
 
 #[post("/v1/dnclient")]
 pub async fn dnclient(
@@ -54,8 +55,18 @@ pub async fn dnclient(
         }
     };
 
-    let host_in_ks = keystore.hosts.iter_mut().find(|u| &u.id == host);
-    let host_in_ks = match host_in_ks {
+    let host_in_ks = match keystore_host::Entity::find().filter(keystore_host::Column::Id.eq(host)).one(&db.conn).await {
+        Ok(maybe_host) => maybe_host,
+        Err(e) => {
+            return HttpResponse::InternalServerError().json(vec![APIError {
+                code: "ERR_DB_ERROR".to_string(),
+                message: "There was an error finding the keys for your host, please contact your administrator".to_string(),
+                path: None
+            }])
+        }
+    };
+
+    let keystore_header = match host_in_ks {
         Some(host) => host,
         None => {
             return HttpResponse::Unauthorized().json(vec![APIError {
@@ -63,18 +74,36 @@ pub async fn dnclient(
                 message: "The host does not exist or you do not have permission to access it."
                     .to_string(),
                 path: None,
+            }]);
+        }
+    };
+
+    let counter = keystore_header.counter;
+
+    // pull their key information
+
+    let key_info = match keystore_entry::Entity::find().filter(keystore_entry::Column::Host.eq(&keystore_header.id)).one(&db.conn).await {
+        Ok(maybe_keys) => maybe_keys,
+        Err(e) => {
+            return HttpResponse::InternalServerError().json(vec![APIError {
+                code: "ERR_DB_ERROR".to_string(),
+                message: "There was an error finding the keys for your host, please contact your administrator".to_string(),
+                path: None
             }])
         }
     };
 
-    let client_keys = host_in_ks
-        .client_keys
-        .iter()
-        .find(|u| u.id == req.counter as u64);
-    let client_keys_2 = host_in_ks
-        .client_keys
-        .iter()
-        .find(|u| u.id == host_in_ks.current_client_key);
+    let keystore_data = match key_info {
+        Some(key_info) => key_info,
+        None => {
+            return HttpResponse::Unauthorized().json(vec![APIError {
+                code: "ERR_HOST_ERROR".to_string(),
+                message: "The host does not exist or you do not have permission to access it."
+                    .to_string(),
+                path: None,
+            }]);
+        }
+    };
 
     let signature = match Signature::from_slice(&req.signature) {
         Ok(sig) => sig,
@@ -90,27 +119,9 @@ pub async fn dnclient(
         }
     };
 
-    let mut valid = false;
-    let mut valid_key = None;
+    let key = VerifyingKey::from(keystore_data.client_signing_key.try_into().unwrap());
 
-    if let Some(client_keys) = client_keys {
-        if client_keys.ed_pub.verify(req.message.as_bytes(), &signature).is_ok() {
-            valid = true;
-            valid_key = Some(client_keys);
-        }
-    }
-    if let Some(client_keys_2) = client_keys_2 {
-        if client_keys_2.ed_pub.verify(req.message.as_bytes(), &signature).is_ok() {
-            valid = true;
-            valid_key = Some(client_keys_2);
-        }
-    }
-
-    if client_keys.is_none() && client_keys_2.is_none() {
-        panic!("No valid keys for host {}", host);
-    }
-
-    if !valid {
+    if !key.verify(&req.message.as_bytes(), &signature).is_ok() {
         // Be intentionally vague as the message is invalid.
         warn!("! invalid signature from {}", host);
         return HttpResponse::Unauthorized().json(vec![APIError {
@@ -121,8 +132,6 @@ pub async fn dnclient(
         }]);
     }
 
-    let client_keys = valid_key.unwrap();
-
     // Sig OK
     // Decode the message from base64
 
@@ -190,7 +199,7 @@ pub async fn dnclient(
 
     // Do a config build
 
-    let info = match collect_info(&db, host, client_keys.dh_pub.as_bytes()).await {
+    let info = match collect_info(&db, host, &keystore_data.client_dh_key).await {
         Ok(i) => i,
         Err(e) => {
             return HttpResponse::InternalServerError().json(EnrollResponse::Error {
@@ -218,13 +227,10 @@ pub async fn dnclient(
         }
     };
 
-    let current_cfg = host_in_ks
-        .config
-        .iter()
-        .find(|u| u.id == host_in_ks.current_config);
+    let current_cfg = keystore_data.config;
+    let generated_config_string = serde_yaml::to_string(&cfg).unwrap();
 
-    let config_update_avail = current_cfg.map(|u| u.config.clone()) != Some(cfg.clone())
-        || req.counter < host_in_ks.current_config as u32;
+    let config_update_avail = current_cfg != generated_config_string || req.counter < keystore_header.counter as u32;
 
     host_am.last_out_of_date = Set(config_update_avail);
     host_am.last_seen_at = Set(SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_secs() as i64);
@@ -315,7 +321,7 @@ pub async fn dnclient(
                 }
             };
 
-            let ks = host_in_ks;
+            let ks = keystore_header;
 
             ks.certs.push(KSCert {
                 id: ks.current_cert + 1,