add rate limiter

Author: joostjager

src/lib.rs

@@ -94,6 +94,7 @@ pub mod logger;
 mod message_handler;
 pub mod payment;
 mod peer_store;
+mod rate_limiter;
 mod runtime;
 mod static_invoice_store;
 mod tx_broadcaster;

src/rate_limiter.rs

@@ -0,0 +1,34 @@
+use std::collections::HashMap;
+use std::time::{Duration, Instant};
+
+pub(crate) struct RateLimiter {
+	users: HashMap<Vec<u8>, Instant>,
+	min_interval: Duration,
+}
+
+impl RateLimiter {
+	pub(crate) fn new(min_interval: Duration) -> Self {
+		Self { users: HashMap::new(), min_interval }
+	}
+
+	pub(crate) fn allow(&mut self, user_id: &[u8]) -> bool {
+		let now = Instant::now();
+		match self.users.get(user_id) {
+			Some(&last) if now.duration_since(last) < self.min_interval => false,
+			_ => {
+				let new_user = self.users.insert(user_id.to_vec(), now).is_none();
+
+				// Each time a new user is added, we take the opportunity to clean up old rate limits.
+				if new_user {
+					self.garbage_collect();
+				}
+				true
+			},
+		}
+	}
+
+	fn garbage_collect(&mut self) {
+		let now = Instant::now();
+		self.users.retain(|_, &mut last| now.duration_since(last) < self.min_interval);
+	}
+}

src/static_invoice_store.rs

@@ -1,26 +1,41 @@
-use crate::{hex_utils, types::DynStore};
+use crate::{hex_utils, rate_limiter::RateLimiter, types::DynStore};
 use bitcoin::hashes::{sha256, Hash};
 use lightning::{offers::static_invoice::StaticInvoice, util::ser::Writeable};
-use std::sync::Arc;
+use std::{
+	sync::{Arc, Mutex},
+	time::Duration,
+};
 
 pub(crate) struct StaticInvoiceStore {
 	kv_store: Arc<DynStore>,
+	rate_limiter: Mutex<RateLimiter>,
 }
 
 impl StaticInvoiceStore {
 	// Static invoices are stored at "static_invoices/<sha256(recipient_id)>/<invoice_slot>".
 	//
 	// Example:
-	// static_invoices/039058c6f2c0cb492c533b0a4d14ef77cc0f78abccced5287d84a1a2011cfb81/0000
+	// static_invoices/039058c6f2c0cb492c533b0a4d14ef77cc0f78abccced5287d84a1a2011cfb81/001f
 	const PRIMARY_NAMESPACE: &str = "static_invoices";
 
 	pub(crate) fn new(kv_store: Arc<DynStore>) -> Self {
-		Self { kv_store }
+		Self { kv_store, rate_limiter: Mutex::new(RateLimiter::new(Duration::from_millis(100))) }
+	}
+
+	fn check_rate_limit(&self, recipient_id: &[u8]) -> Result<(), lightning::io::Error> {
+		let mut limiter = self.rate_limiter.lock().unwrap();
+		if !limiter.allow(recipient_id) {
+			Err(lightning::io::Error::new(lightning::io::ErrorKind::Other, "Rate limit exceeded"))
+		} else {
+			Ok(())
+		}
 	}
 
 	pub(crate) async fn handle_static_invoice_requested(
 		&self, recipient_id: Vec<u8>, invoice_slot: u16,
 	) -> Result<Option<StaticInvoice>, lightning::io::Error> {
+		self.check_rate_limit(&recipient_id)?;
+
 		let (secondary_namespace, key) = Self::get_storage_location(invoice_slot, recipient_id);
 
 		self.kv_store.read(Self::PRIMARY_NAMESPACE, &secondary_namespace, &key).and_then(|data| {
@@ -36,6 +51,8 @@ impl StaticInvoiceStore {
 	pub(crate) async fn handle_persist_static_invoice(
 		&self, invoice: StaticInvoice, invoice_slot: u16, recipient_id: Vec<u8>,
 	) -> Result<(), lightning::io::Error> {
+		self.check_rate_limit(&recipient_id)?;
+
 		let (secondary_namespace, key) = Self::get_storage_location(invoice_slot, recipient_id);
 
 		let mut buf = Vec::new();