feat: add configurable resource limits, fix rate limiter

Signed-off-by: Henry <mail@henrygressmann.de>

Author: explodingcamera

CHANGELOG.md

@@ -22,7 +22,7 @@ Since this is not a library, this changelog focuses on the changes that are rele
 
 ### ⚠️ Breaking Changes
 
-- Proxy-provided client IP headers are no longer trusted by default. Configure `client_ip_headers` to opt in
+- Proxy-provided client IP headers are no longer trusted by default. Configure `client_ip_headers` and `trusted_proxies` to opt in
 
 ### Features
 
@@ -31,6 +31,13 @@ Since this is not a library, this changelog focuses on the changes that are rele
 - Tracker requests now avoid CORS preflight requests by sending JSON as `text/plain` and the event API accepts tracker JSON regardless of content type
 - Added external authentication with OpenID Connect, Google, and Microsoft Entra ID, including optional user creation and provider access restrictions
 
+### Security
+
+- Added input, result, concurrency, and execution-time limits to dashboard reports
+- Dashboard report limits can be adjusted in the `[limits]` configuration section
+- Rate limits now use configured client IP headers when requests come from a trusted reverse proxy
+- Onboarding tokens are now validated in constant time and consumed atomically when creating the initial administrator
+
 ### Other
 
 - Renamed `trusted_headers` to `client_ip_headers` (the old name remains supported)

Cargo.toml

@@ -106,6 +106,9 @@ tower = { version = "*", features = ["util"] }
 default = ["geoip"]
 geoip = ["dep:maxminddb"]
 
+[lints.clippy]
+too_many_arguments = "allow"
+
 [profile.dev]
 opt-level = 1
 

benches/reports.rs

@@ -30,7 +30,7 @@ fn benchmark_reports(c: &mut Criterion) {
         start: Utc::now().checked_sub_days(Days::new(365)).expect("failed to build range start"),
         end: Utc::now(),
     };
-    let day_buckets = reports::build_graph_buckets(&range, GraphInterval::Day, Some("UTC"))
+    let day_buckets = reports::build_graph_buckets(&range, GraphInterval::Day, Some("UTC"), 2000)
         .expect("failed to build day buckets for benchmark");
 
     let conn = app.events_conn().expect("failed to get events connection");
