From 3dd23fd3ba48131cbb991eb0913f51e4ec1e640f Mon Sep 17 00:00:00 2001
From: Thomas Lamprecht <t.lamprecht@proxmox.com>
Date: Mon, 2 Nov 2020 10:12:56 +0100
Subject: [PATCH] server: add CommandoSocket where multiple users can register
 commands

This is a preparatory step to replace the task control socket with it
and provide a "reopen log file" command for the rest server.

Kept it simple by disallowing to register new commands after the
socket gets spawned, this avoids the need for locking.

If we really need that we can always wrap it in a Arc<RWLock<..>> or
something like that, or even nicer, register at compile time.

Signed-off-by: Thomas Lamprecht <t.lamprecht@proxmox.com>
---
 src/server/command_socket.rs | 87 +++++++++++++++++++++++++++++++++---
 1 file changed, 80 insertions(+), 7 deletions(-)

diff --git a/src/server/command_socket.rs b/src/server/command_socket.rs
index fc17edf9..7ef6b88f 100644
--- a/src/server/command_socket.rs
+++ b/src/server/command_socket.rs
@@ -1,13 +1,13 @@
 use anyhow::{bail, format_err, Error};
 
-use futures::*;
-
-use tokio::net::UnixListener;
-
-use std::path::PathBuf;
-use serde_json::Value;
-use std::sync::Arc;
+use std::collections::HashMap;
 use std::os::unix::io::AsRawFd;
+use std::path::PathBuf;
+use std::sync::Arc;
+
+use futures::*;
+use tokio::net::UnixListener;
+use serde_json::Value;
 use nix::sys::socket;
 
 /// Listens on a Unix Socket to handle simple command asynchronously
@@ -140,3 +140,76 @@ pub async fn send_command<P>(
             }
         }).await
 }
+
+/// A callback for a specific commando socket.
+pub type CommandoSocketFn = Box<(dyn Fn(Option<&Value>) -> Result<Value, Error> + Send + Sync + 'static)>;
+
+/// Tooling to get a single control command socket where one can register multiple commands
+/// dynamically.
+/// You need to call `spawn()` to make the socket active.
+pub struct CommandoSocket {
+    socket: PathBuf,
+    commands: HashMap<String, CommandoSocketFn>,
+}
+
+impl CommandoSocket {
+    pub fn new<P>(path: P) -> Self
+        where P: Into<PathBuf>,
+    {
+        CommandoSocket {
+            socket: path.into(),
+            commands: HashMap::new(),
+        }
+    }
+
+    /// Spawn the socket and consume self, meaning you cannot register commands anymore after
+    /// calling this.
+    pub fn spawn(self) -> Result<(), Error> {
+        let control_future = create_control_socket(self.socket.to_owned(), move |param| {
+            let param = param
+                .as_object()
+                .ok_or_else(|| format_err!("unable to parse parameters (expected json object)"))?;
+
+            let command = match param.get("command") {
+                Some(Value::String(command)) => command.as_str(),
+                None => bail!("no command"),
+                _ => bail!("unable to parse command"),
+            };
+
+            if !self.commands.contains_key(command) {
+                bail!("got unknown command '{}'", command);
+            }
+
+            match self.commands.get(command) {
+                None => bail!("got unknown command '{}'", command),
+                Some(handler) => {
+                    let args = param.get("args"); //.unwrap_or(&Value::Null);
+                    (handler)(args)
+                },
+            }
+        })?;
+
+        tokio::spawn(control_future);
+
+        Ok(())
+    }
+
+    /// Register a new command with a callback.
+    pub fn register_command<F>(
+        &mut self,
+        command: String,
+        handler: F,
+    ) -> Result<(), Error>
+    where
+        F: Fn(Option<&Value>) -> Result<Value, Error> + Send + Sync + 'static,
+    {
+
+        if self.commands.contains_key(&command) {
+            bail!("command '{}' already exists!", command);
+        }
+
+        self.commands.insert(command, Box::new(handler));
+
+        Ok(())
+    }
+}