router/src/rules.rs

use crate::{
    iptables::{self, Proto},
    startup::Interfaces,
};
use once_cell::sync::OnceCell;
use sled::{Db, Tree};
use std::net::Ipv4Addr;

static RULES_TREE: OnceCell<Tree> = OnceCell::new();

fn rules_tree(db: &Db) -> &'static Tree {
    RULES_TREE.get_or_init(|| db.open_tree("rules").unwrap())
}

#[derive(Debug, serde::Deserialize, serde::Serialize)]
pub struct Rule {
    pub(crate) proto: Proto,
    pub(crate) port: u16,
    pub(crate) kind: RuleKind,
}

impl Rule {
    pub(crate) fn as_forward(&self) -> Option<(Ipv4Addr, u16)> {
        match &self.kind {
            RuleKind::Forward { dest_ip, dest_port } => Some((*dest_ip, *dest_port)),
            _ => None,
        }
    }
}

#[derive(Debug, serde::Deserialize, serde::Serialize)]
#[serde(tag = "type")]
pub(crate) enum RuleKind {
    Accept,
    Forward {
        dest_ip: Ipv4Addr,
        #[serde(with = "serde_with::rust::display_fromstr")]
        dest_port: u16,
    },
}

pub(crate) async fn apply_all(db: &Db, interfaces: &Interfaces) -> Result<(), anyhow::Error> {
    for (_, rule) in read(db)? {
        apply(interfaces, rule).await?;
    }

    Ok(())
}

pub(crate) fn read(db: &Db) -> Result<Vec<(String, Rule)>, anyhow::Error> {
    rules_tree(db)
        .iter()
        .map(|res| {
            let (id, rule) = res?;

            let id = String::from_utf8_lossy(&id).to_string();
            tide::log::debug!("id: {}", id);
            tide::log::debug!("rule str: {}", String::from_utf8_lossy(&rule));
            let rule: Rule = serde_json::from_slice(&rule)?;
            tide::log::debug!("rule: {:?}", rule);

            Ok((id, rule)) as Result<(String, Rule), anyhow::Error>
        })
        .collect::<Result<Vec<_>, anyhow::Error>>()
}

pub(crate) async fn delete(db: &Db, rule_id: &str) -> Result<Rule, anyhow::Error> {
    let tree = rules_tree(db);

    let rule = tree
        .remove(rule_id.as_bytes())?
        .ok_or(anyhow::anyhow!("No rule with id {}", rule_id))?;

    tree.flush_async().await?;

    let rule: Rule = serde_json::from_slice(&rule)?;

    Ok(rule)
}

pub(crate) async fn unset(interfaces: &Interfaces, rule: Rule) -> Result<(), anyhow::Error> {
    match rule.kind {
        RuleKind::Accept => {
            iptables::delete_input_accept(
                rule.proto,
                &interfaces.external.interface,
                interfaces.external.ip,
                rule.port,
                interfaces.external.mask,
            )
            .await?;
        }
        RuleKind::Forward { dest_ip, dest_port } => {
            for info in &interfaces.internal {
                iptables::delete_forward_accept(
                    &interfaces.external.interface,
                    &info.interface,
                    rule.proto,
                    rule.port,
                )
                .await?;
            }
            iptables::delete_forward_prerouting(
                rule.proto,
                interfaces.external.ip,
                interfaces.external.mask,
                rule.port,
                dest_ip,
                dest_port,
            )
            .await?;
            for iface in &interfaces.internal {
                iptables::delete_forward_postrouting(
                    rule.proto,
                    iface.ip,
                    iface.mask,
                    interfaces.external.ip,
                    rule.port,
                    dest_ip,
                )
                .await?;
            }
            for iface in &interfaces.tunnel {
                iptables::delete_forward_postrouting(
                    rule.proto,
                    iface.ip,
                    iface.mask,
                    interfaces.external.ip,
                    rule.port,
                    dest_ip,
                )
                .await?;
            }
        }
    }

    Ok(())
}

pub(crate) async fn save(db: &Db, rule: &Rule) -> Result<(), anyhow::Error> {
    let tree = rules_tree(db);

    let s = serde_json::to_string(rule)?;
    let id = db.generate_id()?;

    tree.insert(rule_id(id).as_bytes(), s.as_bytes())?;

    tree.flush_async().await?;

    Ok(())
}

pub(crate) async fn apply(interfaces: &Interfaces, rule: Rule) -> Result<(), anyhow::Error> {
    match rule.kind {
        RuleKind::Accept => {
            iptables::input_accept(
                rule.proto,
                &interfaces.external.interface,
                interfaces.external.ip,
                rule.port,
                interfaces.external.mask,
            )
            .await?;
        }
        RuleKind::Forward { dest_ip, dest_port } => {
            for info in &interfaces.internal {
                iptables::forward_accept(
                    &interfaces.external.interface,
                    &info.interface,
                    rule.proto,
                    rule.port,
                )
                .await?;
            }
            iptables::forward_prerouting(
                rule.proto,
                interfaces.external.ip,
                interfaces.external.mask,
                rule.port,
                dest_ip,
                dest_port,
            )
            .await?;
            for iface in &interfaces.internal {
                iptables::forward_postrouting(
                    rule.proto,
                    iface.ip,
                    iface.mask,
                    interfaces.external.ip,
                    rule.port,
                    dest_ip,
                )
                .await?;
            }
            for iface in &interfaces.tunnel {
                iptables::forward_postrouting(
                    rule.proto,
                    iface.ip,
                    iface.mask,
                    interfaces.external.ip,
                    rule.port,
                    dest_ip,
                )
                .await?;
            }
        }
    }

    Ok(())
}

fn rule_id(id: u64) -> String {
    format!("rule-{}", id)
}

#[cfg(test)]
mod tests {
    use crate::{
        iptables::Proto,
        rules::{Rule, RuleKind},
    };

    #[test]
    fn can_serialize() {
        let rule = Rule {
            proto: Proto::Tcp,
            port: 53,
            kind: RuleKind::Forward {
                dest_ip: "192.168.6.84".parse().unwrap(),
                dest_port: 53,
            },
        };

        let s = serde_qs::to_string(&rule).unwrap();

        assert_eq!(
            s,
            "proto=Tcp&port=53&kind[type]=Forward&kind[dest_ip]=192.168.6.84&kind[dest_port]=53"
        );
    }
}