2 files changed, 125 insertions, 3 deletions

diff --git a/corgi/src/main.rs b/corgi/src/main.rs
index cd3b67c..6a3c528 100644
--- a/corgi/src/main.rs
+++ b/corgi/src/main.rs
@@ -1,7 +1,9 @@
 use std::{
 	net::{IpAddr, SocketAddr},
+	path::PathBuf,
 	pin::Pin,
 	process::Stdio,
+	sync::Arc,
 	time::Instant,
 };
 
@@ -17,9 +19,11 @@ use hyper::{
 };
 use hyper_util::rt::TokioIo;
 use regex_lite::Regex;
+use stats::Stats;
 use tokio::{io::AsyncWriteExt, net::TcpListener, process::Command, runtime::Runtime};
 
 mod caller;
+mod stats;
 
 #[derive(Clone, Debug)]
 pub struct Settings {
@@ -64,8 +68,13 @@ fn main() {
 		}
 	}
 
+	let stats = Stats::new(PathBuf::from(
+		conf.get("Server/StatsDb").unwrap().to_owned(),
+	));
+	stats.create_tables();
+
 	let rt = Runtime::new().unwrap();
-	rt.block_on(async { run(settings).await });
+	rt.block_on(async { run(settings, stats).await });
 }
 
 fn parse_script_conf(conf: &Value) -> Script {
@@ -106,12 +115,13 @@ fn parse_script_conf(conf: &Value) -> Script {
 }
 
 // We have tokio::main at home :)
-async fn run(settings: Settings) {
+async fn run(settings: Settings, stats: Stats) {
 	let addr = SocketAddr::from(([0, 0, 0, 0], settings.port));
 	let listen = TcpListener::bind(addr).await.unwrap();
 
 	let svc = Svc {
 		settings,
+		stats: Arc::new(stats),
 		client_addr: addr,
 	};
 
@@ -129,6 +139,7 @@ async fn run(settings: Settings) {
 #[derive(Clone, Debug)]
 struct Svc {
 	settings: Settings,
+	stats: Arc<Stats>,
 	client_addr: SocketAddr,
 }
 
@@ -140,14 +151,16 @@ impl Service<Request<Incoming>> for Svc {
 	fn call(&self, req: Request<Incoming>) -> Self::Future {
 		let settings = self.settings.clone();
 		let caddr = self.client_addr;
+		let stats = self.stats.clone();
 
-		Box::pin(async move { Ok(Self::handle(settings, caddr, req).await) })
+		Box::pin(async move { Ok(Self::handle(settings, stats, caddr, req).await) })
 	}
 }
 
 impl Svc {
 	async fn handle(
 		settings: Settings,
+		stats: Arc<Stats>,
 		caddr: SocketAddr,
 		req: Request<Incoming>,
 	) -> Response<Full<Bytes>> {
@@ -242,6 +255,13 @@ impl Svc {
 			response = response.header(key, value);
 		}
 
+		let db_req = stats::Request {
+			agent: &uagent,
+			ip_address: &client_addr,
+			script: &script.name,
+			path: &path,
+		};
+
 		println!(
 			"served to [{client_addr}]\n\tscript: {}\n\tpath: {path}\n\tcgi took {}ms. total time {}ms\n\tUA: {uagent}",
 			&script.name,
@@ -249,6 +269,8 @@ impl Svc {
 			start.elapsed().as_millis()
 		);
 
+		stats.log_request(db_req);
+
 		let response_body = cgi_response.body.map(|v| Bytes::from(v)).unwrap_or(Bytes::new());
 		response.body(Full::new(response_body)).unwrap()
 	}
diff --git a/corgi/src/stats.rs b/corgi/src/stats.rs
new file mode 100644
index 0000000..0e3b99a
--- /dev/null
+++ b/corgi/src/stats.rs
@@ -0,0 +1,100 @@
+use std::{
+	net::{IpAddr, SocketAddr},
+	path::PathBuf,
+	sync::Mutex,
+};
+
+use base64::{Engine, prelude::BASE64_STANDARD_NO_PAD};
+use rusqlite::{Connection, OptionalExtension, params};
+use sha2::{Digest, Sha256};
+
+#[derive(Debug)]
+pub struct Stats {
+	conn: Mutex<Connection>,
+}
+
+impl Stats {
+	pub fn new(db_path: PathBuf) -> Self {
+		Self {
+			conn: Mutex::new(Connection::open(db_path).unwrap()),
+		}
+	}
+
+	pub fn create_tables(&self) {
+		let conn = self.conn.lock().unwrap();
+		conn.execute(CREATE_TABLE_AGENT, ()).unwrap();
+		conn.execute(CREATE_TABLE_REQUESTS, ()).unwrap();
+	}
+
+	pub fn log_request(&self, request: Request) {
+		let Request {
+			agent,
+			ip_address,
+			script,
+			path,
+		} = request;
+
+		let mut hasher = Sha256::new();
+		hasher.update(agent.as_bytes());
+		let hash = hasher.finalize();
+		let agent_hash = BASE64_STANDARD_NO_PAD.encode(&hash[..]);
+
+		let conn = self.conn.lock().unwrap();
+
+		// Try to get the agent ID from the hash
+		let maybe_agent_id: Option<i64> = conn
+			.query_row(
+				"SELECT id, hash FROM agents WHERE hash=?1",
+				params![&agent_hash],
+				|row| row.get(0),
+			)
+			.optional()
+			.unwrap();
+
+		// Can't find the agent_id? Push a new entry
+		let agent_id = match maybe_agent_id {
+			Some(aid) => aid,
+			None => {
+				conn.execute(
+					"INSERT INTO agents(hash, agent) VALUES(?1, ?2)",
+					params![&agent_hash, agent],
+				)
+				.unwrap();
+
+				conn.last_insert_rowid()
+			}
+		};
+
+		conn.execute(
+			"INSERT INTO requests(agent_id, ip_address, script, path) VALUES(?1, ?2, ?3, ?4)",
+			params![agent_id, ip_address.to_string(), script, path],
+		)
+		.unwrap();
+	}
+}
+
+pub struct Request<'r> {
+	pub agent: &'r str,
+	pub ip_address: &'r IpAddr,
+	pub script: &'r str,
+	pub path: &'r str,
+}
+
+const CREATE_TABLE_AGENT: &'static str = "\
+	CREATE TABLE IF NOT EXISTS agents(
+		id INTEGER PRIMARY KEY AUTOINCREMENT,
+		hash TEXT NOT NULL,
+		agent TEXT NOT NULL
+	);";
+
+const CREATE_TABLE_REQUESTS: &'static str = "\
+	CREATE TABLE IF NOT EXISTS requests(
+		id INTEGER PRIMARY KEY AUTOINCREMENT,
+		timestamp TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP,
+		agent_id INTEGER NOT NULL,
+		ip_address TEXT NOT NULL,
+		script TEXT NOT NULL,
+		path TEXT NOT NULL,
+		FOREIGN KEY (agent_id)
+			REFERENCES agents(id)
+	);";