diff --git a/README.md b/README.md index 1d66b63..ec21fc4 100644 --- a/README.md +++ b/README.md @@ -38,4 +38,83 @@ skubelb --needle some_node_ip \ 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. \ No newline at end of file +they started, and DELETE `http://loadbalancer:8080/register` when they shutdown. + +#### Running as a system service + +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 /etc/nginx \ + --template_dir /etc/nginx-template + --listen 0.0.0.0:8080 + --reload-cmd '/usr/bin/sudo systemctl reload nginx' + +[Install] +WantedBy=multi-user.target +``` + +### 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 + command: ['sh', '-c', 'echo "Wait for heat death of universe" && sleep 999999d'] + lifecycle: + postStart: + exec: + command: ['curl', '-X', 'POST', '34.56.7.198:8888/register'] + preStart: + exec: + command: ['curl', '-X', 'POST', '34.56.7.198:8888/register'] + 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. diff --git a/daemon_set.yaml b/daemon_set.yaml new file mode 100644 index 0000000..65e07a4 --- /dev/null +++ b/daemon_set.yaml @@ -0,0 +1,43 @@ +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 + command: ['sh', '-c', 'echo "Wait for heat death of universe" && sleep 999999d'] + lifecycle: + postStart: + exec: + command: ['curl', '-X', 'POST', '10.128.0.2:8888/register'] + preStop: + exec: + command: ['curl', '-X', 'DELETE', '10.128.0.2:8888/register'] + resources: + limits: + memory: 200Mi + requests: + cpu: 10m + memory: 100Mi + terminationGracePeriodSeconds: 30 \ No newline at end of file diff --git a/src/main.rs b/src/main.rs index 6cf4649..eb2805c 100644 --- a/src/main.rs +++ b/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) -> Result { +fn handle(request: &Request, server_impl: &Mutex) -> Result<(Response, bool)> { router!(request, (POST) (/register) => { server_impl.lock().unwrap().register(request)?; - Ok(Response{status_code: 200, ..Response::empty_204()}) + Ok((Response{status_code: 200, ..Response::empty_204()}, true)) }, (DELETE) (/register) => { server_impl.lock().unwrap().unregister(request)?; - 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)), ) } \ No newline at end of file diff --git a/src/rewriter.rs b/src/rewriter.rs index 45dd169..9c7dc16 100644 --- a/src/rewriter.rs +++ b/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() { - let line = line?; + while let Ok(count) = reader.read_until(b'\n', &mut buff) { + if count == 0 { + break; + } // If the line is not subject to replacement, copy it and // carry on. - if !line.contains(&self.source) { - writeln!(dest_file, "{}", line)?; + let line = &buff[0..count]; + 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); - writeln!(dest_file, "{}", new_line)?; + dest_file.write(start)?; + 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}; diff --git a/src/testdata/testsrc/hello.txt b/src/testdata/testsrc/hello.txt index 444809a..74638e7 100644 --- a/src/testdata/testsrc/hello.txt +++ b/src/testdata/testsrc/hello.txt @@ -1,4 +1,4 @@ This is a line This is to_be_replaced line -This is another line \ No newline at end of file +This is another line diff --git a/src/testdata/testsrc/recursive/world.txt b/src/testdata/testsrc/recursive/world.txt index e1baabd..d789711 100644 --- a/src/testdata/testsrc/recursive/world.txt +++ b/src/testdata/testsrc/recursive/world.txt @@ -1,3 +1,3 @@ This is a to_be_replaced line. -In a nested directory. \ No newline at end of file +In a nested directory.