Feat: Rate limiting (#42)

* feat: Simple queries per IP rate limiting

* extracted rate limiter

* improved rate limiter hasher

* added rate limiting for dht queries

* refactored error handling

* trying to fix broken pipe bug

* reduced logs again

* fixed half invalid pkarr keys

* update cmd argument names

* updated rate limit default value

* reduce dht limit to 10 per second

* added burst sizes

.vscode/launch.json

@@ -20,10 +20,12 @@
                 }
             },
             "args": [
+                "--ip-rate-limit-dht", "1",
+                "--verbose",
             ],
             "cwd": "${workspaceFolder}",
             "env": {
-                "RUST_LOG": "debug"
+                // "RUST_LOG": "debug"
                 // "RUST_LOG": "pkdns=debug,any_dns=debug,pkarr=debug,mainline=debug"
             }
         },

README.md

@@ -96,13 +96,29 @@ Other services might occupy the port 53 already. For example, [Docker Desktop](h
 Usage: pkdns [OPTIONS]
 
 Options:
-  -f, --forward <forward>  ICANN fallback DNS server. IP:Port [default: 8.8.8.8:53]
-  -s, --socket <socket>    Socket the server should listen on. IP:Port [default: 0.0.0.0:53]
-  -v, --verbose            Show verbose output.
-      --min-ttl <min-ttl>  Minimum number of seconds a value is cached for before being refreshed. [default: 300]
-      --max-ttl <max-ttl>  Maximum number of seconds before a cached value gets auto-refreshed. [default: 86400]
-  -h, --help               Print help
-  -V, --version            Print version
+  -f, --forward <forward>
+          ICANN fallback DNS server. IP:Port [default: 8.8.8.8:53]
+  -s, --socket <socket>
+          Socket the server should listen on. IP:Port [default: 0.0.0.0:53]
+  -v, --verbose
+          Show verbose output.
+      --min-ttl <min-ttl>
+          Minimum number of seconds a value is cached for before being refreshed. [default: 300]
+      --max-ttl <max-ttl>
+          Maximum number of seconds before a cached value gets auto-refreshed. [default: 86400]
+      --cache-mb <cache-mb>
+          Maximum size of the pkarr packet cache in megabytes. [default: 100]
+      --query-rate-limit <query-rate-limit>
+          Maximum number of queries per second one IP address can make before it is rate limited. 0 is disabled. [default: 0]
+      --query-rate-limit-burst <query-rate-limit-burst>
+          Short term burst size of the query-rate-limit. 0 is disabled. [default: 0]
+      --dht-rate-limit <dht-rate-limit>
+          Maximum number of queries per second one IP address can make to the DHT before it is rate limited. 0 is disabled. [default: 5]
+      --dht-rate-limit-burst <dht-rate-limit-burst>
+          Short term burst size of the dht-rate-limit. 0 is disabled. [default: 25]
+  -h, --help
+          Print help
+  -V, --version
 ```
 
 For extended logs, see [here](./docs/logging.md).

server/Cargo.toml

@@ -23,4 +23,5 @@ dashmap = "6.1.0"
 
 dyn-clone = "1.0.16"
 thiserror = "1.0.56"
+governor = "0.7.0"
 

server/src/anydns/custom_handler.rs

@@ -2,17 +2,25 @@
 
 use async_trait::async_trait;
 use dyn_clone::DynClone;
-use std::fmt::Debug;
+use std::{fmt::Debug, net::IpAddr};
 
 use super::dns_socket::DnsSocket;
 
+/// Errors that a CustomHandler can return.
 #[derive(thiserror::Error, Debug)]
 pub enum CustomHandlerError {
+    /// Lookup failed. Error will be logged. SRVFAIL will be returned to the user.
     #[error(transparent)]
-    IO(#[from] super::dns_socket::RequestError),
+    Failed(#[from] Box<dyn std::error::Error + Send + Sync>),
 
+    /// Handler does not consider itself responsible for this query.
+    /// Will fallback to ICANN.
     #[error("Query is not processed by handler. Fallback to ICANN.")]
     Unhandled,
+
+    /// Handler rate limited the IP. Will return RCODE::Refused.
+    #[error("Source ip address {0} is rate limited.")]
+    RateLimited(IpAddr)
 }
 
 /**
@@ -25,6 +33,7 @@ pub trait CustomHandler: DynClone + Send + Sync {
         &mut self,
         query: &Vec<u8>,
         socket: DnsSocket,
+        from: Option<IpAddr>
     ) -> Result<Vec<u8>, CustomHandlerError>;
 }
 
@@ -63,8 +72,9 @@ impl HandlerHolder {
         &mut self,
         query: &Vec<u8>,
         socket: DnsSocket,
+        from: Option<IpAddr>
     ) -> Result<Vec<u8>, CustomHandlerError> {
-        self.func.lookup(query, socket).await
+        self.func.lookup(query, socket, from).await
     }
 }
 
@@ -83,6 +93,7 @@ impl CustomHandler for EmptyHandler {
         &mut self,
         _query: &Vec<u8>,
         _socket: DnsSocket,
+        from: Option<IpAddr>
     ) -> Result<Vec<u8>, CustomHandlerError> {
         Err(CustomHandlerError::Unhandled)
     }
@@ -92,7 +103,7 @@ impl CustomHandler for EmptyHandler {
 mod tests {
     use super::super::dns_socket::DnsSocket;
     use async_trait::async_trait;
-    use std::net::SocketAddr;
+    use std::net::{IpAddr, SocketAddr};
 
     use super::{CustomHandler, CustomHandlerError, HandlerHolder};
 
@@ -128,6 +139,7 @@ mod tests {
             &mut self,
             _query: &Vec<u8>,
             _socket: DnsSocket,
+            from: Option<IpAddr>
         ) -> Result<Vec<u8>, CustomHandlerError> {
             println!("value {}", self.value.value);
             Err(CustomHandlerError::Unhandled)
@@ -144,11 +156,13 @@ mod tests {
         let socket = DnsSocket::new(
             "0.0.0.0:18293".parse().unwrap(),
             icann_fallback,
-            holder1.clone()
+            holder1.clone(),
+            None,
+            None
         )
         .await
         .unwrap();
-        let result = cloned.call(&vec![], socket).await;
+        let result = cloned.call(&vec![], socket, None).await;
         assert!(result.is_err());
     }
 }