From b7553573148f6011b703ed1908cea55378934f12 Mon Sep 17 00:00:00 2001
From: mdecimus <mauro@stalw.art>
Date: Sun, 19 May 2024 16:30:16 +0200
Subject: [PATCH] IP allowlists

---
 crates/common/src/config/network.rs   |  4 +-
 crates/common/src/lib.rs              |  6 +-
 crates/common/src/listener/blocked.rs | 94 ++++++++++++++++++++++++---
 3 files changed, 92 insertions(+), 12 deletions(-)

diff --git a/crates/common/src/config/network.rs b/crates/common/src/config/network.rs
index 440f83a1..05fc2b35 100644
--- a/crates/common/src/config/network.rs
+++ b/crates/common/src/config/network.rs
@@ -2,7 +2,7 @@ use utils::config::Config;
 
 use crate::{
     expr::{if_block::IfBlock, tokenizer::TokenMap},
-    listener::blocked::BlockedIps,
+    listener::blocked::{AllowedIps, BlockedIps},
     Network,
 };
 
@@ -12,6 +12,7 @@ impl Default for Network {
     fn default() -> Self {
         Self {
             blocked_ips: Default::default(),
+            allowed_ips: Default::default(),
             url: IfBlock::new::<()>(
                 "server.http.url",
                 [],
@@ -25,6 +26,7 @@ impl Network {
     pub fn parse(config: &mut Config) -> Self {
         let mut network = Network {
             blocked_ips: BlockedIps::parse(config),
+            allowed_ips: AllowedIps::parse(config),
             ..Default::default()
         };
         let token_map = &TokenMap::default().with_variables(CONNECTION_VARS);
diff --git a/crates/common/src/lib.rs b/crates/common/src/lib.rs
index 6c035eb9..a63e3548 100644
--- a/crates/common/src/lib.rs
+++ b/crates/common/src/lib.rs
@@ -38,7 +38,10 @@ use config::{
 };
 use directory::{core::secret::verify_secret_hash, Directory, Principal, QueryBy};
 use expr::if_block::IfBlock;
-use listener::{blocked::BlockedIps, tls::TlsManager};
+use listener::{
+    blocked::{AllowedIps, BlockedIps},
+    tls::TlsManager,
+};
 use mail_send::Credentials;
 use opentelemetry::KeyValue;
 use opentelemetry_sdk::{
@@ -81,6 +84,7 @@ pub struct Core {
 #[derive(Clone)]
 pub struct Network {
     pub blocked_ips: BlockedIps,
+    pub allowed_ips: AllowedIps,
     pub url: IfBlock,
 }
 
diff --git a/crates/common/src/listener/blocked.rs b/crates/common/src/listener/blocked.rs
index 761b3e54..55975a98 100644
--- a/crates/common/src/listener/blocked.rs
+++ b/crates/common/src/listener/blocked.rs
@@ -21,7 +21,11 @@
  * for more details.
 */
 
-use std::{fmt::Debug, net::IpAddr, sync::atomic::AtomicU8};
+use std::{
+    fmt::Debug,
+    net::{IpAddr, Ipv4Addr, Ipv6Addr},
+    sync::atomic::AtomicU8,
+};
 
 use ahash::AHashSet;
 use parking_lot::RwLock;
@@ -41,8 +45,17 @@ pub struct BlockedIps {
     limiter_rate: Option<Rate>,
 }
 
+#[derive(Clone)]
+pub struct AllowedIps {
+    ip_addresses: AHashSet<IpAddr>,
+    ip_networks: Vec<IpAddrMask>,
+    has_networks: bool,
+}
+
 pub const BLOCKED_IP_KEY: &str = "server.blocked-ip";
 pub const BLOCKED_IP_PREFIX: &str = "server.blocked-ip.";
+pub const ALLOWED_IP_KEY: &str = "server.allowed-ip";
+pub const ALLOWED_IP_PREFIX: &str = "server.allowed-ip.";
 
 impl BlockedIps {
     pub fn parse(config: &mut Config) -> Self {
@@ -77,21 +90,57 @@ impl BlockedIps {
     }
 }
 
+impl AllowedIps {
+    pub fn parse(config: &mut Config) -> Self {
+        let mut ip_addresses = AHashSet::new();
+        let mut ip_networks = Vec::new();
+
+        for ip in config
+            .set_values(ALLOWED_IP_KEY)
+            .map(IpAddrOrMask::parse_value)
+            .collect::<Vec<_>>()
+        {
+            match ip {
+                Ok(IpAddrOrMask::Ip(ip)) => {
+                    ip_addresses.insert(ip);
+                }
+                Ok(IpAddrOrMask::Mask(ip)) => {
+                    ip_networks.push(ip);
+                }
+                Err(err) => {
+                    config.new_parse_error(ALLOWED_IP_KEY, err);
+                }
+            }
+        }
+
+        // Add loopback addresses
+        ip_addresses.insert(IpAddr::V4(Ipv4Addr::LOCALHOST));
+        ip_addresses.insert(IpAddr::V6(Ipv6Addr::LOCALHOST));
+
+        AllowedIps {
+            ip_addresses,
+            has_networks: !ip_networks.is_empty(),
+            ip_networks,
+        }
+    }
+}
+
 impl Core {
     pub async fn is_fail2banned(&self, ip: IpAddr, login: String) -> store::Result<bool> {
         if let Some(rate) = &self.network.blocked_ips.limiter_rate {
-            let is_allowed = self
-                .storage
-                .lookup
-                .is_rate_allowed(format!("b:{}", ip).as_bytes(), rate, false)
-                .await?
-                .is_none()
-                && self
+            let is_allowed = self.is_ip_allowed(&ip)
+                || (self
                     .storage
                     .lookup
-                    .is_rate_allowed(format!("b:{}", login).as_bytes(), rate, false)
+                    .is_rate_allowed(format!("b:{}", ip).as_bytes(), rate, false)
                     .await?
-                    .is_none();
+                    .is_none()
+                    && self
+                        .storage
+                        .lookup
+                        .is_rate_allowed(format!("b:{}", login).as_bytes(), rate, false)
+                        .await?
+                        .is_none());
             if !is_allowed {
                 // Add IP to blocked list
                 self.network.blocked_ips.ip_addresses.write().insert(ip);
@@ -129,6 +178,17 @@ impl Core {
                     .iter()
                     .any(|network| network.matches(ip)))
     }
+
+    pub fn is_ip_allowed(&self, ip: &IpAddr) -> bool {
+        self.network.allowed_ips.ip_addresses.contains(ip)
+            || (self.network.allowed_ips.has_networks
+                && self
+                    .network
+                    .allowed_ips
+                    .ip_networks
+                    .iter()
+                    .any(|network| network.matches(ip)))
+    }
 }
 
 impl BlockedIps {
@@ -150,6 +210,20 @@ impl Default for BlockedIps {
     }
 }
 
+impl Default for AllowedIps {
+    fn default() -> Self {
+        // Add IPv4 and IPv6 loopback addresses
+        Self {
+            ip_addresses: AHashSet::from_iter([
+                IpAddr::V4(Ipv4Addr::LOCALHOST),
+                IpAddr::V6(Ipv6Addr::LOCALHOST),
+            ]),
+            ip_networks: Default::default(),
+            has_networks: Default::default(),
+        }
+    }
+}
+
 impl Clone for BlockedIps {
     fn clone(&self) -> Self {
         Self {