From 8a1e9ccb4973577f67f5deabf6250ea79d6e802c Mon Sep 17 00:00:00 2001
From: Daniel Knopik <daniel@dknopik.de>
Date: Tue, 22 Oct 2024 20:21:20 +0200
Subject: [PATCH 1/2] config: better way to specify custom hosts

---
 lib/src/config/shadow.rs | 11 +----------
 lib/src/error.rs         |  4 ++++
 lib/src/lib.rs           | 24 +++++++++++++++---------
 3 files changed, 20 insertions(+), 19 deletions(-)

diff --git a/lib/src/config/shadow.rs b/lib/src/config/shadow.rs
index e76bbcf..65d77f1 100644
--- a/lib/src/config/shadow.rs
+++ b/lib/src/config/shadow.rs
@@ -129,7 +129,7 @@ pub struct HostsMut<'a> {
 }
 
 impl<'a> Iterator for HostsMut<'a> {
-    type Item = Result<UntypedHost<&'a mut Mapping>, Error>;
+    type Item = Result<&'a mut Mapping, Error>;
 
     fn next(&mut self) -> Option<Self::Item> {
         self.hosts
@@ -138,15 +138,6 @@ impl<'a> Iterator for HostsMut<'a> {
             .map(|(_, host)| {
                 host.as_mapping_mut()
                     .ok_or_else(|| Error::ExpectedOtherType("host".to_string()))
-                    .map(UntypedHost)
             })
     }
 }
-
-pub struct UntypedHost<T>(T);
-
-impl<'a> UntypedHost<&'a mut Mapping> {
-    pub fn network_node_id_mut(&mut self) -> Option<&mut Value> {
-        self.0.get_mut("network_node_id")
-    }
-}
diff --git a/lib/src/error.rs b/lib/src/error.rs
index 1acf1e8..0562ff9 100644
--- a/lib/src/error.rs
+++ b/lib/src/error.rs
@@ -51,4 +51,8 @@ pub enum Error {
             validator client"
     )]
     MissingValidatorCount,
+    #[error(
+        "You must specify a reliability and location, and no network_node_id in your custom hosts"
+    )]
+    InvalidShadowHost,
 }
diff --git a/lib/src/lib.rs b/lib/src/lib.rs
index d12b911..77365e8 100644
--- a/lib/src/lib.rs
+++ b/lib/src/lib.rs
@@ -4,7 +4,6 @@ use crate::network_graph::{generate_network_graph, GeneratedNetworkGraph};
 use crate::node::NodeManager;
 use crate::validators::ValidatorManager;
 use log::{debug, info};
-use serde_yaml::Value;
 use std::borrow::Cow;
 use std::ffi::{OsStr, OsString};
 use std::fs::{create_dir, File};
@@ -110,15 +109,22 @@ pub fn generate<T: TryInto<FullConfig, Error = Error>>(
 
     // postprocessing given shadow config values: overwrite string network ids
     for host in shadow_config.hosts_mut()? {
-        if let Some(node_id) = host?.network_node_id_mut() {
-            if let Some((location, reliability)) = node_id.as_str().and_then(|s| s.split_once('-'))
-            {
-                let node = network_graph.assign_network_node(location, reliability)?;
-                *node_id = Value::Number(node.id().into());
-            } else {
-                return Err(Error::ExpectedOtherType("network_node_id".to_string()));
-            }
+        let mapping = host?;
+        if mapping.get("network_node_id").is_some() {
+            return Err(Error::InvalidShadowHost);
         }
+        let location = mapping
+            .get("location")
+            .ok_or(Error::InvalidShadowHost)?
+            .as_str()
+            .ok_or_else(|| Error::ExpectedOtherType("location".to_string()))?;
+        let reliability = mapping
+            .get("reliability")
+            .ok_or(Error::InvalidShadowHost)?
+            .as_str()
+            .ok_or_else(|| Error::ExpectedOtherType("reliability".to_string()))?;
+        let node = network_graph.assign_network_node(location, reliability)?;
+        mapping.insert("network_node_id".into(), node.id().into());
     }
 
     info!("Generating nodes");

From 7a6f95ee47d745aa11d035d63b979f9a0119e034 Mon Sep 17 00:00:00 2001
From: Daniel Knopik <daniel@dknopik.de>
Date: Tue, 22 Oct 2024 20:34:32 +0200
Subject: [PATCH 2/2] remove reliability and location from custom shadow hosts
 during conversion

---
 lib/src/lib.rs | 21 +++++++++++----------
 1 file changed, 11 insertions(+), 10 deletions(-)

diff --git a/lib/src/lib.rs b/lib/src/lib.rs
index 77365e8..3a92728 100644
--- a/lib/src/lib.rs
+++ b/lib/src/lib.rs
@@ -4,6 +4,7 @@ use crate::network_graph::{generate_network_graph, GeneratedNetworkGraph};
 use crate::node::NodeManager;
 use crate::validators::ValidatorManager;
 use log::{debug, info};
+use serde_yaml::Value;
 use std::borrow::Cow;
 use std::ffi::{OsStr, OsString};
 use std::fs::{create_dir, File};
@@ -113,17 +114,17 @@ pub fn generate<T: TryInto<FullConfig, Error = Error>>(
         if mapping.get("network_node_id").is_some() {
             return Err(Error::InvalidShadowHost);
         }
-        let location = mapping
-            .get("location")
+        let Value::String(location) = mapping.remove("location").ok_or(Error::InvalidShadowHost)?
+        else {
+            return Err(Error::ExpectedOtherType("location".to_string()));
+        };
+        let Value::String(reliability) = mapping
+            .remove("reliability")
             .ok_or(Error::InvalidShadowHost)?
-            .as_str()
-            .ok_or_else(|| Error::ExpectedOtherType("location".to_string()))?;
-        let reliability = mapping
-            .get("reliability")
-            .ok_or(Error::InvalidShadowHost)?
-            .as_str()
-            .ok_or_else(|| Error::ExpectedOtherType("reliability".to_string()))?;
-        let node = network_graph.assign_network_node(location, reliability)?;
+        else {
+            return Err(Error::ExpectedOtherType("reliability".to_string()));
+        };
+        let node = network_graph.assign_network_node(&location, &reliability)?;
         mapping.insert("network_node_id".into(), node.id().into());
     }