From 9f681ff4ac939329f39d508617d5b1e45e7f737f Mon Sep 17 00:00:00 2001
From: Lyle Mantooth <lyle@menteeth.us>
Date: Sat, 20 May 2023 12:29:31 -0400
Subject: [PATCH] Use locat for Geo-IP locating.

---
 .envrc      |  1 +
 .gitignore  |  3 ++-
 Cargo.lock  | 32 ++++++++++++++++++++++++++++++++
 Cargo.toml  |  1 +
 src/main.rs | 28 +++++++++++++++++++++++++++-
 5 files changed, 63 insertions(+), 2 deletions(-)

diff --git a/.envrc b/.envrc
index 8f9838c..ffccafd 100644
--- a/.envrc
+++ b/.envrc
@@ -1,3 +1,4 @@
 watch_file flake.nix flake.lock
 dotenv_if_exists
+export GEOLITE2_COUNTRY_DB="db/ip_country_sample.mmdb"
 use flake
diff --git a/.gitignore b/.gitignore
index c5dd462..4c0b809 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,2 +1,3 @@
-target
+/db
+/target
 .env
diff --git a/Cargo.lock b/Cargo.lock
index bc451b0..b626896 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -334,6 +334,7 @@ dependencies = [
  "axum",
  "color-eyre",
  "image",
+ "locat",
  "opentelemetry",
  "opentelemetry-honeycomb",
  "pretty-hex",
@@ -1122,6 +1123,15 @@ version = "2.7.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "12b6ee2129af8d4fb011108c73d99a1b83a85977f23b82460c0ae2e25bb4b57f"
 
+[[package]]
+name = "ipnetwork"
+version = "0.18.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4088d739b183546b239688ddbc79891831df421773df95e236daf7867866d355"
+dependencies = [
+ "serde",
+]
+
 [[package]]
 name = "itoa"
 version = "1.0.6"
@@ -1197,6 +1207,16 @@ version = "0.3.7"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "ece97ea872ece730aed82664c424eb4c8291e1ff2480247ccf7409044bc6479f"
 
+[[package]]
+name = "locat"
+version = "0.3.1"
+source = "registry+ssh://git@git.shipyard.rs/menteeth/crate-index.git"
+checksum = "5710a18ee2dbe5aff9a408bc48e25a9c53116b2e3a0a334b23388437806c0884"
+dependencies = [
+ "maxminddb",
+ "thiserror",
+]
+
 [[package]]
 name = "lock_api"
 version = "0.4.9"
@@ -1223,6 +1243,18 @@ version = "0.7.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "b87248edafb776e59e6ee64a79086f65890d3510f2c656c000bf2a7e8a0aea40"
 
+[[package]]
+name = "maxminddb"
+version = "0.23.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fe2ba61113f9f7a9f0e87c519682d39c43a6f3f79c2cc42c3ba3dda83b1fa334"
+dependencies = [
+ "ipnetwork",
+ "log",
+ "memchr",
+ "serde",
+]
+
 [[package]]
 name = "memchr"
 version = "2.5.0"
diff --git a/Cargo.toml b/Cargo.toml
index b9c8e80..9d8f4c8 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -9,6 +9,7 @@ artem = { version = "=1.1.5", default-features = false }
 axum = "0.6"
 color-eyre = "0.6"
 image = "0.24"
+locat = { version = "0.3.1", registry = "menteeth" }
 opentelemetry = { version = "0.18", features = ["rt-tokio"] }
 opentelemetry-honeycomb = { git = "https://github.com/fasterthanlime/opentelemetry-honeycomb-rs", branch = "simplified", version = "0.1.0" }
 pretty-hex = "0.3"
diff --git a/src/main.rs b/src/main.rs
index 93e764b..e3713eb 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,4 +1,6 @@
+use std::net::IpAddr;
 use std::str::FromStr;
+use std::sync::Arc;
 
 use axum::{
     body::BoxBody,
@@ -15,12 +17,15 @@ use opentelemetry::{
 };
 use reqwest::StatusCode;
 use serde::Deserialize;
-use tracing::{info, Level};
+use tracing::{info, warn, Level};
 use tracing_subscriber::{filter::Targets, layer::SubscriberExt, util::SubscriberInitExt};
 
+use locat::Locat;
+
 #[derive(Clone)]
 struct ServerState {
     client: reqwest::Client,
+    locat: Arc<Locat>,
 }
 
 #[tokio::main]
@@ -40,8 +45,12 @@ async fn main() {
         .with(filter)
         .init();
 
+    let cdb_env_var = "GEOLITE2_COUNTRY_DB";
+    let cdb_path = std::env::var(cdb_env_var)
+        .unwrap_or_else(|_| panic!("${cdb_env_var} must be set"));
     let state = ServerState {
         client: Default::default(),
+        locat: Arc::new(Locat::new(&cdb_path, "todo.db").unwrap()),
     };
 
     let app = Router::new().route("/", get(root_get)).with_state(state);
@@ -60,6 +69,13 @@ async fn main() {
         .unwrap();
 }
 
+fn get_client_addr(headers: &HeaderMap) -> Option<IpAddr> {
+    let header = headers.get("x-remote-ip")?;
+    let header = header.to_str().ok()?;
+    let addr = header.parse::<IpAddr>().ok()?;
+    Some(addr)
+}
+
 async fn root_get(headers: HeaderMap, State(state): State<ServerState>) -> Response<BoxBody> {
     let tracer = global::tracer("");
     let mut span = tracer.start("root_get");
@@ -71,6 +87,16 @@ async fn root_get(headers: HeaderMap, State(state): State<ServerState>) -> Respo
             .unwrap_or_default(),
     ));
 
+    if let Some(addr) = get_client_addr(&headers) {
+        match state.locat.ip_to_iso_code(addr).await {
+            Some(country) => {
+                info!("Got request from {country}");
+                span.set_attribute(KeyValue::new("country", country.to_string()));
+            }
+            None => warn!("Could not determine country for IP address"),
+        }
+    }
+
     root_get_inner(state)
         .with_context(Context::current_with_span(span))
         .await