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: [PATCH] Add proxy support

---
 .dockerignore |  1 +
 Cargo.lock    |  2 +-
 Cargo.toml    |  2 +-
 Dockerfile    |  8 ++++++--
 README.md     |  7 ++++++-
 src/get.rs    | 17 +++++++++++++++--
 src/main.rs   |  5 ++++-
 7 files changed, 34 insertions(+), 8 deletions(-)
 create mode 100644 .dockerignore

diff --git a/.dockerignore b/.dockerignore
new file mode 100644
index 0000000..ea8c4bf
--- /dev/null
+++ b/.dockerignore
@@ -0,0 +1 @@
+/target
diff --git a/Cargo.lock b/Cargo.lock
index 0a4af45..35d1164 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -194,7 +194,7 @@ checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd"
 
 [[package]]
 name = "chela"
-version = "0.2.0"
+version = "0.3.0"
 dependencies = [
  "axum",
  "color-eyre",
diff --git a/Cargo.toml b/Cargo.toml
index 7853f0f..a3733f2 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "chela"
-version = "0.2.0"
+version = "0.3.0"
 edition = "2021"
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
diff --git a/Dockerfile b/Dockerfile
index 9be2e51..1f0feac 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -1,6 +1,10 @@
-FROM rust:1.76.0
+FROM rust:1.76.0 AS builder
 WORKDIR /usr/src/chela
 COPY . .
 RUN cargo build -r
-CMD ["./target/release/chela"]
+
+FROM gcr.io/distroless/cc-debian12
+WORKDIR /usr/src/chela
+COPY --from=builder /usr/src/chela/target/release/chela ./
+CMD ["./chela"]
 EXPOSE 3000
diff --git a/README.md b/README.md
index 48afbad..2f24b88 100644
--- a/README.md
+++ b/README.md
@@ -11,7 +11,7 @@ You can create a redirect by navigating to the `/create` page and filling out th
 ### With Docker
 #### CLI
 ```bash
-docker run -d \
+$ docker run -d \
     -p 3000:3000 \
     -e DATABASE_URL=postgres://chela:password@dbhost/postgres?sslmode=disable \
     -e CHELA_HOST=a.com \
@@ -37,6 +37,7 @@ services:
             - DATABASE_URL=postgres://chela:password@chela-postgres/postgres?sslmode=disable
             - CHELA_HOST=a.com
             - CHELA_MAIN_PAGE_REDIRECT='https://example.com'
+            - CHELA_BEHIND_PROXY=1
         depends_on:
             - chela-postgres
         restart: unless-stopped
@@ -59,6 +60,9 @@ The address that Chela should listen on. Defaults to `0.0.0.0`.
 ##### `CHELA_MAIN_PAGE_REDIRECT`
 A page that Chela will redirect to when `/` is requested instead of replying with the default homepage.
 
+##### `CHELA_BEHIND_PROXY`
+If this variable is set, Chela will use the `X-Real-IP` header as the client IP address rather than the connection address.
+
 ### Manually
 #### Build
 ```bash
@@ -86,6 +90,7 @@ server {
     
     location / {
         proxy_pass http://localhost:3000/;
+        proxy_set_header X-Real-IP $remote_addr;
 
         limit_except GET HEAD {
             auth_basic 'Restricted';
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
 )