diff --git a/.cargo/config.toml b/.cargo/config.toml new file mode 100644 index 0000000..529b14d --- /dev/null +++ b/.cargo/config.toml @@ -0,0 +1,8 @@ +[unstable] +registry-auth = true + +[registries.menteeth] +index = "ssh://git@git.shipyard.rs/menteeth/crate-index.git" + +[net] +git-fetch-with-cli = true diff --git a/.envrc b/.envrc index 8f9838c..5066f56 100644 --- a/.envrc +++ b/.envrc @@ -1,3 +1,5 @@ watch_file flake.nix flake.lock dotenv_if_exists +export ANALYTICS_DB="db/analytics.db" +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..1c25a32 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -17,6 +17,17 @@ version = "1.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f26201604c87b1e01bd3d98f8d5d9a8fcbb815e8cedb41ffccbeb4bf593a35fe" +[[package]] +name = "ahash" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2c99f64d1e06488f620f932677e24bc6e2897582980441ae90a671415bd7ec2f" +dependencies = [ + "cfg-if", + "once_cell", + "version_check", +] + [[package]] name = "aho-corasick" version = "1.0.1" @@ -334,6 +345,7 @@ dependencies = [ "axum", "color-eyre", "image", + "locat", "opentelemetry", "opentelemetry-honeycomb", "pretty-hex", @@ -642,6 +654,18 @@ dependencies = [ "once_cell", ] +[[package]] +name = "fallible-iterator" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4443176a9f2c162692bd3d352d745ef9413eec5782a80d8fd6f8a1ac692a07f7" + +[[package]] +name = "fallible-streaming-iterator" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7360491ce676a36bf9bb3c56c1aa791658183a54d2744120f27285738d90465a" + [[package]] name = "fastrand" version = "1.9.0" @@ -892,6 +916,24 @@ version = "0.12.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8a9ee70c43aaf417c914396645a0fa852624801b24ebb7ae78fe8272889ac888" +[[package]] +name = "hashbrown" +version = "0.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "43a3c133739dddd0d2990f9a4bdf8eb4b21ef50e4851ca85ab661199821d510e" +dependencies = [ + "ahash", +] + +[[package]] +name = "hashlink" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0761a1b9491c4f2e3d66aa0f62d0fba0af9a0e2852e4d48ea506632a4b56e6aa" +dependencies = [ + "hashbrown 0.13.2", +] + [[package]] name = "hazy" version = "0.1.1" @@ -1093,7 +1135,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bd070e393353796e801d209ad339e89596eb4c8d430d18ede6a1cced8fafbd99" dependencies = [ "autocfg", - "hashbrown", + "hashbrown 0.12.3", ] [[package]] @@ -1122,6 +1164,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" @@ -1191,12 +1242,33 @@ dependencies = [ "serde_json", ] +[[package]] +name = "libsqlite3-sys" +version = "0.25.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29f835d03d717946d28b1d1ed632eb6f0e24a299388ee623d0c23118d3e8a7fa" +dependencies = [ + "pkg-config", + "vcpkg", +] + [[package]] name = "linux-raw-sys" version = "0.3.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ece97ea872ece730aed82664c424eb4c8291e1ff2480247ccf7409044bc6479f" +[[package]] +name = "locat" +version = "0.4.0" +source = "registry+ssh://git@git.shipyard.rs/menteeth/crate-index.git" +checksum = "8be9ea47c9293870c87ce3820d1abfcea3de226168f2ab01e395ab1995394a83" +dependencies = [ + "maxminddb", + "rusqlite", + "thiserror", +] + [[package]] name = "lock_api" version = "0.4.9" @@ -1223,6 +1295,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" @@ -1803,6 +1887,20 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b833d8d034ea094b1ea68aa6d5c740e0d04bad9d16568d08ba6f76823a114316" +[[package]] +name = "rusqlite" +version = "0.28.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "01e213bc3ecb39ac32e81e51ebe31fd888a940515173e3a18a35f8c6e896422a" +dependencies = [ + "bitflags", + "fallible-iterator", + "fallible-streaming-iterator", + "hashlink", + "libsqlite3-sys", + "smallvec", +] + [[package]] name = "rustc-demangle" version = "0.1.23" diff --git a/Cargo.toml b/Cargo.toml index b9c8e80..191284b 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.4.0", 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/flake.nix b/flake.nix index 3129a67..fcea8e0 100644 --- a/flake.nix +++ b/flake.nix @@ -37,16 +37,18 @@ src = if inShell then null else ./.; buildInputs = - [ rustc - cargo - openssl + [ openssl pkg-config + sqlite ] ++ (if inShell then [ # In 'nix develop', provide some developer tools. rust-analyzer rustfmt + rustup clippy ] else [ + rustc + cargo (import-cargo.builders.importCargo { lockFile = ./Cargo.lock; inherit pkgs; diff --git a/rust-toolchain.toml b/rust-toolchain.toml new file mode 100644 index 0000000..9706c9a --- /dev/null +++ b/rust-toolchain.toml @@ -0,0 +1,2 @@ +[toolchain] +channel = "nightly-2023-04-30" diff --git a/src/main.rs b/src/main.rs index 93e764b..9caeb3f 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,11 +45,23 @@ 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 adb_env_var = "ANALYTICS_DB"; + let adb_path = std::env::var(adb_env_var) + .unwrap_or_else(|_| panic!("${adb_env_var} must be set")); + let state = ServerState { client: Default::default(), + locat: Arc::new(Locat::new(&cdb_path, &adb_path).unwrap()), }; - let app = Router::new().route("/", get(root_get)).with_state(state); + let app = Router::new() + .route("/", get(root_get)) + .route("/analytics", get(analytics_get)) + .with_state(state); let quit_signal = async { _ = tokio::signal::ctrl_c().await; @@ -60,6 +77,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 +95,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 @@ -166,3 +200,13 @@ async fn download_file(client: &reqwest::Client, url: &str) -> color_eyre::Resul .await?; Ok(bytes.to_vec()) } + +async fn analytics_get(State(state): State) -> Response { + let analytics = state.locat.get_analytics().await.unwrap(); + let mut response = String::new(); + use std::fmt::Write; + for (country, count) in analytics { + _ = writeln!(&mut response, "{country}: {count}"); + } + response.into_response() +}