refactor enable_services

aspectron · Jul 22, 2024 · 0a83833 · 0a83833
1 parent 9461be4
commit 0a83833
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 25 deletions.
diff --git a/src/resolver.rs b/src/resolver.rs
@@ -156,7 +156,9 @@ pub fn install(ctx: &mut Context) -> Result<()> {
         return Ok(());
     }
 
-    log::remark("Installing Kaspa wPRC resolver...")?;
+    check_resolver_key(ctx)?;
+
+    log::remark("Installing 'kaspa-resolver'")?;
 
     let config = &ctx.config.resolver;
 
@@ -347,6 +349,10 @@ fn resolver_config_file(version: usize) -> String {
     format!("resolver.{version}.bin")
 }
 
+fn key_exists() -> bool {
+    resolver_config_folder().join(key_file()).exists()
+}
+
 fn load_key() -> Result<Secret> {
     Ok(Secret::from(fs::read(
         resolver_config_folder().join(key_file()),
@@ -392,6 +398,13 @@ pub fn update_resolver_config() {
     }
 }
 
+pub fn check_resolver_key(ctx: &mut Context) -> Result<()> {
+    if !key_exists() {
+        init_resolver_config(ctx)?;
+    }
+    Ok(())
+}
+
 pub fn init_resolver_config(ctx: &mut Context) -> Result<()> {
     if !ctx.config.resolver.enabled {
         log::warning("Resolver service is not enabled, please enable it.")?;

diff --git a/src/service.rs b/src/service.rs
@@ -1,6 +1,6 @@
 use crate::imports::*;
 
-#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize, Hash)]
+#[derive(Debug, Clone, Copy, Eq, PartialEq, Serialize, Deserialize, Hash)]
 pub enum ServiceKind {
     Kaspad(Network),
     Resolver,
@@ -69,41 +69,42 @@ pub trait Service {
 }
 
 pub fn enable_services(ctx: &mut Context, services: Vec<ServiceDetail>) -> Result<()> {
-    let mut kaspad_networks = Vec::new();
+    let kinds = services
+        .iter()
+        .map(|service| service.kind)
+        .collect::<Vec<_>>();
+    let networks = services
+        .iter()
+        .filter_map(|service| match service.kind {
+            ServiceKind::Kaspad(network) => Some(network),
+            _ => None,
+        })
+        .collect::<Vec<_>>();
 
-    let mut resolver_enabled = false;
-    for service in services {
-        match service.kind {
-            ServiceKind::Kaspad(network) => kaspad_networks.push(network),
-            ServiceKind::Resolver => {
-                resolver_enabled = true;
-                ctx.config.resolver.enabled = true;
-                ctx.config.save()?;
-
-                if !resolver::is_installed(ctx) {
-                    resolver::install(ctx)?;
-                }
+    if kinds.contains(&ServiceKind::Resolver) {
+        ctx.config.resolver.enabled = true;
+        ctx.config.save()?;
 
-                if !systemd::is_enabled(&ctx.config.resolver)? {
-                    systemd::daemon_reload()?;
-                    systemd::enable(&ctx.config.resolver)?;
-                    systemd::start(&ctx.config.resolver)?;
-                }
-            }
-            _ => {}
+        if !resolver::is_installed(ctx) {
+            resolver::install(ctx)?;
         }
-    }
 
-    if !resolver_enabled {
+        if !systemd::is_enabled(&ctx.config.resolver)? {
+            systemd::daemon_reload()?;
+            systemd::enable(&ctx.config.resolver)?;
+            systemd::start(&ctx.config.resolver)?;
+        }
+    } else {
         ctx.config.resolver.enabled = false;
         ctx.config.save()?;
+
         if systemd::is_enabled(&ctx.config.resolver)? {
             systemd::stop(&ctx.config.resolver)?;
             systemd::disable(&ctx.config.resolver)?;
         }
     }
 
-    kaspad::configure_networks(ctx, kaspad_networks)?;
+    kaspad::configure_networks(ctx, networks)?;
     nginx::reconfigure(ctx)?;
 
     Ok(())