Skip to content

Commit

Permalink
WIP: Refactor kernel interface to remove mock and global var
Browse files Browse the repository at this point in the history 
When a unit testing scheme was originally designed for the kernel
interface module in 2017 it was turned into a stateful struct to
accomodate a mock test environment that allowed the tests to call the
real functions and specify what exact commands they expected to be run.

This structure has not been maintained, new tests are not written in
that style and instead use separate testing functions. But the original
tests have not been modified much since that time. Leaving Rita with a
KI global variable that in production doesn't actually do anything all
to support mocking during tests.

This commit removes the mock structure and turns kernel interface into
just a grab bag of utility functions where tests functions are all split
out internally rather than relying on mock.

While we are touching all of these files it will also be renamed to
system_interface rather than kernel interface as we now interface with a
lot more than just linux kernel functions from this module.
  • Loading branch information
@jkilpatr
jkilpatr committed Sep 22, 2024
1 parent 368e5e2 commit dc920bd
Show file tree
Hide file tree
Showing 10 changed files with 493 additions and 606 deletions.
8 changes: 3 additions & 5 deletions althea_kernel_interface/src/babel.rs
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,5 @@
use crate::KernelInterface;
use crate::run_command;

impl dyn KernelInterface {
pub fn restart_babel(&self) {
let _res = self.run_command("/etc/init.d/babeld", &["restart"]);
}
pub fn restart_babel() {
let _res = run_command("/etc/init.d/babeld", &["restart"]);
}
14 changes: 6 additions & 8 deletions althea_kernel_interface/src/bridge_tools.rs
View file
Original file line number Diff line number Diff line change
@@ -1,15 +1,13 @@
//! Simple helper functions for brctl
use super::KernelInterface;
use crate::run_command;
use crate::KernelInterfaceError as Error;
use std::process::Output;

impl dyn KernelInterface {
pub fn add_if_to_bridge(&self, br: &str, iface: &str) -> Result<Output, Error> {
self.run_command("brctl", &["addif", br, iface])
}
pub fn add_if_to_bridge(br: &str, iface: &str) -> Result<Output, Error> {
run_command("brctl", &["addif", br, iface])
}

pub fn del_if_from_bridge(&self, br: &str, iface: &str) -> Result<Output, Error> {
self.run_command("brctl", &["delif", br, iface])
}
pub fn del_if_from_bridge(br: &str, iface: &str) -> Result<Output, Error> {
run_command("brctl", &["delif", br, iface])
}
25 changes: 11 additions & 14 deletions althea_kernel_interface/src/check_cron.rs
View file
Original file line number Diff line number Diff line change
@@ -1,19 +1,16 @@
use super::KernelInterface;
use crate::KernelInterfaceError as Error;
use std::process::{Command, Stdio};

impl dyn KernelInterface {
/// Checks if the cron service is running and starts it if it's not
pub fn check_cron(&self) -> Result<(), Error> {
Command::new("/etc/init.d/cron")
.args(["enable"])
.stdout(Stdio::piped())
.output()?;
Command::new("/etc/init.d/cron")
.args(["start"])
.stdout(Stdio::piped())
.output()?;
/// Checks if the cron service is running and starts it if it's not
pub fn check_cron() -> Result<(), Error> {
Command::new("/etc/init.d/cron")
.args(["enable"])
.stdout(Stdio::piped())
.output()?;
Command::new("/etc/init.d/cron")
.args(["start"])
.stdout(Stdio::piped())
.output()?;

Ok(())
}
Ok(())
}
260 changes: 125 additions & 135 deletions althea_kernel_interface/src/counter.rs
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,4 @@
use super::KernelInterface;
use crate::KernelInterfaceError as Error;
use crate::{run_command, KernelInterfaceError as Error};
use regex::Regex;
use std::collections::HashMap;
use std::net::IpAddr;
Expand Down Expand Up @@ -131,143 +130,136 @@ add zxcv 1234:5678:9801:2345:6789:0123:4567:8902,wg0 packets 123456789 bytes 987
}
}

impl dyn KernelInterface {
pub fn init_counter(&self, target: &FilterTarget) -> Result<(), Error> {
if self.does_nftables_exist() {
info!("Trying to init a counter!");
self.nft_init_counters(target.set_name(), target.chain(), target.nft_interface())?;
} else {
self.run_command(
"ipset",
&[
"create",
target.set_name(),
"hash:net,iface",
"family",
"inet6",
"counters",
],
)?;
self.add_iptables_rule(
"ip6tables",
&[
"-w",
"-I",
target.table(),
"1",
"-m",
"set",
"!",
"--match-set",
target.set_name(),
&format!("dst,{}", target.interface()),
"-j",
"SET",
"--add-set",
target.set_name(),
&format!("dst,{}", target.interface()),
],
)?;
}

Ok(())
pub fn init_counter(target: &FilterTarget) -> Result<(), Error> {
if does_nftables_exist() {
info!("Trying to init a counter!");
nft_init_counters(target.set_name(), target.chain(), target.nft_interface())?;
} else {
run_command(
"ipset",
&[
"create",
target.set_name(),
"hash:net,iface",
"family",
"inet6",
"counters",
],
)?;
add_iptables_rule(
"ip6tables",
&[
"-w",
"-I",
target.table(),
"1",
"-m",
"set",
"!",
"--match-set",
target.set_name(),
&format!("dst,{}", target.interface()),
"-j",
"SET",
"--add-set",
target.set_name(),
&format!("dst,{}", target.interface()),
],
)?;
}

fn parse_nft_set_counters(
&self,
set_name: &str,
) -> Result<HashMap<(IpAddr, String), u64>, Error> {
let mut ret_map = HashMap::new();
let out = self.run_command("nft", &["list", "set", "inet", "fw4", set_name])?;
// flush the list immediately to not missing accounting for any bytes
self.run_command("nft", &["flush", "set", "inet", "fw4", set_name])?;

let out = out.stdout;
let out = String::from_utf8(out).expect("fix command");
for line in out.lines() {
if line.contains("packets") {
let ret = line.replace("elements = { ", "");
let mut ret = ret.split_ascii_whitespace();

// line is in the form:
// ff02::1:6 . \"wg0\" counter packets 3 bytes 204,

let ip_addr = IpAddr::from_str(ret.next().unwrap_or(""))?;

ret.next();

let iface = ret.next();
let iface = if let Some(iface) = iface {
iface.replace('\"', "")
} else {
return Err(Error::ParseError(format!(
"No interface to parse for counter string {:?}",
line
)));
};

ret.next();
ret.next();
let packets: u64 = ret.next().unwrap_or("").parse()?;

ret.next();
let bytes: u64 = ret.next().unwrap_or("").replace(',', "").parse()?;

let total_bytes = (packets * 40) + bytes;

trace!(
"ipaddr, iface, packets, bytes, total: {:?}, {:?} {:?}, {:?}, {:?}",
ip_addr,
iface,
packets,
bytes,
total_bytes
);

ret_map.insert((ip_addr, iface), total_bytes);
}
}
Ok(())
}

Ok(ret_map)
}
fn parse_nft_set_counters(set_name: &str) -> Result<HashMap<(IpAddr, String), u64>, Error> {
let mut ret_map = HashMap::new();
let out = run_command("nft", &["list", "set", "inet", "fw4", set_name])?;
// flush the list immediately to not missing accounting for any bytes
run_command("nft", &["flush", "set", "inet", "fw4", set_name])?;

let out = out.stdout;
let out = String::from_utf8(out).expect("fix command");
for line in out.lines() {
if line.contains("packets") {
let ret = line.replace("elements = { ", "");
let mut ret = ret.split_ascii_whitespace();

// line is in the form:
// ff02::1:6 . \"wg0\" counter packets 3 bytes 204,

let ip_addr = IpAddr::from_str(ret.next().unwrap_or(""))?;

ret.next();

let iface = ret.next();
let iface = if let Some(iface) = iface {
iface.replace('\"', "")
} else {
return Err(Error::ParseError(format!(
"No interface to parse for counter string {:?}",
line
)));
};

ret.next();
ret.next();
let packets: u64 = ret.next().unwrap_or("").parse()?;

ret.next();
let bytes: u64 = ret.next().unwrap_or("").replace(',', "").parse()?;

let total_bytes = (packets * 40) + bytes;

trace!(
"ipaddr, iface, packets, bytes, total: {:?}, {:?} {:?}, {:?}, {:?}",
ip_addr,
iface,
packets,
bytes,
total_bytes
);

pub fn read_counters(
&self,
target: &FilterTarget,
) -> Result<HashMap<(IpAddr, String), u64>, Error> {
if self.does_nftables_exist() {
self.parse_nft_set_counters(target.set_name())
} else {
self.run_command(
"ipset",
&[
"create",
&format!("tmp_{}", target.set_name()),
"hash:net,iface",
"family",
"inet6",
"counters",
],
)?;

self.run_command(
"ipset",
&[
"swap",
&format!("tmp_{}", target.set_name()),
target.set_name(),
],
)?;

let output =
self.run_command("ipset", &["save", &format!("tmp_{}", target.set_name())])?;
let res = parse_ipset(&String::from_utf8(output.stdout)?);
trace!("ipset parsed into {:?}", res);

self.run_command("ipset", &["destroy", &format!("tmp_{}", target.set_name())])?;
res
ret_map.insert((ip_addr, iface), total_bytes);
}
}

Ok(ret_map)
}

pub fn read_counters(
target: &FilterTarget,
) -> Result<HashMap<(IpAddr, String), u64>, Error> {
if does_nftables_exist() {
parse_nft_set_counters(target.set_name())
} else {
run_command(
"ipset",
&[
"create",
&format!("tmp_{}", target.set_name()),
"hash:net,iface",
"family",
"inet6",
"counters",
],
)?;

run_command(
"ipset",
&[
"swap",
&format!("tmp_{}", target.set_name()),
target.set_name(),
],
)?;

let output = run_command("ipset", &["save", &format!("tmp_{}", target.set_name())])?;
let res = parse_ipset(&String::from_utf8(output.stdout)?);
trace!("ipset parsed into {:?}", res);

run_command("ipset", &["destroy", &format!("tmp_{}", target.set_name())])?;
res
}
}

#[test]
Expand All @@ -276,8 +268,6 @@ fn test_init_counter() {
use std::process::ExitStatus;
use std::process::Output;

use crate::KI;

let mut counter = 0;

KI.set_mock(Box::new(move |program, args| {
Expand Down
Loading

0 comments on commit dc920bd

Please sign in to comment.