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, } #[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 { + let header = headers.get("x-remote-ip")?; + let header = header.to_str().ok()?; + let addr = header.parse::().ok()?; + Some(addr) +} + async fn root_get(headers: HeaderMap, State(state): State) -> Response { let tracer = global::tracer(""); let mut span = tracer.start("root_get"); @@ -71,6 +87,16 @@ async fn root_get(headers: HeaderMap, State(state): State) -> 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