Compare commits

...

6 commits

8 changed files with 310 additions and 33 deletions

1
.gitignore vendored
View file

@ -6,3 +6,4 @@ perf.data*
dhat.out.* dhat.out.*
settings.yaml settings.yaml
tarpaulin-report.html tarpaulin-report.html
GeoLite2-Country.mmdb

132
Cargo.lock generated
View file

@ -1,7 +1,5 @@
# This file is automatically @generated by Cargo. # This file is automatically @generated by Cargo.
# It is not intended for manual editing. # It is not intended for manual editing.
version = 3
[[package]] [[package]]
name = "actix-codec" name = "actix-codec"
version = "0.4.0" version = "0.4.0"
@ -9,7 +7,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1d5dbeb2d9e51344cb83ca7cc170f1217f9fe25bfc50160e6e200b5c31c1019a" checksum = "1d5dbeb2d9e51344cb83ca7cc170f1217f9fe25bfc50160e6e200b5c31c1019a"
dependencies = [ dependencies = [
"bitflags", "bitflags",
"bytes", "bytes 1.0.1",
"futures-core", "futures-core",
"futures-sink", "futures-sink",
"log", "log",
@ -33,7 +31,7 @@ dependencies = [
"base64", "base64",
"bitflags", "bitflags",
"brotli2", "brotli2",
"bytes", "bytes 1.0.1",
"bytestring", "bytestring",
"derive_more", "derive_more",
"encoding_rs", "encoding_rs",
@ -170,7 +168,7 @@ dependencies = [
"actix-utils", "actix-utils",
"actix-web-codegen", "actix-web-codegen",
"ahash 0.7.4", "ahash 0.7.4",
"bytes", "bytes 1.0.1",
"cfg-if", "cfg-if",
"cookie", "cookie",
"derive_more", "derive_more",
@ -388,6 +386,16 @@ version = "1.4.3"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "14c189c53d098945499cdfa7ecc63567cf3886b3332b312a5b4585d8d3a6a610" checksum = "14c189c53d098945499cdfa7ecc63567cf3886b3332b312a5b4585d8d3a6a610"
[[package]]
name = "bytes"
version = "0.4.12"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "206fdffcfa2df7cbe15601ef46c813fce0965eb3286db6b56c583b814b51c81c"
dependencies = [
"byteorder",
"iovec",
]
[[package]] [[package]]
name = "bytes" name = "bytes"
version = "1.0.1" version = "1.0.1"
@ -400,14 +408,14 @@ version = "1.0.0"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "90706ba19e97b90786e19dc0d5e2abd80008d99d4c0c5d1ad0b5e72cec7c494d" checksum = "90706ba19e97b90786e19dc0d5e2abd80008d99d4c0c5d1ad0b5e72cec7c494d"
dependencies = [ dependencies = [
"bytes", "bytes 1.0.1",
] ]
[[package]] [[package]]
name = "cc" name = "cc"
version = "1.0.68" version = "1.0.69"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "4a72c244c1ff497a746a7e1fb3d14bd08420ecda70c8f25c7112f2781652d787" checksum = "e70cc2f62c6ce1868963827bd677764c62d07c3d9a3e1fb1177ee1a9ab199eb2"
dependencies = [ dependencies = [
"jobserver", "jobserver",
] ]
@ -499,9 +507,9 @@ checksum = "6245d59a3e82a7fc217c5828a6692dbc6dfb63a0c8c90495621f7b9d79704a0e"
[[package]] [[package]]
name = "cookie" name = "cookie"
version = "0.15.0" version = "0.15.1"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ffdf8865bac3d9a3bde5bde9088ca431b11f5d37c7a578b8086af77248b76627" checksum = "d5f1c7727e460397e56abc4bddc1d49e07a1ad78fc98eb2e1c8f032a58a2f80d"
dependencies = [ dependencies = [
"percent-encoding", "percent-encoding",
"time 0.2.27", "time 0.2.27",
@ -674,6 +682,18 @@ dependencies = [
"syn", "syn",
] ]
[[package]]
name = "filetime"
version = "0.2.14"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1d34cfa13a63ae058bfa601fe9e313bbdb3746427c1459185464ce0fcf62e1e8"
dependencies = [
"cfg-if",
"libc",
"redox_syscall",
"winapi",
]
[[package]] [[package]]
name = "flate2" name = "flate2"
version = "1.0.20" version = "1.0.20"
@ -682,8 +702,10 @@ checksum = "cd3aec53de10fe96d7d8c565eb17f2c687bb5518a2ec453b5b1252964526abe0"
dependencies = [ dependencies = [
"cfg-if", "cfg-if",
"crc32fast", "crc32fast",
"futures 0.1.31",
"libc", "libc",
"miniz_oxide", "miniz_oxide",
"tokio-io",
] ]
[[package]] [[package]]
@ -708,6 +730,12 @@ version = "1.1.0"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "fed34cd105917e91daa4da6b3728c47b068749d6a62c59811f06ed2ac71d9da7" checksum = "fed34cd105917e91daa4da6b3728c47b068749d6a62c59811f06ed2ac71d9da7"
[[package]]
name = "futures"
version = "0.1.31"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3a471a38ef8ed83cd6e40aa59c1ffe17db6855c18e3604d9c4ed8c08ebc28678"
[[package]] [[package]]
name = "futures" name = "futures"
version = "0.3.15" version = "0.3.15"
@ -854,7 +882,7 @@ version = "0.3.3"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "825343c4eef0b63f541f8903f395dc5beb362a979b5799a84062527ef1e37726" checksum = "825343c4eef0b63f541f8903f395dc5beb362a979b5799a84062527ef1e37726"
dependencies = [ dependencies = [
"bytes", "bytes 1.0.1",
"fnv", "fnv",
"futures-core", "futures-core",
"futures-sink", "futures-sink",
@ -924,7 +952,7 @@ version = "0.2.4"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "527e8c9ac747e28542699a951517aa9a6945af506cd1f2e1b53a576c17b6cc11" checksum = "527e8c9ac747e28542699a951517aa9a6945af506cd1f2e1b53a576c17b6cc11"
dependencies = [ dependencies = [
"bytes", "bytes 1.0.1",
"fnv", "fnv",
"itoa", "itoa",
] ]
@ -935,7 +963,7 @@ version = "0.4.2"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "60daa14be0e0786db0f03a9e57cb404c9d756eed2b6c62b9ea98ec5743ec75a9" checksum = "60daa14be0e0786db0f03a9e57cb404c9d756eed2b6c62b9ea98ec5743ec75a9"
dependencies = [ dependencies = [
"bytes", "bytes 1.0.1",
"http", "http",
"pin-project-lite", "pin-project-lite",
] ]
@ -958,7 +986,7 @@ version = "0.14.10"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7728a72c4c7d72665fde02204bcbd93b247721025b222ef78606f14513e0fd03" checksum = "7728a72c4c7d72665fde02204bcbd93b247721025b222ef78606f14513e0fd03"
dependencies = [ dependencies = [
"bytes", "bytes 1.0.1",
"futures-channel", "futures-channel",
"futures-core", "futures-core",
"futures-util", "futures-util",
@ -1014,13 +1042,22 @@ dependencies = [
[[package]] [[package]]
name = "instant" name = "instant"
version = "0.1.9" version = "0.1.10"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "61124eeebbd69b8190558df225adf7e4caafce0d743919e5d6b19652314ec5ec" checksum = "bee0328b1209d157ef001c94dd85b4f8f64139adb0eac2659f4b08382b2f474d"
dependencies = [ dependencies = [
"cfg-if", "cfg-if",
] ]
[[package]]
name = "iovec"
version = "0.1.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b2b3ea6ff95e175473f8ffe6a7eb7c00d054240321b84c57051175fe3c1e075e"
dependencies = [
"libc",
]
[[package]] [[package]]
name = "ipnet" name = "ipnet"
version = "2.3.1" version = "2.3.1"
@ -1196,16 +1233,18 @@ dependencies = [
"async-trait", "async-trait",
"base64", "base64",
"bincode", "bincode",
"bytes", "bytes 1.0.1",
"chacha20", "chacha20",
"chrono", "chrono",
"clap", "clap",
"ctrlc", "ctrlc",
"dotenv", "dotenv",
"futures", "flate2",
"futures 0.3.15",
"lfu_cache", "lfu_cache",
"log", "log",
"lru", "lru",
"maxminddb",
"md-5", "md-5",
"once_cell", "once_cell",
"parking_lot", "parking_lot",
@ -1218,6 +1257,7 @@ dependencies = [
"serde_yaml", "serde_yaml",
"sodiumoxide", "sodiumoxide",
"sqlx", "sqlx",
"tar",
"tempfile", "tempfile",
"thiserror", "thiserror",
"tokio", "tokio",
@ -1251,6 +1291,17 @@ version = "0.1.8"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7ffc5c5338469d4d3ea17d269fa8ea3512ad247247c30bd2df69e68309ed0a08" checksum = "7ffc5c5338469d4d3ea17d269fa8ea3512ad247247c30bd2df69e68309ed0a08"
[[package]]
name = "maxminddb"
version = "0.20.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "835288c768ea07fedcccf84f1c9a463fe7e36856b5868d5ea050f858e20b01b8"
dependencies = [
"log",
"memchr",
"serde",
]
[[package]] [[package]]
name = "md-5" name = "md-5"
version = "0.9.1" version = "0.9.1"
@ -1665,7 +1716,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "246e9f61b9bb77df069a947682be06e31ac43ea37862e244a69f177694ea6d22" checksum = "246e9f61b9bb77df069a947682be06e31ac43ea37862e244a69f177694ea6d22"
dependencies = [ dependencies = [
"base64", "base64",
"bytes", "bytes 1.0.1",
"encoding_rs", "encoding_rs",
"futures-core", "futures-core",
"futures-util", "futures-util",
@ -2021,7 +2072,7 @@ dependencies = [
"atoi", "atoi",
"bitflags", "bitflags",
"byteorder", "byteorder",
"bytes", "bytes 1.0.1",
"chrono", "chrono",
"crc", "crc",
"crossbeam-channel", "crossbeam-channel",
@ -2065,7 +2116,7 @@ checksum = "47e4a2349d1ffd60a03ca0de3f116ba55d7f406e55a0d84c64a5590866d94c06"
dependencies = [ dependencies = [
"dotenv", "dotenv",
"either", "either",
"futures", "futures 0.3.15",
"heck", "heck",
"hex", "hex",
"once_cell", "once_cell",
@ -2204,6 +2255,17 @@ version = "1.0.1"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "55937e1799185b12863d447f42597ed69d9928686b8d88a1df17376a097d8369" checksum = "55937e1799185b12863d447f42597ed69d9928686b8d88a1df17376a097d8369"
[[package]]
name = "tar"
version = "0.4.35"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7d779dc6aeff029314570f666ec83f19df7280bb36ef338442cfa8c604021b80"
dependencies = [
"filetime",
"libc",
"xattr",
]
[[package]] [[package]]
name = "tempfile" name = "tempfile"
version = "3.2.0" version = "3.2.0"
@ -2346,7 +2408,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "98c8b05dc14c75ea83d63dd391100353789f5f24b8b3866542a5e85c8be8e985" checksum = "98c8b05dc14c75ea83d63dd391100353789f5f24b8b3866542a5e85c8be8e985"
dependencies = [ dependencies = [
"autocfg", "autocfg",
"bytes", "bytes 1.0.1",
"libc", "libc",
"memchr", "memchr",
"mio", "mio",
@ -2359,6 +2421,17 @@ dependencies = [
"winapi", "winapi",
] ]
[[package]]
name = "tokio-io"
version = "0.1.13"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "57fc868aae093479e3131e3d165c93b1c7474109d13c90ec0dda2a1bbfff0674"
dependencies = [
"bytes 0.4.12",
"futures 0.1.31",
"log",
]
[[package]] [[package]]
name = "tokio-macros" name = "tokio-macros"
version = "1.3.0" version = "1.3.0"
@ -2399,7 +2472,7 @@ version = "0.6.7"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1caa0b0c8d94a049db56b5acf8cba99dc0623aab1b26d5b5f5e2d945846b3592" checksum = "1caa0b0c8d94a049db56b5acf8cba99dc0623aab1b26d5b5f5e2d945846b3592"
dependencies = [ dependencies = [
"bytes", "bytes 1.0.1",
"futures-core", "futures-core",
"futures-sink", "futures-sink",
"log", "log",
@ -2592,9 +2665,9 @@ checksum = "f1bddf1187be692e79c5ffeab891132dfb0f236ed36a43c7ed39f1165ee20191"
[[package]] [[package]]
name = "vergen" name = "vergen"
version = "5.1.12" version = "5.1.13"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5f710158a7c9449f400cf734e97cb417a09a7b5f0cf54b133718b6b2951c9f79" checksum = "542f37b4798c879409865dde7908e746d836f77839c3a6bea5c8b4e4dcf6620b"
dependencies = [ dependencies = [
"anyhow", "anyhow",
"cfg-if", "cfg-if",
@ -2794,6 +2867,15 @@ version = "0.2.0"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "85e60b0d1b5f99db2556934e21937020776a5d31520bf169e851ac44e6420214" checksum = "85e60b0d1b5f99db2556934e21937020776a5d31520bf169e851ac44e6420214"
[[package]]
name = "xattr"
version = "0.2.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "244c3741f4240ef46274860397c7c74e50eb23624996930e484c16679633a54c"
dependencies = [
"libc",
]
[[package]] [[package]]
name = "yaml-rust" name = "yaml-rust"
version = "0.4.5" version = "0.4.5"

View file

@ -25,11 +25,13 @@ chrono = { version = "0.4", features = [ "serde" ] }
clap = { version = "3.0.0-beta.2", features = [ "wrap_help" ] } clap = { version = "3.0.0-beta.2", features = [ "wrap_help" ] }
ctrlc = "3" ctrlc = "3"
dotenv = "0.15" dotenv = "0.15"
flate2 = { version = "1", features = [ "tokio" ] }
futures = "0.3" futures = "0.3"
once_cell = "1" once_cell = "1"
log = { version = "0.4", features = [ "serde" ] } log = { version = "0.4", features = [ "serde" ] }
lfu_cache = "1" lfu_cache = "1"
lru = "0.6" lru = "0.6"
maxminddb = "0.20"
md-5 = "0.9" md-5 = "0.9"
parking_lot = "0.11" parking_lot = "0.11"
prometheus = { version = "0.12", features = [ "process" ] } prometheus = { version = "0.12", features = [ "process" ] }
@ -41,6 +43,7 @@ serde_repr = "0.1"
serde_yaml = "0.8" serde_yaml = "0.8"
sodiumoxide = "0.2" sodiumoxide = "0.2"
sqlx = { version = "0.5", features = [ "runtime-actix-rustls", "sqlite", "time", "chrono", "macros", "offline" ] } sqlx = { version = "0.5", features = [ "runtime-actix-rustls", "sqlite", "time", "chrono", "macros", "offline" ] }
tar = "0.4"
thiserror = "1" thiserror = "1"
tokio = { version = "1", features = [ "rt-multi-thread", "macros", "fs", "sync", "parking_lot" ] } tokio = { version = "1", features = [ "rt-multi-thread", "macros", "fs", "sync", "parking_lot" ] }
tokio-stream = { version = "0.1", features = [ "sync" ] } tokio-stream = { version = "0.1", features = [ "sync" ] }

View file

@ -58,3 +58,25 @@ Note that the client secret (`CLIENT_SECRET`) is the only configuration option
that can only can be provided from the environment, an `.env` file, or the that can only can be provided from the environment, an `.env` file, or the
`settings.yaml` file. In other words, you _cannot_ provide this value from the `settings.yaml` file. In other words, you _cannot_ provide this value from the
command line. command line.
## Special thanks
This project could not have been completed without the assistance of the
following:
#### Development Assistance (Alphabetical Order)
- carbotaniuman#6974
- LFlair#1337
- Plykiya#1738
- Tristan 9#6752
- The Rust Discord community
#### Beta testers
- NigelVH#7162
---
If using the geo IP logging feature, then this product includes GeoLite2 data
created by MaxMind, available from https://www.maxmind.com.

View file

@ -200,6 +200,12 @@ impl Config {
#[derive(Deserialize, Serialize, Clone)] #[derive(Deserialize, Serialize, Clone)]
pub struct ClientSecret(String); pub struct ClientSecret(String);
impl ClientSecret {
pub fn as_str(&self) -> &str {
self.0.as_ref()
}
}
impl std::fmt::Debug for ClientSecret { impl std::fmt::Debug for ClientSecret {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
write!(f, "[client secret]") write!(f, "[client secret]")
@ -301,8 +307,11 @@ struct CliArgs {
/// respectively. /// respectively.
#[clap(short, long, parse(from_occurrences), conflicts_with = "verbose")] #[clap(short, long, parse(from_occurrences), conflicts_with = "verbose")]
pub quiet: usize, pub quiet: usize,
/// Unstable options. Intentionally not documented.
#[clap(short = 'Z', long)] #[clap(short = 'Z', long)]
pub unstable_options: Vec<UnstableOptions>, pub unstable_options: Vec<UnstableOptions>,
/// Override the image server with the one provided. Do not set this unless
/// you know what you're doing.
#[clap(long)] #[clap(long)]
pub override_upstream: Option<Url>, pub override_upstream: Option<Url>,
/// Enables ephemeral disk encryption. Items written to disk are first /// Enables ephemeral disk encryption. Items written to disk are first
@ -310,8 +319,11 @@ struct CliArgs {
/// performance, privacy, and usability with this flag enabled. /// performance, privacy, and usability with this flag enabled.
#[clap(short, long)] #[clap(short, long)]
pub ephemeral_disk_encryption: bool, pub ephemeral_disk_encryption: bool,
/// The path to the config file. Default value is `./settings.yaml`.
#[clap(short, long)] #[clap(short, long)]
pub config_path: Option<PathBuf>, pub config_path: Option<PathBuf>,
/// Whether to use an in-memory cache in addition to the disk cache. Default
/// value is "on_disk", other options are "lfu" and "lru".
#[clap(short = 't', long)] #[clap(short = 't', long)]
pub cache_type: Option<CacheType>, pub cache_type: Option<CacheType>,
} }

View file

@ -5,17 +5,21 @@
use std::env::VarError; use std::env::VarError;
use std::error::Error; use std::error::Error;
use std::fmt::Display; use std::fmt::Display;
use std::net::SocketAddr;
use std::num::ParseIntError; use std::num::ParseIntError;
use std::str::FromStr;
use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc; use std::sync::Arc;
use std::time::Duration; use std::time::Duration;
use actix_web::dev::Service;
use actix_web::rt::{spawn, time, System}; use actix_web::rt::{spawn, time, System};
use actix_web::web::{self, Data}; use actix_web::web::{self, Data};
use actix_web::{App, HttpResponse, HttpServer}; use actix_web::{App, HttpResponse, HttpServer};
use cache::{Cache, DiskCache}; use cache::{Cache, DiskCache};
use chacha20::Key; use chacha20::Key;
use config::Config; use config::Config;
use maxminddb::geoip2;
use parking_lot::RwLock; use parking_lot::RwLock;
use rustls::{NoClientAuth, ServerConfig}; use rustls::{NoClientAuth, ServerConfig};
use sodiumoxide::crypto::stream::xchacha20::gen_key; use sodiumoxide::crypto::stream::xchacha20::gen_key;
@ -27,6 +31,7 @@ use tracing::{debug, error, info, warn};
use crate::cache::mem::{Lfu, Lru}; use crate::cache::mem::{Lfu, Lru};
use crate::cache::{MemoryCache, ENCRYPTION_KEY}; use crate::cache::{MemoryCache, ENCRYPTION_KEY};
use crate::config::{CacheType, UnstableOptions, OFFLINE_MODE}; use crate::config::{CacheType, UnstableOptions, OFFLINE_MODE};
use crate::metrics::{record_country_visit, GEOIP_DATABASE};
use crate::state::DynamicServerCert; use crate::state::DynamicServerCert;
mod cache; mod cache;
@ -107,6 +112,12 @@ async fn main() -> Result<(), Box<dyn Error>> {
metrics::init(); metrics::init();
} }
if let Some(key) = config.geoip_license_key.clone() {
if let Err(e) = metrics::load_geo_ip_data(key).await {
error!("Failed to initialize geo ip db: {}", e);
}
}
// HTTP Server init // HTTP Server init
let server = if OFFLINE_MODE.load(Ordering::Acquire) { let server = if OFFLINE_MODE.load(Ordering::Acquire) {
@ -171,6 +182,23 @@ async fn main() -> Result<(), Box<dyn Error>> {
let server = HttpServer::new(move || { let server = HttpServer::new(move || {
App::new() App::new()
.wrap(actix_web::middleware::Compress::default()) .wrap(actix_web::middleware::Compress::default())
.wrap_fn(|req, srv| {
if let Some(reader) = GEOIP_DATABASE.get() {
let maybe_country = req
.connection_info()
.realip_remote_addr()
.map(SocketAddr::from_str)
.and_then(Result::ok)
.as_ref()
.map(SocketAddr::ip)
.map(|ip| reader.lookup::<geoip2::Country>(ip))
.and_then(Result::ok);
record_country_visit(maybe_country);
}
srv.call(req)
})
.service(routes::index) .service(routes::index)
.service(routes::token_data) .service(routes::token_data)
.service(routes::token_data_saver) .service(routes::token_data_saver)

View file

@ -1,5 +1,28 @@
use once_cell::sync::Lazy; use std::fs::metadata;
use prometheus::{register_int_counter, IntCounter}; use std::hint::unreachable_unchecked;
use std::time::SystemTime;
use chrono::Duration;
use flate2::read::GzDecoder;
use maxminddb::geoip2::Country;
use once_cell::sync::{Lazy, OnceCell};
use prometheus::{register_int_counter, register_int_counter_vec, IntCounter, IntCounterVec};
use tar::Archive;
use thiserror::Error;
use tracing::{debug, field::debug, info, warn};
use crate::config::ClientSecret;
pub static GEOIP_DATABASE: OnceCell<maxminddb::Reader<Vec<u8>>> = OnceCell::new();
static COUNTRY_VISIT_COUNTER: Lazy<IntCounterVec> = Lazy::new(|| {
register_int_counter_vec!(
"country_visits_total",
"The number of visits from a country",
&["country"]
)
.unwrap()
});
macro_rules! init_counters { macro_rules! init_counters {
($(($counter:ident, $ty:ty, $name:literal, $desc:literal),)*) => { ($(($counter:ident, $ty:ty, $name:literal, $desc:literal),)*) => {
@ -11,7 +34,11 @@ macro_rules! init_counters {
#[allow(clippy::shadow_unrelated)] #[allow(clippy::shadow_unrelated)]
pub fn init() { pub fn init() {
// These need to be called at least once, otherwise the macro never
// called and thus the metrics don't get logged
$(let _a = $counter.get();)* $(let _a = $counter.get();)*
init_other();
} }
}; };
} }
@ -20,13 +47,13 @@ init_counters!(
( (
CACHE_HIT_COUNTER, CACHE_HIT_COUNTER,
IntCounter, IntCounter,
"cache_hit", "cache_hit_total",
"The number of cache hits." "The number of cache hits."
), ),
( (
CACHE_MISS_COUNTER, CACHE_MISS_COUNTER,
IntCounter, IntCounter,
"cache_miss", "cache_miss_total",
"The number of cache misses." "The number of cache misses."
), ),
( (
@ -38,19 +65,120 @@ init_counters!(
( (
REQUESTS_DATA_COUNTER, REQUESTS_DATA_COUNTER,
IntCounter, IntCounter,
"requests_data", "requests_data_total",
"The number of requests served from the /data endpoint." "The number of requests served from the /data endpoint."
), ),
( (
REQUESTS_DATA_SAVER_COUNTER, REQUESTS_DATA_SAVER_COUNTER,
IntCounter, IntCounter,
"requests_data_saver", "requests_data_saver_total",
"The number of requests served from the /data-saver endpoint." "The number of requests served from the /data-saver endpoint."
), ),
( (
REQUESTS_OTHER_COUNTER, REQUESTS_OTHER_COUNTER,
IntCounter, IntCounter,
"requests_other", "requests_other_total",
"The total number of request not served by primary endpoints." "The total number of request not served by primary endpoints."
), ),
); );
// initialization for any other counters that aren't simple int counters
fn init_other() {
let _a = COUNTRY_VISIT_COUNTER.local();
}
#[derive(Error, Debug)]
pub enum DbLoadError {
#[error(transparent)]
Reqwest(#[from] reqwest::Error),
#[error(transparent)]
Io(#[from] std::io::Error),
#[error(transparent)]
MaxMindDb(#[from] maxminddb::MaxMindDBError),
}
pub async fn load_geo_ip_data(license_key: ClientSecret) -> Result<(), DbLoadError> {
const DB_PATH: &str = "./GeoLite2-Country.mmdb";
// Check date of db
let db_date_created = metadata(DB_PATH)
.ok()
.and_then(|metadata| match metadata.created() {
Ok(time) => Some(time),
Err(_) => {
debug("fs didn't report birth time, fall back to last modified instead");
metadata.modified().ok()
}
})
.unwrap_or(SystemTime::UNIX_EPOCH);
let duration = match SystemTime::now().duration_since(dbg!(db_date_created)) {
Ok(time) => Duration::from_std(time).expect("duration to fit"),
Err(_) => {
warn!("Clock may have gone backwards?");
Duration::max_value()
}
};
// DB expired, fetch a new one
if duration > Duration::weeks(1) {
fetch_db(license_key).await?;
} else {
info!("Geo IP database isn't old enough, not updating.");
}
// Result literally cannot panic here, buuuuuut if it does we'll panic
GEOIP_DATABASE
.set(maxminddb::Reader::open_readfile(DB_PATH)?)
.map_err(|_| ())
.expect("to set the geo ip db singleton");
Ok(())
}
async fn fetch_db(license_key: ClientSecret) -> Result<(), DbLoadError> {
let resp = reqwest::get(format!("https://download.maxmind.com/app/geoip_download?edition_id=GeoLite2-Country&license_key={}&suffix=tar.gz", license_key.as_str()))
.await?
.bytes()
.await?;
let mut decoder = Archive::new(GzDecoder::new(resp.as_ref()));
let mut decoded_paths: Vec<_> = decoder
.entries()?
.filter_map(Result::ok)
.filter_map(|mut entry| {
let path = entry.path().ok()?.to_path_buf();
let file_name = path.file_name()?;
if file_name != "GeoLite2-Country.mmdb" {
return None;
}
entry.unpack(file_name).ok()?;
Some(path)
})
.collect();
assert_eq!(decoded_paths.len(), 1);
let path = match decoded_paths.pop() {
Some(path) => path,
None => unsafe { unreachable_unchecked() },
};
debug!("Extracted {}", path.as_path().to_string_lossy());
Ok(())
}
pub fn record_country_visit(country: Option<Country>) {
let iso_code = if let Some(country) = country {
country
.country
.and_then(|c| c.iso_code)
.unwrap_or("unknown")
} else {
"unknown"
};
COUNTRY_VISIT_COUNTER
.get_metric_with_label_values(&[iso_code])
.unwrap()
.inc();
}

View file

@ -22,7 +22,7 @@ use crate::CLIENT_API_VERSION;
pub const CONTROL_CENTER_PING_URL: &str = "https://api.mangadex.network/ping"; pub const CONTROL_CENTER_PING_URL: &str = "https://api.mangadex.network/ping";
#[derive(Serialize)] #[derive(Serialize, Debug)]
pub struct Request<'a> { pub struct Request<'a> {
secret: &'a ClientSecret, secret: &'a ClientSecret,
port: Port, port: Port,
@ -177,6 +177,7 @@ pub async fn update_server_state(
data: &mut Arc<RwLockServerState>, data: &mut Arc<RwLockServerState>,
) { ) {
let req = Request::from_config_and_state(secret, cli); let req = Request::from_config_and_state(secret, cli);
debug!("Sending ping request: {:?}", req);
let client = reqwest::Client::new(); let client = reqwest::Client::new();
let resp = client.post(CONTROL_CENTER_PING_URL).json(&req).send().await; let resp = client.post(CONTROL_CENTER_PING_URL).json(&req).send().await;
match resp { match resp {