124 lines
3.9 KiB
Rust
124 lines
3.9 KiB
Rust
use rustproxy_config::{NfTablesOptions, NfTablesProtocol};
|
|
|
|
/// Build nftables DNAT rule for port forwarding.
|
|
pub fn build_dnat_rule(
|
|
table_name: &str,
|
|
chain_name: &str,
|
|
source_port: u16,
|
|
target_host: &str,
|
|
target_port: u16,
|
|
options: &NfTablesOptions,
|
|
) -> Vec<String> {
|
|
let protocol = match options.protocol.as_ref().unwrap_or(&NfTablesProtocol::Tcp) {
|
|
NfTablesProtocol::Tcp => "tcp",
|
|
NfTablesProtocol::Udp => "udp",
|
|
NfTablesProtocol::All => "tcp", // TODO: handle "all"
|
|
};
|
|
|
|
let mut rules = Vec::new();
|
|
|
|
// DNAT rule
|
|
rules.push(format!(
|
|
"nft add rule ip {} {} {} dport {} dnat to {}:{}",
|
|
table_name, chain_name, protocol, source_port, target_host, target_port,
|
|
));
|
|
|
|
// SNAT rule if preserving source IP is not enabled
|
|
if !options.preserve_source_ip.unwrap_or(false) {
|
|
rules.push(format!(
|
|
"nft add rule ip {} postrouting {} dport {} masquerade",
|
|
table_name, protocol, target_port,
|
|
));
|
|
}
|
|
|
|
// Rate limiting
|
|
if let Some(max_rate) = &options.max_rate {
|
|
rules.push(format!(
|
|
"nft add rule ip {} {} {} dport {} limit rate {} accept",
|
|
table_name, chain_name, protocol, source_port, max_rate,
|
|
));
|
|
}
|
|
|
|
rules
|
|
}
|
|
|
|
/// Build the initial table and chain setup commands.
|
|
pub fn build_table_setup(table_name: &str) -> Vec<String> {
|
|
vec![
|
|
format!("nft add table ip {}", table_name),
|
|
format!("nft add chain ip {} prerouting {{ type nat hook prerouting priority 0 \\; }}", table_name),
|
|
format!("nft add chain ip {} postrouting {{ type nat hook postrouting priority 100 \\; }}", table_name),
|
|
]
|
|
}
|
|
|
|
/// Build cleanup commands to remove the table.
|
|
pub fn build_table_cleanup(table_name: &str) -> Vec<String> {
|
|
vec![format!("nft delete table ip {}", table_name)]
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
fn make_options() -> NfTablesOptions {
|
|
NfTablesOptions {
|
|
preserve_source_ip: None,
|
|
protocol: None,
|
|
max_rate: None,
|
|
priority: None,
|
|
table_name: None,
|
|
use_ip_sets: None,
|
|
use_advanced_nat: None,
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_basic_dnat_rule() {
|
|
let options = make_options();
|
|
let rules = build_dnat_rule("rustproxy", "prerouting", 443, "10.0.0.1", 8443, &options);
|
|
assert!(rules.len() >= 1);
|
|
assert!(rules[0].contains("dnat to 10.0.0.1:8443"));
|
|
assert!(rules[0].contains("dport 443"));
|
|
}
|
|
|
|
#[test]
|
|
fn test_preserve_source_ip() {
|
|
let mut options = make_options();
|
|
options.preserve_source_ip = Some(true);
|
|
let rules = build_dnat_rule("rustproxy", "prerouting", 443, "10.0.0.1", 8443, &options);
|
|
// When preserving source IP, no masquerade rule
|
|
assert!(rules.iter().all(|r| !r.contains("masquerade")));
|
|
}
|
|
|
|
#[test]
|
|
fn test_without_preserve_source_ip() {
|
|
let options = make_options();
|
|
let rules = build_dnat_rule("rustproxy", "prerouting", 443, "10.0.0.1", 8443, &options);
|
|
assert!(rules.iter().any(|r| r.contains("masquerade")));
|
|
}
|
|
|
|
#[test]
|
|
fn test_rate_limited_rule() {
|
|
let mut options = make_options();
|
|
options.max_rate = Some("100/second".to_string());
|
|
let rules = build_dnat_rule("rustproxy", "prerouting", 80, "10.0.0.1", 8080, &options);
|
|
assert!(rules.iter().any(|r| r.contains("limit rate 100/second")));
|
|
}
|
|
|
|
#[test]
|
|
fn test_table_setup_commands() {
|
|
let commands = build_table_setup("rustproxy");
|
|
assert_eq!(commands.len(), 3);
|
|
assert!(commands[0].contains("add table ip rustproxy"));
|
|
assert!(commands[1].contains("prerouting"));
|
|
assert!(commands[2].contains("postrouting"));
|
|
}
|
|
|
|
#[test]
|
|
fn test_table_cleanup() {
|
|
let commands = build_table_cleanup("rustproxy");
|
|
assert_eq!(commands.len(), 1);
|
|
assert!(commands[0].contains("delete table ip rustproxy"));
|
|
}
|
|
}
|