smartproxy/rust/crates/rustproxy-nftables/src/rule_builder.rs

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"));
    }
}