From a24c3c5da46e7b8e3da6913169fb8a7032a8f6eb Mon Sep 17 00:00:00 2001 From: Sachymetsu Date: Wed, 7 Oct 2026 20:02:39 +0200 Subject: [PATCH] IP Range block feature --- Cargo.lock | 105 +++++++++++++++++++++--------- Cargo.toml | 1 + crates/nailconfig/Cargo.toml | 1 + crates/nailconfig/src/lib.rs | 3 + crates/nailrater/Cargo.toml | 2 + crates/nailrater/src/futures.rs | 55 +++++++++++----- crates/nailrater/src/lib.rs | 26 +++++++- crates/nailrater/src/scheduler.rs | 7 +- crates/nailserve/src/lib.rs | 1 + defaults/pit.default.toml | 4 ++ 10 files changed, 151 insertions(+), 54 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 7f7a869..4b4fe52 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -144,9 +144,9 @@ checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5" [[package]] name = "cc" -version = "1.5.0" +version = "1.6.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "207aedd8d29a0eb804986bc99eed0d43719fec84a9b7f9c8b70cf83b97cffcc1" +checksum = "f74872d07caf508b30a21f6836e7d7016a2eaf7d9ff4f48deaa58cd8a0407630" dependencies = [ "find-msvc-tools", "jobserver", @@ -253,9 +253,9 @@ dependencies = [ [[package]] name = "config" -version = "0.15.26" +version = "0.15.27" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "790374c727a4cb20877dd2cc99ce7011cbe03b0b8df083acccd37df0d161304d" +checksum = "38e9fbc53afcec54441422f1c596cd054655b32e70e1cbac956651b7b175a404" dependencies = [ "pathdiff", "serde_core", @@ -380,9 +380,9 @@ dependencies = [ [[package]] name = "either" -version = "1.18.0" +version = "1.19.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "252afb9ae5eaa683babdc6a068b3f5726eb19e05070c731f9b2a23a7c3e8ed34" +checksum = "0e9c71c2167ca323c882b99918929403426e2373ea17242ff5653e0d5e1058be" [[package]] name = "equivalent" @@ -618,9 +618,9 @@ checksum = "e4eba85ea1d0a966a983acd07deee566e67395d2d96b6fb39e62b5a833f1eb0b" [[package]] name = "h2" -version = "0.4.19" +version = "0.4.20" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ef8e5e5a340588f4452631496976cf8636d4a7ecf600239fdc27615d2530bc16" +checksum = "7d29020232d6aa3fb1daca64c1127cf662cf97f254ae16c18c05b8ab635fc118" dependencies = [ "atomic-waker", "bytes", @@ -719,9 +719,9 @@ checksum = "df3b46402a9d5adb4c86a0cf463f42e19994e3ee891101b1841f30a545cb49a9" [[package]] name = "hyper" -version = "1.11.1" +version = "1.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "27b501faa50e7a26c3d3560ca625132f4078a17771f4810baf70475ae48cbe43" +checksum = "2c3e324da4c95177d6291d4c8730197c0d1822f8a9766814a4a44fa5ab797c9c" dependencies = [ "atomic-waker", "bytes", @@ -789,6 +789,25 @@ dependencies = [ "hashbrown 0.17.1", ] +[[package]] +name = "ipnet" +version = "2.12.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "791930b43c0d5973160d90a8f3894509f2b273430f5c5c73b668636d0287c5c0" +dependencies = [ + "serde", +] + +[[package]] +name = "ipnet-trie" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0cb57916213e04ad56f36e59f27cf13a874615dcf7e6f0a7a5708581cc84080d" +dependencies = [ + "ipnet", + "prefix-trie", +] + [[package]] name = "itertools" version = "0.13.0" @@ -845,24 +864,25 @@ dependencies = [ [[package]] name = "lazy_static" -version = "1.5.0" +version = "1.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" +checksum = "20870f649af7073d53e38067b2a84312175d56ea15217e1b15bc83506ec50afb" [[package]] name = "libabort" -version = "0.1.9" +version = "0.1.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cec5d1db7977801dd3e593f88a15906cac3c4eb7a69e38b6cc162cb020b22d7d" +checksum = "f8832aae8349f14a90350b08208f8c1874a047cd54eef898fb6738b2e90e9555" dependencies = [ + "autocfg", "rustversion-detect", ] [[package]] name = "libc" -version = "0.2.189" +version = "0.2.190" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3eaf3ede3fee6db1a4c2ee091bf8a8b4dccdc6d17f656fb07896ee72867612f2" +checksum = "ce5d3ddc6d3fa000eb1536d85e147bfe31aacaba692ed6a876f95cb7c855be78" [[package]] name = "libm" @@ -937,9 +957,9 @@ dependencies = [ [[package]] name = "mio" -version = "1.2.3" +version = "1.2.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4b18443e9c262bfe8fa82f51666e2642c53393f7e5c27b3e1aeab922cff5b9d8" +checksum = "1788edb87fdc09c7e26304471e2f5be8cdefb1b6930d6e3985fc02ff53bf86ee" dependencies = [ "libc", "wasi", @@ -975,6 +995,7 @@ version = "0.1.0" dependencies = [ "color-eyre", "config", + "ipnet", "nailbox", "serde", ] @@ -1088,6 +1109,8 @@ version = "0.1.0" dependencies = [ "futures-lite", "hyper", + "ipnet", + "ipnet-trie", "nailbody", "nailbox", "nailconfig", @@ -1265,9 +1288,9 @@ dependencies = [ [[package]] name = "objc2" -version = "0.6.4" +version = "0.6.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3a12a8ed07aefc768292f076dc3ac8c48f3781c8f2d5851dd3d98950e8c5a89f" +checksum = "08849bbd4767dfae9457696856ae1c84fe4e0281bbe4a7abff2d0e06fb7981f8" dependencies = [ "objc2-encode", ] @@ -1524,6 +1547,16 @@ dependencies = [ "zerocopy", ] +[[package]] +name = "prefix-trie" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eb5f930995ba4986bd239ba8d8fded67cad82d1db329c4f316f312847cba16aa" +dependencies = [ + "ipnet", + "num-traits", +] + [[package]] name = "proc-macro2" version = "1.0.107" @@ -1724,9 +1757,9 @@ checksum = "cf54715a573b99ac80df0bc206da022bcd442c974952c7b9720069370852e21f" [[package]] name = "rustversion-detect" -version = "0.1.3" +version = "0.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4cfa9e87e97427c3a1b472eace073b2bc577ad0e1444c128d938b3d5bcdacb17" +checksum = "5853ad4cb8296e825e04116aefddec19e63bc2841b41b7c93e3947edc8cb8b74" [[package]] name = "saa" @@ -1935,9 +1968,9 @@ dependencies = [ [[package]] name = "tokio" -version = "1.53.1" +version = "1.53.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "202caea871b69668250d242070849eb495be178ed697a3e98aebce5bc81a0bed" +checksum = "e95f91fcc7a621e8b030f6aa23c71fe9838ae2fb4d8118b75602a328f5144044" dependencies = [ "bytes", "libc", @@ -1979,10 +2012,12 @@ version = "1.1.6+spec-1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "920602543f0911ab71da12c50d59701da54c196d1a2bf5cb4b75667f137a406a" dependencies = [ + "indexmap", "serde_core", "serde_spanned", "toml_datetime", "toml_parser", + "toml_writer", "winnow", ] @@ -2004,6 +2039,12 @@ dependencies = [ "winnow", ] +[[package]] +name = "toml_writer" +version = "1.1.2+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7d56353a2a665ad0f41a421187180aab746c8c325620617ad883a99a1cbe66d2" + [[package]] name = "tonic" version = "0.14.6" @@ -2206,9 +2247,9 @@ checksum = "c6f5d3c3b1bf09027a88a6bc961fc00497d651009560b5463668dc81b0fa87a8" [[package]] name = "uuid" -version = "1.26.1" +version = "1.27.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2ef6dac1e96601b4fb3acccccff2139741fcb757cb9a36089bf5be91cfb285ce" +checksum = "97277d36b9c3ace13e58fa6e753f8b0bbbf302a18dd193240d70f9e29681059a" dependencies = [ "getrandom 0.4.3", "js-sys", @@ -2234,9 +2275,9 @@ dependencies = [ [[package]] name = "want" -version = "0.3.1" +version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bfa7760aed19e106de2c7c0b581b509f2f25d3dacaf737cb82ac61bc6d760b0e" +checksum = "ec4cdd0dd910afe868b7ef477227d8d538b46b3075031afee8a9f2acb0a2ed0b" dependencies = [ "try-lock", ] @@ -2393,18 +2434,18 @@ dependencies = [ [[package]] name = "zerocopy" -version = "0.8.58" +version = "0.8.61" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c17e8fafad82b542ff3717217ecdc736231b59e387768c9630123b4ce4d2db44" +checksum = "879fb705ce98c32e41ebdb970fbe1204f8492423b314c6ab0354c3e7b5542866" dependencies = [ "zerocopy-derive", ] [[package]] name = "zerocopy-derive" -version = "0.8.58" +version = "0.8.61" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "595f56e044df4f46a0c9a626f65c3d99eb8488f7e8a8baa12dd76326d9710bf2" +checksum = "708882a28301d604fa039cc7727607a98d04c4b86dc76ec9cc683f805d709759" dependencies = [ "proc-macro2", "quote", diff --git a/Cargo.toml b/Cargo.toml index 5c83f72..3a4ba38 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -57,6 +57,7 @@ opentelemetry_sdk = { version = "0.32.0", default-features = false, features = [ ] } opentelemetry-semantic-conventions = "0.32" winnow = { version = "1.0.3", features = ["simd"] } +ipnet = "2.12.2" [package] name = "nailpit" diff --git a/crates/nailconfig/Cargo.toml b/crates/nailconfig/Cargo.toml index 46ae10d..d8e3564 100644 --- a/crates/nailconfig/Cargo.toml +++ b/crates/nailconfig/Cargo.toml @@ -15,6 +15,7 @@ serde.workspace = true nailbox = { path = "../nailbox" } config = { version = "0.15.11", default-features = false, features = ["toml"] } color-eyre.workspace = true +ipnet = { workspace = true, features = ["serde"] } [lib] bench = false diff --git a/crates/nailconfig/src/lib.rs b/crates/nailconfig/src/lib.rs index 104b962..19e1aa5 100644 --- a/crates/nailconfig/src/lib.rs +++ b/crates/nailconfig/src/lib.rs @@ -7,6 +7,7 @@ use core::num::NonZero; use std::{ops::Deref, sync::Arc}; use color_eyre::{Result, eyre::Context}; +use ipnet::IpNet; use nailbox::try_arc_within; #[derive(Debug, serde::Serialize, serde::Deserialize)] @@ -40,6 +41,8 @@ pub struct ServerConfig { pub pit_routes: Vec, pub socket_addr: String, pub worker_threads: NonZero, + #[serde(default)] + pub ip_blocks: Vec, } #[derive(Debug, serde::Serialize, serde::Deserialize)] diff --git a/crates/nailrater/Cargo.toml b/crates/nailrater/Cargo.toml index 6989466..e6580cf 100644 --- a/crates/nailrater/Cargo.toml +++ b/crates/nailrater/Cargo.toml @@ -25,6 +25,8 @@ tracing-futures.workspace = true scc.workspace = true tokio.workspace = true rapidhash.workspace = true +ipnet = "2.12.2" +ipnet-trie = "0.3.1" [lints] workspace = true diff --git a/crates/nailrater/src/futures.rs b/crates/nailrater/src/futures.rs index 4e2c4ed..ff5a91f 100644 --- a/crates/nailrater/src/futures.rs +++ b/crates/nailrater/src/futures.rs @@ -2,7 +2,10 @@ use core::{ pin::Pin, task::{Context, Poll}, }; -use std::{sync::Arc, time::Instant}; +use std::{ + sync::Arc, + time::{Duration, Instant}, +}; use futures_lite::{FutureExt, future::Boxed, ready}; use hyper::{ @@ -10,6 +13,7 @@ use hyper::{ body::{Bytes as ArcBytes, Incoming}, header::{CONTENT_ENCODING, CONTENT_TYPE, HeaderValue}, }; +use ipnet::IpNet; use nailbody::Body; use nailbox::boxed_future_within; use nailip::IdentifiedPeer; @@ -19,7 +23,9 @@ use rapidhash::quality::RandomState; use scc::hash_map::Entry; use tokio::time::{Sleep, sleep}; -use crate::{PEERS, PeerRecord, PeerState, modes::LimitModes, scheduler::PRUNING_SCHEDULER}; +use crate::{ + BLOCKS, PEERS, PeerRecord, PeerState, modes::LimitModes, scheduler::PRUNING_SCHEDULER, +}; pin_project! { pub struct NailedResponseFuture { @@ -102,8 +108,10 @@ fn prune() -> Boxed<()> { boxed_future_within(async move || { tracing::trace!("PRUNING STARTED"); + // Keep recorded peers for at least 2 hours, pruning them if not visited + // within that period of time. PEERS - .retain_async(|_, v| v.last_seen.elapsed() < crate::PEER_TIMEOUT) + .retain_async(|_, v| v.last_seen.elapsed() < Duration::from_hours(2)) .await }) } @@ -165,24 +173,41 @@ where let req = req.take().unwrap(); // If this panics, it is because the future was polled twice in the wrong state. let peer_record = entry .and_modify(|p| { - p.count += 1; p.last_seen = Instant::now(); - let next_mode = mode.limit(&p.count); + if matches!(p.state, PeerState::Ready | PeerState::Delay(_)) { + p.count += 1; - if p.state != next_mode && matches!(next_mode, PeerState::SpicyDrop) { - tracing::info!("remote.peer" = %peer, "spicy.mode" = ?&p.supports_spicy, "Spicy mode activated!"); - } + let next_mode = mode.limit(&p.count); - p.state = next_mode; + if p.state != next_mode && matches!(next_mode, PeerState::SpicyDrop) { + tracing::info!("remote.peer" = %peer, "spicy.mode" = ?&p.supports_spicy, "Spicy mode activated!"); + } + + p.state = next_mode; + } }) .or_insert_with(|| { - tracing::info!("remote.peer" = %peer, "New remote peer"); - PeerRecord { - count: 1, - state: mode.limit(&1), - last_seen: Instant::now(), - supports_spicy: SpicyPayloadKind::accepts_encoding(req.headers()), + // Safe to unwrap here, because we initialise the locked data before we start + // issuing rater futures. + let blocked = BLOCKS.get().unwrap().longest_match(&IpNet::from(peer.peer)); + + if let Some((blocked_net, _)) = blocked { + tracing::info!("remote.peer" = %peer, "blocked.network" = %blocked_net, "Blocked remote peer"); + PeerRecord { + count: 1, + state: if spicy_payloads.is_some() { PeerState::SpicyDrop } else { PeerState::Drop }, + last_seen: Instant::now(), + supports_spicy: SpicyPayloadKind::accepts_encoding(req.headers()), + } + } else { + tracing::info!("remote.peer" = %peer, "New remote peer"); + PeerRecord { + count: 1, + state: mode.limit(&1), + last_seen: Instant::now(), + supports_spicy: SpicyPayloadKind::accepts_encoding(req.headers()), + } } }); diff --git a/crates/nailrater/src/lib.rs b/crates/nailrater/src/lib.rs index 069b9e2..e904913 100644 --- a/crates/nailrater/src/lib.rs +++ b/crates/nailrater/src/lib.rs @@ -3,7 +3,7 @@ mod modes; mod scheduler; use std::{ - sync::{Arc, LazyLock}, + sync::{Arc, LazyLock, OnceLock}, time::{Duration, Instant}, }; @@ -11,17 +11,18 @@ use futures::NailedResponseFuture; use hyper::{Request, Response, body::Incoming}; pub use modes::LimitModes; use nailbody::Body; +use nailconfig::NailConfig; use nailip::IdentifiedPeer; use nailspicy::{SpicyPayloadKind, SpicyPayloads}; use rapidhash::quality::RandomState; use scc::HashMap; use tracing_futures::{Instrument, Instrumented}; -const PEER_TIMEOUT: Duration = Duration::from_secs(60 * 2); - static PEERS: LazyLock> = LazyLock::new(Default::default); +static BLOCKS: OnceLock> = OnceLock::new(); + #[derive(Debug, Default, Clone, Copy, PartialEq, Eq)] enum PeerState { #[default] @@ -51,9 +52,28 @@ impl NailRater { pub fn new( mode: impl Into, spicy_payload: Option>, + app_config: Arc, inner: S, ) -> Self { LazyLock::force(&PEERS); + BLOCKS.get_or_init(|| { + let mut blocks: ipnet_trie::IpnetTrie<_> = Default::default(); + + if !app_config.server.ip_blocks.is_empty() { + for blocked_network in app_config.server.ip_blocks.iter().copied() { + blocks.insert(blocked_network, ()); + } + + if !blocks.is_empty() { + tracing::info!( + "blocked.network" = ?blocks.ip_count(), + "Initialised IP range blocker" + ); + } + } + + blocks + }); Self { mode: mode.into(), diff --git a/crates/nailrater/src/scheduler.rs b/crates/nailrater/src/scheduler.rs index 7aaac52..24c2a2c 100644 --- a/crates/nailrater/src/scheduler.rs +++ b/crates/nailrater/src/scheduler.rs @@ -1,9 +1,7 @@ -use std::time::Instant; +use std::time::{Duration, Instant}; use parking_lot::Mutex; -use crate::PEER_TIMEOUT; - pub(crate) struct Scheduler { value: Mutex>, } @@ -23,7 +21,8 @@ impl Scheduler { let mut inner = self.value.lock(); let elapsed = inner.get_or_insert_with(Instant::now).elapsed(); - if elapsed >= PEER_TIMEOUT { + // Schedule a pruning every 10 minutes + if elapsed >= Duration::from_mins(10) { inner.take(); drop(inner); diff --git a/crates/nailserve/src/lib.rs b/crates/nailserve/src/lib.rs index eff8c0f..e299cf5 100644 --- a/crates/nailserve/src/lib.rs +++ b/crates/nailserve/src/lib.rs @@ -32,6 +32,7 @@ async fn server_loop( Rc::new(NailRater::new( LimitModes::from(&app.config.rate_limiting), app.spicy_payloads.get(), + app.config.clone_inner(), Rc::new(NailResponder::new(app)), )), )))); diff --git a/defaults/pit.default.toml b/defaults/pit.default.toml index c77822a..64b982a 100644 --- a/defaults/pit.default.toml +++ b/defaults/pit.default.toml @@ -10,6 +10,10 @@ socket_addr = "0.0.0.0:3000" # will benefit from more threads to balance the generation load better. worker_threads = 1 +# A list of CIDR notation network addresses that are to be blocked. Can be +# specific addresses, or address ranges, such as `10.168.2.99/32` or `10.0.1.0/24` respectively +ip_blocks = [] + [generator] # Input text files to feed the Markov Chain. Uses glob format. Multiple files that are # included will be turned into multiple different chains, so generated pages have randomised -- 2.54.0