216 lines
6.0 KiB
Rust
216 lines
6.0 KiB
Rust
|
use std::{
|
||
|
env::VarError,
|
||
|
ffi::OsString,
|
||
|
io::Write,
|
||
|
net::{Ipv4Addr, SocketAddr},
|
||
|
};
|
||
|
|
||
|
use axum::{
|
||
|
extract::{Path, State},
|
||
|
http::{header::REFERER, HeaderMap},
|
||
|
routing::{get, MethodRouter},
|
||
|
Router,
|
||
|
};
|
||
|
use clap::Parser;
|
||
|
use sqlx::sqlite::{SqliteConnectOptions, SqlitePool, SqlitePoolOptions};
|
||
|
use tokio::net::TcpListener;
|
||
|
use url::Url;
|
||
|
|
||
|
#[tokio::main]
|
||
|
async fn main() {
|
||
|
init();
|
||
|
|
||
|
let pool = db().await;
|
||
|
|
||
|
// the core application, defining the routes and handlers
|
||
|
let app = Router::new()
|
||
|
.stripped_clone("/count/", get(get_count))
|
||
|
.route("/count/:period", get(get_count))
|
||
|
.with_state(pool.clone())
|
||
|
.into_make_service();
|
||
|
let listener = mklistener().await;
|
||
|
axum::serve(listener, app)
|
||
|
.with_graceful_shutdown(shutdown_signal())
|
||
|
.await
|
||
|
.unwrap();
|
||
|
|
||
|
pool.close().await;
|
||
|
}
|
||
|
|
||
|
#[axum::debug_handler]
|
||
|
async fn get_count(
|
||
|
State(db): State<SqlitePool>,
|
||
|
period: Option<Path<String>>,
|
||
|
headers: HeaderMap,
|
||
|
) -> String {
|
||
|
let now = chrono::Utc::now();
|
||
|
let referer = headers.get(REFERER);
|
||
|
let page = if let Some(referer) = referer {
|
||
|
let p = referer.to_str().unwrap_or("/").to_string();
|
||
|
if let Ok(path) = Url::parse(&p) {
|
||
|
path.path().to_string()
|
||
|
} else {
|
||
|
return "".to_string();
|
||
|
}
|
||
|
} else {
|
||
|
return "".to_string();
|
||
|
};
|
||
|
let page = &page;
|
||
|
|
||
|
let count = if let Some(Path(period)) = period {
|
||
|
// don't increment the counter, just fetch the requested period
|
||
|
match period.as_str() {
|
||
|
"day" => {
|
||
|
let then = now - chrono::Duration::try_hours(24).unwrap();
|
||
|
let then = then.to_rfc3339();
|
||
|
let then = &then;
|
||
|
sqlx::query_scalar!(
|
||
|
"select count(*) from hits where page = ? and accessed > ?",
|
||
|
page,
|
||
|
then
|
||
|
)
|
||
|
.fetch_one(&db)
|
||
|
.await
|
||
|
.unwrap_or(1)
|
||
|
}
|
||
|
"week" => {
|
||
|
let then = now - chrono::Duration::try_days(7).unwrap();
|
||
|
let then = then.to_rfc3339();
|
||
|
let then = &then;
|
||
|
sqlx::query_scalar!(
|
||
|
"select count(*) from hits where page = ? and accessed > ?",
|
||
|
page,
|
||
|
then
|
||
|
)
|
||
|
.fetch_one(&db)
|
||
|
.await
|
||
|
.unwrap_or(1)
|
||
|
}
|
||
|
_ => 2,
|
||
|
}
|
||
|
} else {
|
||
|
// increment the counter and return the all-time total
|
||
|
let now = now.to_rfc3339();
|
||
|
let now = &now;
|
||
|
sqlx::query!("insert into hits (page, accessed) values (?, ?)", page, now)
|
||
|
.execute(&db)
|
||
|
.await
|
||
|
.unwrap_or_default();
|
||
|
sqlx::query_scalar!("select count(*) from hits where page = ?", page)
|
||
|
.fetch_one(&db)
|
||
|
.await
|
||
|
.unwrap_or(1)
|
||
|
};
|
||
|
|
||
|
format!("{count}")
|
||
|
}
|
||
|
|
||
|
//-************************************************************************
|
||
|
// li'l helpers
|
||
|
//-************************************************************************
|
||
|
#[derive(Debug, Parser)]
|
||
|
#[clap(version, about)]
|
||
|
struct Cli {
|
||
|
#[clap(
|
||
|
long,
|
||
|
short,
|
||
|
help = "Path to environment file.",
|
||
|
default_value = ".env"
|
||
|
)]
|
||
|
pub env: OsString,
|
||
|
}
|
||
|
|
||
|
fn init() {
|
||
|
let cli = Cli::parse();
|
||
|
dotenvy::from_path_override(cli.env).expect("Could not read .env file.");
|
||
|
env_logger::builder()
|
||
|
.format(|buf, record| {
|
||
|
let ts = buf.timestamp();
|
||
|
writeln!(buf, "{}: {}", ts, record.args())
|
||
|
})
|
||
|
.init();
|
||
|
}
|
||
|
|
||
|
async fn db() -> SqlitePool {
|
||
|
let dbfile = std::env::var("DATABASE_FILE").unwrap();
|
||
|
let opts = SqliteConnectOptions::new()
|
||
|
.foreign_keys(true)
|
||
|
.create_if_missing(true)
|
||
|
.journal_mode(sqlx::sqlite::SqliteJournalMode::Wal)
|
||
|
.filename(&dbfile)
|
||
|
.optimize_on_close(true, None);
|
||
|
let pool = SqlitePoolOptions::new()
|
||
|
.connect_with(opts.clone())
|
||
|
.await
|
||
|
.unwrap();
|
||
|
|
||
|
sqlx::migrate!().run(&pool).await.unwrap();
|
||
|
|
||
|
let count = sqlx::query_scalar!("select count(*) from hits")
|
||
|
.fetch_one(&pool)
|
||
|
.await
|
||
|
.expect("could not get customer count from DB");
|
||
|
log::info!("Connected to DB, found {count} customers.");
|
||
|
|
||
|
pool
|
||
|
}
|
||
|
|
||
|
async fn mklistener() -> TcpListener {
|
||
|
let ip =
|
||
|
std::env::var("LISTENING_ADDR").expect("Could not find $LISTENING_ADDR in environment");
|
||
|
let ip: Ipv4Addr = ip
|
||
|
.parse()
|
||
|
.unwrap_or_else(|_| panic!("Could not parse {ip} as an IP address"));
|
||
|
let port: u16 = std::env::var("LISTENING_PORT")
|
||
|
.and_then(|p| p.parse().map_err(|_| VarError::NotPresent))
|
||
|
.unwrap_or_else(|_| {
|
||
|
panic!("Could not find LISTENING_PORT in env or parse if present");
|
||
|
});
|
||
|
let addr = SocketAddr::from((ip, port));
|
||
|
TcpListener::bind(&addr).await.unwrap()
|
||
|
}
|
||
|
|
||
|
/// Adds both routes, with and without a trailing slash.
|
||
|
trait RouterPathStrip<S>
|
||
|
where
|
||
|
S: Clone + Send + Sync + 'static,
|
||
|
{
|
||
|
fn stripped_clone(self, path: &str, method_router: MethodRouter<S>) -> Self;
|
||
|
}
|
||
|
|
||
|
impl<S> RouterPathStrip<S> for Router<S>
|
||
|
where
|
||
|
S: Clone + Send + Sync + 'static,
|
||
|
{
|
||
|
fn stripped_clone(self, path: &str, method_router: MethodRouter<S>) -> Self {
|
||
|
assert!(path.ends_with('/'));
|
||
|
self.route(path, method_router.clone())
|
||
|
.route(path.trim_end_matches('/'), method_router)
|
||
|
}
|
||
|
}
|
||
|
|
||
|
async fn shutdown_signal() {
|
||
|
use tokio::signal;
|
||
|
let ctrl_c = async {
|
||
|
signal::ctrl_c()
|
||
|
.await
|
||
|
.expect("failed to install Ctrl+C handler");
|
||
|
};
|
||
|
|
||
|
#[cfg(unix)]
|
||
|
let terminate = async {
|
||
|
signal::unix::signal(signal::unix::SignalKind::terminate())
|
||
|
.expect("failed to install signal handler")
|
||
|
.recv()
|
||
|
.await;
|
||
|
};
|
||
|
|
||
|
#[cfg(not(unix))]
|
||
|
let terminate = std::future::pending::<()>();
|
||
|
|
||
|
tokio::select! {
|
||
|
_ = ctrl_c => {log::info!("shutting down")},
|
||
|
_ = terminate => {},
|
||
|
}
|
||
|
}
|