add: implement basic server

Cargo.toml

@@ -5,9 +5,11 @@ edition = "2024"
 
 [dependencies]
 anyhow = "1.0.97"
+chrono = "0.4.40"
 clap = { version = "4.5.34", features = ["derive"] }
 dir-diff = "0.3.3"
 env_logger = "0.11.7"
 include_directory = "0.1.1"
 log = "0.4.27"
+rouille = "3.6.2"
 tempdir = "0.3.7"

README.md

@@ -12,8 +12,30 @@ tier N1 VM. That extra hop doesn't matter.
 ## Config
 
 Configure nginx to do what you want, test it. Use any Node IP for your testing.
-This tool accepts an argument (rewrite_string) which will be replaced using these rules:
+This will become the 'template_dir' in the argument to the LB.
 
-1. For each line that contains rewrite string:
+Move that directory to somewhere new, i.e. `/etc/nginx-template/`. Make
-2. Copy the line once per node IP, replacing the string with host IPs
+a symlink from that new directory to the old one (i.e.,
+`ln -s /etc/nginx-template /etc/nginx/`).
 
+Make a workspace directory for this tool; it will write configs to this folder
+before updating the symlink you created above. It needs to be persistent so on
+server reboot the service starts ok (i.e., `mkdir /var/skubelb/`).
+
+Make sure the user running the tool has read access to the template folder, read-write
+access to the workspace folder and config symlink.
+
+Run the server with a command like:
+
+```sh
+skubelb --needle some_node_ip \
+    --workspace_dir /var/skubelb \
+    --config_symlink /etc/nginx \
+    --template_dir /etc/nginx-template
+    --listen 0.0.0.0:8080
+```
+
+Replacing `some_node_ip` with the node IP you used during the initial setup.
+
+Next, configure the Kubernetes nodes to POST `http://loadbalancer:8080/register` when
+they started, and DELETE `http://loadbalancer:8080/register` when they shutdown.

src/lib.rs

@@ -0,0 +1,5 @@
+mod rewriter;
+mod server;
+
+pub use rewriter::*;
+pub use server::*;

src/main.rs

@@ -1,10 +1,14 @@
+use std::sync::Mutex;
+
 use clap::Parser;
 
-mod rewriter;
+use skubelb::Rewriter;
+use skubelb::Server;
 
 use env_logger::Env;
-use log::info;
+use log::{info, warn};
-use rewriter::Rewriter;
+use anyhow::Result;
+use rouille::{router, Request, Response};
 
 /// Implements a HTTP server which allows clients to 'register'
 /// themselves. Their IP address will be used to replace a
@@ -29,11 +33,20 @@ struct Args {
     /// The folder which contains the templates that
     /// will be be searched for the needle.
     #[arg(short, long)]
-    source_folder: String,
+    template_dir: String,
 
-    /// Where to write the replaced lines.
+    /// The symlink that should be updated each time the config changes.
+    /// 
+    /// Symlinks are used because file updates are not atomic.
     #[arg(short, long)]
-    dest_folder: String,
+    config_symlink: String,
+
+    /// Where to actually store the generated configs.
+    #[arg(short, long)]
+    workspace_dir: String,
+
+    #[arg(short, long, default_value_t = String::from("0.0.0.0:8080"))]
+    listen: String,
 }
 
 fn main() {
@@ -41,11 +54,31 @@ fn main() {
     env_logger::Builder::from_env(Env::default().default_filter_or("info")).init();
     let args = Args::parse();
 
-    let mut rewriter = Rewriter::new(args.rewrite_string);
+    let rewriter = Rewriter::new(args.rewrite_string);
+    let server_impl = Mutex::new(Server::new(rewriter, args.workspace_dir, args.template_dir, args.config_symlink));
 
-    rewriter.add_replacement(String::from("abc"));
+    rouille::start_server(args.listen, move |request| {
-    rewriter
+        info!("Processing request: {:?}", request);
-        .rewrite_folder(&args.source_folder, &args.dest_folder)
+        match handle(request, &server_impl) {
-        .unwrap();
+            Ok(resp) => resp,
-    info!("Finished writing new config to {}", args.dest_folder);
+            Err(e) => {
+                warn!("{:?}", e);
+                Response{status_code: 500, ..Response::empty_400()}
+            }
+        }
+    });
+}
+
+fn handle(request: &Request, server_impl: &Mutex<Server>) -> Result<Response> {
+    router!(request,
+        (POST) (/register) => {
+            server_impl.lock().unwrap().register(request)?;
+            Ok(Response{status_code: 200, ..Response::empty_204()})
+        },
+        (DELETE) (/register) => {
+            server_impl.lock().unwrap().unregister(request)?;
+            Ok(Response{status_code: 200, ..Response::empty_204()})
+        },
+        _ => Ok(Response::empty_404()),
+    )
 }

src/server.rs

@@ -0,0 +1,61 @@
+use std::{fs, path::Path};
+
+use chrono::Utc;
+use log::info;
+use rouille::Request;
+use anyhow::{Context, Result};
+
+use crate::Rewriter;
+
+pub struct Server {
+    rewriter: Rewriter,
+    // Where we write temporary files
+    workspace_dir: String,
+    // Directory to read configs from
+    template_dir: String,
+    // The symlink that is updated
+    config_dir: String,
+}
+
+impl Server {
+    pub fn new(rewriter: Rewriter, workspace_dir: String, template_dir: String, config_dir: String) -> Self {
+        Self {
+            rewriter,
+            workspace_dir,
+            template_dir,
+            config_dir,
+        }
+    }
+
+    pub fn register(&mut self, request: &Request) -> Result<()> {
+        let ip = request.remote_addr().ip().to_string();
+        info!("Registering {} as a handler", ip);
+        self.rewriter.add_replacement(ip);
+        self.generate_config()?;
+        Ok(())
+    }
+
+    pub fn unregister(&mut self, request: &Request) -> Result<()> {
+        let ip = request.remote_addr().ip().to_string();
+        info!("Deregistering {} as a handler", ip);
+        self.rewriter.remove_replacement(&ip);
+        self.generate_config()?;
+        Ok(())
+    }
+
+    pub fn generate_config(&self) -> Result<()> {
+        // Create a new directory in our workspace
+        let now = Utc::now();
+        // Writes into 2020/01/01/<unix timestamp>
+        // This will fail if we have multiple requests per second
+        let path = Path::new(&self.workspace_dir).join(&now.format("%Y/%m/%d/%s").to_string());
+        let path = path.as_os_str().to_str().unwrap();
+        fs::create_dir_all(path).with_context(|| "creating directory")?;
+        self.rewriter.rewrite_folder(&self.template_dir, path).with_context(|| "generating configs")?;
+        // Finally, symlink it to the output folder; only support Linux for now
+        let symlink = Path::new(&self.workspace_dir).join("symlink.tmp");
+        std::os::unix::fs::symlink(path, &symlink).with_context(|| "creating symlink")?;
+        fs::rename(symlink, &self.config_dir).with_context(|| "renaming symlink")?;
+        Ok(())
+    }
+}