router/src/rules.rs

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

#[derive(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(serde::Deserialize, serde::Serialize)]
#[serde(tag = "type")]
pub(crate) enum RuleKind {
    Accept,
    Forward { dest_ip: Ipv4Addr, 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> {
    db.iter()
        .map(|res| {
            let (id, rule) = res?;

            let id = String::from_utf8_lossy(&id).to_string();
            let rule: Rule = serde_json::from_slice(&rule)?;

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

pub(crate) fn delete(db: &Db, rule_id: String) -> Result<Rule, anyhow::Error> {
    let rule = db
        .remove(rule_id.as_bytes())?
        .ok_or(anyhow::anyhow!("No rule with id {}", rule_id))?;

    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(
                &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?;
            iptables::delete_forward_postrouting(interfaces.external.ip, rule.port, dest_ip)
                .await?;
        }
    }

    Ok(())
}

pub(crate) fn save(db: &Db, rule: &Rule) -> Result<(), anyhow::Error> {
    let s = serde_json::to_string(rule)?;
    let id = db.generate_id()?;

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

    Ok(())
}

pub(crate) async fn apply(interfaces: &Interfaces, rule: Rule) -> Result<(), anyhow::Error> {
    match rule.kind {
        RuleKind::Accept => {
            iptables::input_accept(
                &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?;
            iptables::forward_postrouting(interfaces.external.ip, rule.port, dest_ip).await?;
        }
    }

    Ok(())
}

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