diff --git a/Cargo.lock b/Cargo.lock index 28a39b6..bd2c17b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -45,6 +45,12 @@ version = "1.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a" +[[package]] +name = "bumpalo" +version = "3.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "37ccbd214614c6783386c1af30caf03192f17891059cecc394b4fb119e363de3" + [[package]] name = "bytes" version = "1.2.1" @@ -96,6 +102,28 @@ dependencies = [ "os_str_bytes", ] +[[package]] +name = "crossbeam-utils" +version = "0.8.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "51887d4adc7b564537b15adcfb307936f8075dfcd5f00dde9a9f1d29383682bc" +dependencies = [ + "cfg-if", + "once_cell", +] + +[[package]] +name = "dashmap" +version = "5.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3495912c9c1ccf2e18976439f4443f3fee0fd61f424ff99fde6a66b15ecb448f" +dependencies = [ + "cfg-if", + "hashbrown", + "lock_api", + "parking_lot_core", +] + [[package]] name = "data-encoding" version = "2.3.2" @@ -104,12 +132,14 @@ checksum = "3ee2393c4a91429dffb4bedf19f4d6abf27d8a732c8ce4980305d782e5426d57" [[package]] name = "dns-rebinder" -version = "0.2.1" +version = "0.2.2" dependencies = [ "clap", "env_logger", + "governor", "hex", "log", + "nonzero_ext", "rand", "tokio", "trust-dns-proto", @@ -150,6 +180,21 @@ dependencies = [ "percent-encoding", ] +[[package]] +name = "futures" +version = "0.3.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f73fe65f54d1e12b726f517d3e2135ca3125a437b6d998caf1962961f7172d9e" +dependencies = [ + "futures-channel", + "futures-core", + "futures-executor", + "futures-io", + "futures-sink", + "futures-task", + "futures-util", +] + [[package]] name = "futures-channel" version = "0.3.21" @@ -157,6 +202,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c3083ce4b914124575708913bca19bfe887522d6e2e6d0952943f5eac4a74010" dependencies = [ "futures-core", + "futures-sink", ] [[package]] @@ -165,26 +211,65 @@ version = "0.3.21" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0c09fd04b7e4073ac7156a9539b57a484a8ea920f79c7c675d05d289ab6110d3" +[[package]] +name = "futures-executor" +version = "0.3.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9420b90cfa29e327d0429f19be13e7ddb68fa1cccb09d65e5706b8c7a749b8a6" +dependencies = [ + "futures-core", + "futures-task", + "futures-util", +] + [[package]] name = "futures-io" version = "0.3.21" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fc4045962a5a5e935ee2fdedaa4e08284547402885ab326734432bed5d12966b" +[[package]] +name = "futures-macro" +version = "0.3.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "33c1e13800337f4d4d7a316bf45a567dbcb6ffe087f16424852d97e97a91f512" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "futures-sink" +version = "0.3.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "21163e139fa306126e6eedaf49ecdb4588f939600f0b1e770f4205ee4b7fa868" + [[package]] name = "futures-task" version = "0.3.21" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "57c66a976bf5909d801bbef33416c41372779507e7a6b3a5e25e4749c58f776a" +[[package]] +name = "futures-timer" +version = "3.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e64b03909df88034c26dc1547e8970b91f98bdb65165d6a4e9110d94263dbb2c" + [[package]] name = "futures-util" version = "0.3.21" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d8b7abd5d659d9b90c8cba917f6ec750a74e2dc23902ef9cd4cc8c8b22e6036a" dependencies = [ + "futures-channel", "futures-core", + "futures-io", + "futures-macro", + "futures-sink", "futures-task", + "memchr", "pin-project-lite", "pin-utils", "slab", @@ -198,7 +283,24 @@ checksum = "4eb1a864a501629691edf6c15a593b7a51eebaa1e8468e9ddc623de7c9b58ec6" dependencies = [ "cfg-if", "libc", - "wasi", + "wasi 0.11.0+wasi-snapshot-preview1", +] + +[[package]] +name = "governor" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "19775995ee20209163239355bc3ad2f33f83da35d9ef72dea26e5af753552c87" +dependencies = [ + "dashmap", + "futures", + "futures-timer", + "no-std-compat", + "nonzero_ext", + "parking_lot", + "quanta", + "rand", + "smallvec", ] [[package]] @@ -261,6 +363,15 @@ version = "2.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "879d54834c8c76457ef4293a689b2a8c59b076067ad77b15efafbb05f92a592b" +[[package]] +name = "js-sys" +version = "0.3.59" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "258451ab10b34f8af53416d1fdab72c22e805f0c92a1136d59470ec0b11138b2" +dependencies = [ + "wasm-bindgen", +] + [[package]] name = "lazy_static" version = "1.4.0" @@ -292,6 +403,15 @@ dependencies = [ "cfg-if", ] +[[package]] +name = "mach" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b823e83b2affd8f40a9ee8c29dbc56404c1e34cd2710921f2801e2cf29527afa" +dependencies = [ + "libc", +] + [[package]] name = "matches" version = "0.1.9" @@ -312,10 +432,22 @@ checksum = "57ee1c23c7c63b0c9250c339ffdc69255f110b298b901b9f6c82547b7b87caaf" dependencies = [ "libc", "log", - "wasi", + "wasi 0.11.0+wasi-snapshot-preview1", "windows-sys", ] +[[package]] +name = "no-std-compat" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b93853da6d84c2e3c7d730d6473e8817692dd89be387eb01b94d7f108ecb5b8c" + +[[package]] +name = "nonzero_ext" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "38bf9645c8b145698bb0b18a4637dcacbc421ea49bef2317e4fd8065a387cf21" + [[package]] name = "num_cpus" version = "1.13.1" @@ -418,6 +550,22 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "quanta" +version = "0.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "20afe714292d5e879d8b12740aa223c6a88f118af41870e8b6196e39a02238a8" +dependencies = [ + "crossbeam-utils", + "libc", + "mach", + "once_cell", + "raw-cpuid", + "wasi 0.10.2+wasi-snapshot-preview1", + "web-sys", + "winapi", +] + [[package]] name = "quote" version = "1.0.20" @@ -457,6 +605,15 @@ dependencies = [ "getrandom", ] +[[package]] +name = "raw-cpuid" +version = "10.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2c49596760fce12ca21550ac21dc5a9617b2ea4b6e0aa7d8dab8ff2824fc2bba" +dependencies = [ + "bitflags", +] + [[package]] name = "redox_syscall" version = "0.2.16" @@ -686,12 +843,82 @@ version = "0.9.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "49874b5167b65d7193b8aba1567f5c7d93d001cafc34600cee003eda787e483f" +[[package]] +name = "wasi" +version = "0.10.2+wasi-snapshot-preview1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fd6fbd9a79829dd1ad0cc20627bf1ed606756a7f77edff7b66b7064f9cb327c6" + [[package]] name = "wasi" version = "0.11.0+wasi-snapshot-preview1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9c8d87e72b64a3b4db28d11ce29237c246188f4f51057d65a7eab63b7987e423" +[[package]] +name = "wasm-bindgen" +version = "0.2.82" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc7652e3f6c4706c8d9cd54832c4a4ccb9b5336e2c3bd154d5cccfbf1c1f5f7d" +dependencies = [ + "cfg-if", + "wasm-bindgen-macro", +] + +[[package]] +name = "wasm-bindgen-backend" +version = "0.2.82" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "662cd44805586bd52971b9586b1df85cdbbd9112e4ef4d8f41559c334dc6ac3f" +dependencies = [ + "bumpalo", + "log", + "once_cell", + "proc-macro2", + "quote", + "syn", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-macro" +version = "0.2.82" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b260f13d3012071dfb1512849c033b1925038373aea48ced3012c09df952c602" +dependencies = [ + "quote", + "wasm-bindgen-macro-support", +] + +[[package]] +name = "wasm-bindgen-macro-support" +version = "0.2.82" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5be8e654bdd9b79216c2929ab90721aa82faf65c48cdf08bdc4e7f51357b80da" +dependencies = [ + "proc-macro2", + "quote", + "syn", + "wasm-bindgen-backend", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-shared" +version = "0.2.82" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6598dd0bd3c7d51095ff6531a5b23e02acdc81804e30d8f07afb77b7215a140a" + +[[package]] +name = "web-sys" +version = "0.3.59" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed055ab27f941423197eb86b2035720b1a3ce40504df082cac2ecc6ed73335a1" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + [[package]] name = "winapi" version = "0.3.9" diff --git a/Cargo.toml b/Cargo.toml index 991d805..a9701f3 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,7 +1,7 @@ [package] name = "dns-rebinder" authors = ["sinderella"] -version = "0.2.1" +version = "0.2.2" edition = "2021" # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html @@ -9,8 +9,10 @@ edition = "2021" [dependencies] clap = { version = "3.2.16", features = ["derive"] } env_logger = "0.9.0" +governor = "0.4.2" hex = "0.4.3" log = "0.4.17" +nonzero_ext = "0.3.0" rand = "0.8.5" tokio = { version = "1", features = ["full"] } trust-dns-proto = "0.21.2" diff --git a/src/main.rs b/src/main.rs index e03f6ff..c70ed00 100644 --- a/src/main.rs +++ b/src/main.rs @@ -13,13 +13,18 @@ use crate::cli::Commands; use clap::Parser; use env_logger::Builder; +use governor::clock::DefaultClock; +use governor::state::keyed::DashMapStateStore; +use governor::{Quota, RateLimiter}; use log::LevelFilter; use log::{error, info, warn}; +use nonzero_ext::nonzero; use rand::Rng; +use std::net::IpAddr; use std::net::{Ipv4Addr, SocketAddr}; use std::sync::Arc; - use tokio::net::UdpSocket; +use trust_dns_proto::op::MessageType; use trust_dns_proto::op::{Message, Query}; use trust_dns_proto::rr::rdata::SOA; @@ -159,6 +164,7 @@ fn process_query( pub(crate) async fn handle_connection( socket: &UdpSocket, + limiter: &RateLimiter, DefaultClock>, domain: &str, ns_records: &Option>, ns_public_ip: &Option, @@ -167,8 +173,19 @@ pub(crate) async fn handle_connection( let (len, addr) = socket.recv_from(&mut buffer).await.expect("receive failed"); let request = Message::from_vec(&buffer[0..len]).expect("failed parse of request"); + if let IpAddr::V4(ipv4) = addr.ip() { + match limiter.check_key(&ipv4) { + Ok(_) => {} + Err(_blocked) => { + warn!("blocked: {:?}", ipv4); + return Ok(()); + } + } + } + let mut message = Message::new(); message.set_id(request.id()); + message.set_message_type(MessageType::Response); message.set_recursion_desired(request.recursion_desired()); message.set_recursion_available(false); @@ -229,6 +246,13 @@ async fn main() -> Result<()> { .expect("couldn't bind to address"); info!("started listening on port: {:?}", port); + let quota = Quota::per_minute(nonzero!(20u32)); + let ratelimiter = Arc::new(RateLimiter::< + Ipv4Addr, + DashMapStateStore, + DefaultClock, + >::keyed(quota)); + let s = Arc::new(socket); let d = Arc::new(domain); let nsr = Arc::new(ns_records); @@ -239,10 +263,12 @@ async fn main() -> Result<()> { let domain_param = Arc::clone(&d); let ns_records_param = Arc::clone(&nsr); let ns_public_ip_param = Arc::clone(&npip); + let limiter_param = Arc::clone(&ratelimiter); let handler = tokio::spawn(async move { match handle_connection( &sock_param, + &limiter_param, &domain_param, &ns_records_param, &ns_public_ip_param,