add: multiple fixes, docs for deployment

README.md

@@ -6,22 +6,27 @@ open a port on all nodes. When the request lands on any node,
 it is forwarded to the correct pod via the network mesh kubernetes
 is using. In theory, there is one a hop penalty.
 
-But lets be honest. You're running with a single LB, probably a GCE free
+But lets be honest. You're running with a single LB, probably on a
-tier N1 VM. That extra hop doesn't matter.
+GCE free 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 will become the 'template_dir' in the argument to the LB.
 
-Move that directory to somewhere new, i.e. `/etc/nginx-template/`. Make
+Move that directory (i.e., `/etc/nginx`) to somewhere new,
-a symlink from that new directory to the old one (i.e.,
+(i.e. `/etc/nginx-template/`).
-`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/`).
 
+Create a symlink in the workspace which will point to the 'active' configuration;
+this will be updated by the tool (i.e., `ln -s /etc/nginx-template /var/skubelb/nginx`).
+Make a symlink from the old config directory to that symlink (i.e., 
+`ln -s /var/skubelb/nginx /etc/nginx`). Two layers are symlinks are used so we can
+have a non-root user run the rool when we setup the service.
+
 Make sure the user running the tool has read access to the template folder, read-write
 access to the workspace folder and config symlink.
 
@@ -30,12 +35,112 @@ Run the server with a command like:
 ```sh
 skubelb --needle some_node_ip \
     --workspace_dir /var/skubelb \
-    --config_symlink /etc/nginx \
+    --config_symlink /var/skubelb/nginx \
     --template_dir /etc/nginx-template
-    --listen 0.0.0.0:8080
+    --listen 0.0.0.0:8888
 ```
 
 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