@@ -80,7 +80,7 @@ fn benchmark_reports(c: &mut Criterion) {
                 &metric,
                 |b, metric| {
                     b.iter(|| {
-                        reports::dimension_report(&conn, &entities, "pageview", &range, &dimension, &[], metric)
+                        reports::dimension_report(&conn, &entities, "pageview", &range, &dimension, &[], metric, 1000)
                             .expect("dimension_report failed")
                     });
                 },

data/config.example.toml

@@ -17,7 +17,7 @@ listen = 9042  # The port to listen on (http)
 # client_ip_headers=["cloudflare", "fly", "X-Forwarded-For", "X-Client-IP"]
 
 # # Trusted reverse proxies (IP or CIDR).
-# # Empty by default, meaning all proxies are trusted.
+# # Empty by default, meaning forwarded client IP and GeoIP headers are not trusted.
 # trusted_proxies=["127.0.0.1", "10.0.0.0/8"]
 
 # # Hour of day, in the server's local timezone, when visitor grouping salts rotate.
@@ -42,7 +42,18 @@ listen = 9042  # The port to listen on (http)
 # # Otherwise, the database will be downloaded automatically and stored in the data_dir
 # maxmind_db_path="./GeoLite2-City.mmdb"
 
+# Database settings
 [duckdb]
 # # See https://liwan.dev/guides/duckdb for guidance
 # threads=2
 # memory_limit=2G
+
+# Resource limits for dashboard reports.
+[limits]
+# report_timeout_seconds=30
+# report_max_concurrency=8
+# report_max_range_days=3660
+# report_max_dimension_results=1000
+# report_max_datapoints=2000
+# report_max_filters=20
+# report_max_filter_value_bytes=2048

src/app/core/onboarding.rs

@@ -1,12 +1,14 @@
 use anyhow::Result;
-use arc_swap::ArcSwapOption;
-use std::sync::Arc;
+use std::sync::{Arc, Mutex};
 
-use crate::{app::SqlitePool, utils::hash::onboarding_token};
+use crate::{
+    app::SqlitePool,
+    utils::hash::{onboarding_token, onboarding_token_matches},
+};
 
 #[derive(Clone)]
 pub struct LiwanOnboarding {
-    token: Arc<ArcSwapOption<String>>,
+    token: Arc<Mutex<Option<String>>>,
 }
 
 impl LiwanOnboarding {
@@ -15,20 +17,69 @@ impl LiwanOnboarding {
             tracing::debug!("Checking if an onboarding token needs to be generated");
             let conn = pool.get()?;
             let onboarded = conn.prepare("select 1 from users limit 1")?.exists([])?;
-            ArcSwapOption::new((!onboarded).then(|| onboarding_token().into()))
+            (!onboarded).then(onboarding_token)
         };
 
-        Ok(Self { token: onboarding.into() })
+        Ok(Self { token: Arc::new(Mutex::new(onboarding)) })
     }
 
     /// Get the onboarding token, if it exists
-    pub fn token(&self) -> Result<Option<String>> {
-        Ok((self.token.load()).as_ref().map(|v| (**v).clone()))
+    pub fn token(&self) -> Option<String> {
+        self.token.lock().unwrap_or_else(|poisoned| poisoned.into_inner()).clone()
     }
 
     /// Clear the onboarding token to prevent it from being used again
-    pub fn clear(&self) -> Result<()> {
-        self.token.store(None);
-        Ok(())
+    pub fn clear(&self) {
+        *self.token.lock().unwrap_or_else(|poisoned| poisoned.into_inner()) = None;
+    }
+
+    /// Complete setup once when the supplied token is valid.
+    pub fn complete_setup(&self, supplied_token: &str, setup: impl FnOnce() -> Result<()>) -> Result<bool> {
+        let mut token = self.token.lock().unwrap_or_else(|poisoned| poisoned.into_inner());
+        if !token.as_deref().is_some_and(|token| onboarding_token_matches(token, supplied_token)) {
+            return Ok(false);
+        }
+
+        setup()?;
+        *token = None;
+        Ok(true)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::sync::atomic::{AtomicUsize, Ordering};
+
+    #[test]
+    fn complete_allows_only_one_setup() {
+        let onboarding = LiwanOnboarding { token: Arc::new(Mutex::new(Some("token".to_string()))) };
+        let setup_count = Arc::new(AtomicUsize::new(0));
+        let attempts = (0..2)
+            .map(|_| {
+                let onboarding = onboarding.clone();
+                let setup_count = setup_count.clone();
+                std::thread::spawn(move || {
+                    onboarding
+                        .complete_setup("token", || {
+                            setup_count.fetch_add(1, Ordering::Relaxed);
+                            Ok(())
+                        })
+                        .unwrap()
+                })
+            })
+            .map(|attempt| attempt.join().unwrap())
+            .collect::<Vec<_>>();
+
+        assert_eq!(attempts.iter().filter(|completed| **completed).count(), 1);
+        assert_eq!(setup_count.load(Ordering::Relaxed), 1);
+    }
+
+    #[test]
+    fn complete_keeps_token_when_setup_fails() {
+        let onboarding = LiwanOnboarding { token: Arc::new(Mutex::new(Some("token".to_string()))) };
+
+        assert!(onboarding.complete_setup("token", || anyhow::bail!("failed")).is_err());
+        assert_eq!(onboarding.token().as_deref(), Some("token"));
     }
 }

src/app/core/reports/dimension.rs

@@ -17,6 +17,7 @@ pub fn dimension_report(
     dimension: &Dimension,
     filters: &[DimensionFilter],
     metric: &Metric,
+    max_results: usize,
 ) -> Result<ReportTable> {
     if entities.is_empty() {
         return Ok(BTreeMap::new());
@@ -87,7 +88,8 @@ pub fn dimension_report(
 			{metric_column} as metric_value
 		from session_data sd
 		group by dimension_value
-		order by metric_value desc;
+		order by metric_value desc
+		limit {max_results};
 	"
     );
 
@@ -112,3 +114,41 @@ pub fn dimension_report(
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::app::Liwan;
+    use crate::config::Config;
+    use chrono::{Duration, TimeZone, Utc};
+
+    #[test]
+    fn dimension_report_limits_results() {
+        let app = Liwan::new_memory(Config::default()).expect("failed to create app");
+        let max_results = app.config.limits.report_max_dimension_results;
+        let conn = app.events_conn().expect("failed to get events conn");
+        conn.execute_batch(&format!(
+            "insert into events (entity_id, visitor_group_id, event, created_at, fqdn, path)
+             select 'entity-1', 'visitor-' || i, 'pageview', '2024-01-01'::timestamp, 'example.com', '/' || i
+             from range({}) values(i)",
+            max_results + 1
+        ))
+        .expect("failed to insert events");
+        let start = Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0).unwrap();
+        let range = DateRange { start, end: start + Duration::days(1) };
+
+        let report = dimension_report(
+            &conn,
+            &["entity-1".to_string()],
+            "pageview",
+            &range,
+            &Dimension::Path,
+            &[],
+            &Metric::Views,
+            max_results,
+        )
+        .expect("failed to build dimension report");
+
+        assert_eq!(report.len(), max_results);
+    }
+}

src/app/core/reports/graph.rs

@@ -52,6 +52,7 @@ pub fn build_graph_buckets(
     range: &DateRange,
     interval: GraphInterval,
     timezone: Option<&str>,
+    max_datapoints: usize,
 ) -> Result<Vec<DateRange>> {
     if range.start >= range.end {
         return Ok(Vec::new());
@@ -71,9 +72,31 @@ pub fn build_graph_buckets(
         GraphInterval::Day => resolve_local_day_start(timezone, range.start.with_timezone(&timezone).date_naive())?,
     };
 
-    let mut buckets = Vec::new();
+    let mut bucket_count = 0;
     let mut bucket_start = aligned_start;
 
+    while bucket_start < range.end {
+        let next_bucket_start = match interval {
+            GraphInterval::Hour => bucket_start + Duration::hours(1),
+            GraphInterval::Day => {
+                let next_date = bucket_start
+                    .with_timezone(&timezone)
+                    .date_naive()
+                    .checked_add_days(Days::new(1))
+                    .context("Failed to advance bucket date")?;
+                resolve_local_day_start(timezone, next_date)?
+            }
+        };
+
+        bucket_count += 1;
+        if bucket_count > max_datapoints {
+            anyhow::bail!("Too many data points");
+        }
+        bucket_start = next_bucket_start;
+    }
+
+    let mut buckets = Vec::with_capacity(bucket_count);
+    bucket_start = aligned_start;
     while bucket_start < range.end {
         let next_bucket_start = match interval {
             GraphInterval::Hour => bucket_start + Duration::hours(1),
@@ -248,7 +271,7 @@ mod tests {
         };
 
         let buckets =
-            build_graph_buckets(&range, GraphInterval::Hour, Some(timezone)).expect("failed to build buckets");
+            build_graph_buckets(&range, GraphInterval::Hour, Some(timezone), 2000).expect("failed to build buckets");
         let local_starts = buckets
             .iter()
             .map(|bucket| bucket.start.with_timezone(&Tz::Asia__Kolkata).format("%Y-%m-%d %H:%M").to_string())
@@ -266,7 +289,8 @@ mod tests {
             end: local_datetime(Tz::America__New_York, 2024, 1, 3, 12, 0),
         };
 
-        let buckets = build_graph_buckets(&range, GraphInterval::Day, Some(timezone)).expect("failed to build buckets");
+        let buckets =
+            build_graph_buckets(&range, GraphInterval::Day, Some(timezone), 2000).expect("failed to build buckets");
         let local_starts = buckets
             .iter()
             .map(|bucket| bucket.start.with_timezone(&Tz::America__New_York).format("%Y-%m-%d %H:%M").to_string())
@@ -299,7 +323,8 @@ mod tests {
             .expect("failed to append events");
 
         let range = DateRange { start, end };
-        let buckets = build_graph_buckets(&range, GraphInterval::Day, Some("UTC")).expect("failed to build buckets");
+        let buckets =
+            build_graph_buckets(&range, GraphInterval::Day, Some("UTC"), 2000).expect("failed to build buckets");
         let conn = app.events_conn().expect("failed to get events conn");
         let report =
             overall_report(&conn, &["entity-1".to_string()], "pageview", &range, &buckets, &[], &Metric::Views)
@@ -308,4 +333,13 @@ mod tests {
         let values = report.iter().map(|point| point.value).collect::<Vec<_>>();
         assert_eq!(values, vec![1.0, 1.0]);
     }
+
+    #[test]
+    fn build_graph_buckets_rejects_excessive_ranges() {
+        let start = local_datetime(Tz::UTC, 2024, 1, 1, 0, 0);
+        let max_datapoints = 2000;
+        let range = DateRange { start, end: start + Duration::hours(max_datapoints as i64 + 1) };
+
+        assert!(build_graph_buckets(&range, GraphInterval::Hour, Some("UTC"), max_datapoints).is_err());
+    }
 }