From 1075ab0ed187a10f2eb3faf1247da8661531bbdb Mon Sep 17 00:00:00 2001
From: Shav Kinderlehrer <shav@trinket.icu>
Date: Sun, 7 Apr 2024 14:44:03 -0400
Subject: Add proxy support

---
 src/get.rs  | 17 +++++++++++++++--
 src/main.rs |  5 ++++-
 2 files changed, 19 insertions(+), 3 deletions(-)

(limited to 'src')

diff --git a/src/get.rs b/src/get.rs
index d96e05f..49868a8 100644
--- a/src/get.rs
+++ b/src/get.rs
@@ -93,7 +93,20 @@ pub async fn id(
 
 async fn save_analytics(headers: HeaderMap, item: UrlRow, addr: SocketAddr, state: ServerState) {
     let id = item.id;
-    let ip = addr.ip().to_string();
+    let ip: Option<String> = if state.behind_proxy {
+        match headers.get("x-real-ip") {
+            Some(it) => {
+                if let Ok(i) = it.to_str() {
+                    Some(i.to_string())
+                } else {
+                    None
+                }
+            }
+            None => None,
+        }
+    } else {
+        Some(addr.ip().to_string())
+    };
     let referer = match headers.get("referer") {
         Some(it) => {
             if let Ok(i) = it.to_str() {
@@ -129,7 +142,7 @@ VALUES ($1,$2,$3,$4)
     .await;
 
     if res.is_ok() {
-        log!("Saved analytics for '{id}' from {ip}");
+        log!("Saved analytics for '{id}' from {}", ip.unwrap_or_default());
     }
 }
 
diff --git a/src/main.rs b/src/main.rs
index ac16eb4..ccad3e8 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -23,6 +23,7 @@ pub struct ServerState {
     pub host: String,
     pub sqids: Sqids,
     pub main_page_redirect: Option<Url>,
+    pub behind_proxy: bool,
 }
 
 #[derive(Debug, Clone, sqlx::FromRow, PartialEq, Eq)]
@@ -53,11 +54,13 @@ async fn main() -> eyre::Result<()> {
         .blocklist(["create".to_string()].into())
         .build()?;
     let main_page_redirect = std::env::var("CHELA_MAIN_PAGE_REDIRECT").unwrap_or_default();
+    let behind_proxy = std::env::var("CHELA_BEHIND_PROXY").is_ok();
     let server_state = ServerState {
         db_pool,
         host,
         sqids,
         main_page_redirect: Url::parse(&main_page_redirect).ok(),
+        behind_proxy,
     };
 
     let address = std::env::var("CHELA_LISTEN_ADDRESS").unwrap_or("0.0.0.0".to_string());
@@ -108,7 +111,7 @@ CREATE TABLE IF NOT EXISTS chela.urls (
 CREATE TABLE IF NOT EXISTS chela.tracking (
     timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
     id TEXT NOT NULL,
-    ip TEXT NOT NULL,
+    ip TEXT,
     referrer TEXT,
     user_agent TEXT
 )
-- 
cgit v1.2.3