+Next, configure the Kubernetes nodes to POST `http://loadbalancer:8888/register/${NODE_IP}` when
-they started, and DELETE `http://loadbalancer:8080/register` when they shutdown.
+they started, and DELETE `http://loadbalancer:8888/register/${NODE_IP}` when they shutdown. The easiest
+way to do this is with a daemonset; see the example below, or in daemon_set.yaml.
+
+#### Running as a system service
+
+Setup a user to run the service; make that user
+with `useradd -M skubelb`, Prevent logins with `usermod -L skubelb`.
+
+Make a workspace dir, `mkdir /var/skubelb/`, and give access to the
+daemon user, `chown skubelb:skubelb /var/skubelb/`.
+
+Add the systemd config to `/etc/systemd/system/skubelb.service`:
+
+```toml
+[Unit]
+Description=Simple Kubernetes Load Balancer
+After=network.target
+StartLimitIntervalSec=0
+
+[Service]
+Type=simple
+Restart=always
+RestartSec=1
+User=skubelb
+ExecStart=/usr/local/bin/skubelb --needle some_node_ip \
+    --workspace-dir /var/skubelb \
+    --config-symlink /var/skubelb/nginx \
+    --template-dir /etc/nginx-template \
+    --listen 0.0.0.0:8888 \
+    --reload-cmd '/usr/bin/sudo systemctl reload nginx'
+
+[Install]
+WantedBy=multi-user.target
+```
+
+Make sure you update `--needle some_node_ip` with something
+like `--needle 123.44.55.123`. The IP of node you tested with.
+
+### Sample Kubernets configuration
+
+Deploy this [daemon set](https://kubernetes.io/docs/concepts/workloads/controllers/daemonset/)
+to your cluster, replacing `lb_address` with the address of your load balancer.
+
+```yaml
+apiVersion: apps/v1
+kind: DaemonSet
+metadata:
+  name: skubelb
+  namespace: skubelb
+  labels:
+    k8s-app: skubelb
+spec:
+  selector:
+    matchLabels:
+      name: skubelb
+  template:
+    metadata:
+      labels:
+        name: skubelb
+    spec:
+      tolerations:
+      # these tolerations are to have the daemonset runnable on control plane nodes
+      # remove them if your control plane nodes should not run pods
+      - key: node-role.kubernetes.io/control-plane
+        operator: Exists
+        effect: NoSchedule
+      - key: node-role.kubernetes.io/master
+        operator: Exists
+        effect: NoSchedule
+      containers:
+      - name: skubelb
+        image: alpine/curl:latest
+        env:
+        - name: NODE_IP
+          valueFrom:
+            fieldRef:
+              fieldPath: status.hostIP
+        command: ['sh', '-c', 'echo "Wait for heat death of universe" && sleep 999999d']
+        lifecycle:
+          preStop:
+            exec:
+              command:
+              - sh
+              - "-c"
+              - "curl -X POST 10.128.0.2:8888/register/${NODE_IP}"
+          postStart:
+            exec:
+              command:
+              - sh
+              - "-c"
+              - "curl -X POST 10.128.0.2:8888/register/${NODE_IP}"
+        resources:
+          limits:
+            memory: 200Mi
+          requests:
+            cpu: 10m
+            memory: 100Mi
+      terminationGracePeriodSeconds: 30
+```
+
+NOTE: you should need to make an entry in the firewall to allow this request through. It is very important that the firewall entry has a source filter; it should only be allowed from the Kubernetes cluster. Nginx will forward traffic to any host that registers, and this could easily become a MitM vulnerability.

daemon_set.yaml

@@ -0,0 +1,54 @@
+apiVersion: apps/v1
+kind: DaemonSet
+metadata:
+  name: skubelb
+  namespace: skubelb
+  labels:
+    k8s-app: skubelb
+spec:
+  selector:
+    matchLabels:
+      name: skubelb
+  template:
+    metadata:
+      labels:
+        name: skubelb
+    spec:
+      tolerations:
+      # these tolerations are to have the daemonset runnable on control plane nodes
+      # remove them if your control plane nodes should not run pods
+      - key: node-role.kubernetes.io/control-plane
+        operator: Exists
+        effect: NoSchedule
+      - key: node-role.kubernetes.io/master
+        operator: Exists
+        effect: NoSchedule
+      containers:
+      - name: skubelb
+        image: alpine/curl:latest
+        env:
+        - name: NODE_IP
+          valueFrom:
+            fieldRef:
+              fieldPath: status.hostIP
+        command: ['sh', '-c', 'echo "Wait for heat death of universe" && sleep 999999d']
+        lifecycle:
+          preStop:
+            exec:
+              command:
+              - sh
+              - "-c"
+              - "curl -X POST 10.128.0.2:8888/register/${NODE_IP}"
+          postStart:
+            exec:
+              command:
+              - sh
+              - "-c"
+              - "curl -X POST 10.128.0.2:8888/register/${NODE_IP}"
+        resources:
+          limits:
+            memory: 200Mi
+          requests:
+            cpu: 10m
+            memory: 100Mi
+      terminationGracePeriodSeconds: 30

src/main.rs

@@ -1,4 +1,5 @@
 use std::sync::Mutex;
+use std::process::Command;
 
 use clap::Parser;
 
@@ -28,7 +29,7 @@ struct Args {
     /// and instead N lines (one per replacement) is added to
     /// the output.
     #[arg(short, long)]
-    rewrite_string: String,
+    needle: String,
 
     /// The folder which contains the templates that
     /// will be be searched for the needle.
@@ -45,8 +46,13 @@ struct Args {
     #[arg(short, long)]
     workspace_dir: String,
 
+    /// Address to listen for http requests on.
     #[arg(short, long, default_value_t = String::from("0.0.0.0:8080"))]
     listen: String,
+
+    /// Command to reload nginx.
+    #[arg(short, long, default_value_t = String::from("sudo nginx -s reload"))]
+    reload_cmd: String,
 }
 
 fn main() {
@@ -54,13 +60,32 @@ fn main() {
     env_logger::Builder::from_env(Env::default().default_filter_or("info")).init();
     let args = Args::parse();
 
-    let rewriter = Rewriter::new(args.rewrite_string);
+    let rewriter = Rewriter::new(args.needle);
     let server_impl = Mutex::new(Server::new(rewriter, args.workspace_dir, args.template_dir, args.config_symlink));
+    let reload_command = args.reload_cmd.leak();
+    let reload_command: Vec<&str> = reload_command.split_ascii_whitespace().collect();
 
     rouille::start_server(args.listen, move |request| {
         info!("Processing request: {:?}", request);
         match handle(request, &server_impl) {
-            Ok(resp) => resp,
+            Ok((resp, reload)) => {
+                if reload && reload_command.len() > 0 {
+                    let output = Command::new(reload_command[0])
+                        .args(&reload_command[1..])
+                        .output();
+                    match output {
+                        Ok(o) => {
+                            info!("Ran {:?}; exit code: {}", reload_command, o.status);
+                            info!("Ran {:?}; stdout: {}", reload_command, String::from_utf8_lossy(&o.stdout));
+                            info!("Ran {:?}; stderr: {}", reload_command, String::from_utf8_lossy(&o.stderr));
+                        },
+                        Err(e) => {
+                            warn!("Failed to run {:?}: {:?}", reload_command, e);
+                        }
+                    };
+                }
+                resp
+            },
             Err(e) => {
                 warn!("{:?}", e);
                 Response{status_code: 500, ..Response::empty_400()}
@@ -69,16 +94,16 @@ fn main() {
     });
 }
 
-fn handle(request: &Request, server_impl: &Mutex<Server>) -> Result<Response> {
+fn handle(request: &Request, server_impl: &Mutex<Server>) -> Result<(Response, bool)> {
     router!(request,
-        (POST) (/register) => {
+        (POST) (/register/{ip: String}) => {
-            server_impl.lock().unwrap().register(request)?;
+            server_impl.lock().unwrap().register(request, &ip)?;
-            Ok(Response{status_code: 200, ..Response::empty_204()})
+            Ok((Response{status_code: 200, ..Response::empty_204()}, true))
         },
-        (DELETE) (/register) => {
+        (DELETE) (/register/{ip: String}) => {
-            server_impl.lock().unwrap().unregister(request)?;
+            server_impl.lock().unwrap().unregister(request, &ip)?;
-            Ok(Response{status_code: 200, ..Response::empty_204()})
+            Ok((Response{status_code: 200, ..Response::empty_204()}, true))
         },
-        _ => Ok(Response::empty_404()),
+        _ => Ok((Response::empty_404(), false)),
     )
 }

src/rewriter.rs

@@ -58,23 +58,34 @@ impl Rewriter {
                 // Open 2 files; one to read and translate, and one to write.
                 let source_file = File::open(&src_path)?;
                 let mut dest_file = File::create(&dst_path)?;
-                let reader = BufReader::new(source_file);
+                let mut buff = Vec::with_capacity(2048);
+                let mut reader = BufReader::new(source_file);
 
-                for line in reader.lines() {
+                while let Ok(count) = reader.read_until(b'\n', &mut buff) {
-                    let line = line?;
+                    if count == 0 {
+                        break;
+                    }
                     // If the line is not subject to replacement, copy it and
                     // carry on.
-                    if !line.contains(&self.source) {
+                    let line = &buff[0..count];
-                        writeln!(dest_file, "{}", line)?;
+                    let m = contains(&self.source.as_bytes(), line);
+                    if m.is_none() {
+                        dest_file.write(&line)?;
+                        buff.clear();
                         continue;
                     }
+                    let m = m.unwrap();
+                    let start = &line[0..m.0];
+                    let end = &line[m.1..];
 
                     // Else, repeat the line multiple times, replacing the string
                     // in question
                     for replacement in &replacements {
-                        let new_line = line.replace(&self.source, &replacement);
+                        dest_file.write(start)?;
-                        writeln!(dest_file, "{}", new_line)?;
+                        dest_file.write(replacement.as_bytes())?;
+                        dest_file.write(end)?;
                     }
+                    buff.clear();
                 }
             }
         }
@@ -82,6 +93,21 @@ impl Rewriter {
     }
 }
 
+fn contains(needle: &[u8], haystack: &[u8]) -> Option<(usize, usize)> {
+    let mut i = 0;
+    while i + needle.len() <= haystack.len() {
+        let mut j = 0;
+        while i+j < haystack.len() && j < needle.len() && haystack[i+j] == needle[j] {
+            j += 1;
+        }
+        if j == needle.len() {
+            return Some((i, i+j));
+        }
+        i += 1;
+    }
+    None
+}
+
 #[cfg(test)]
 mod tests {
     use include_directory::{Dir, include_directory};

src/server.rs

@@ -27,18 +27,16 @@ impl Server {
         }
     }
 
-    pub fn register(&mut self, request: &Request) -> Result<()> {
+    pub fn register(&mut self, _request: &Request, ip: &str) -> Result<()> {
-        let ip = request.remote_addr().ip().to_string();
         info!("Registering {} as a handler", ip);
-        self.rewriter.add_replacement(ip);
+        self.rewriter.add_replacement(ip.to_string());
         self.generate_config()?;
         Ok(())
     }
 
-    pub fn unregister(&mut self, request: &Request) -> Result<()> {
+    pub fn unregister(&mut self, _request: &Request, ip: &str) -> Result<()> {
-        let ip = request.remote_addr().ip().to_string();
         info!("Deregistering {} as a handler", ip);
-        self.rewriter.remove_replacement(&ip);
+        self.rewriter.remove_replacement(ip);
         self.generate_config()?;
         Ok(())
     }

src/testdata/testsrc/hello.txt

@@ -1,4 +1,4 @@
 This is a line
 This is to_be_replaced line
 
-This is another line
+This is another line

src/testdata/testsrc/recursive/world.txt

@@ -1,3 +1,3 @@
 This is a to_be_replaced line.
 
-In a nested directory.
+In a nested directory.