use crate::agency::Agency; use querying::target::Target; use querying::{Check, CheckVerdict, Checker}; use rocket::http::Status; use rocket::outcome::{IntoOutcome, try_outcome}; use rocket::request::{FromRequest, Outcome}; use rocket::tokio::sync::RwLockReadGuard; use rocket::{Request, State}; use rocket_client_addr::ClientRealAddr; use serde::Serialize; use sqlx::types::Uuid; use sqlx::types::chrono::NaiveDateTime; pub async fn save_query( db: &mut sqlx::PgConnection, target: &Target, check: &Check, addr: &ClientRealAddr, checker: RwLockReadGuard<'_, Checker>, ) -> Result { let (cdn_networks, cdn_providers, rkn_domain): (Vec<_>, Vec<_>, Option<_>) = if let CheckVerdict::Blocked { cdn_provider_subnets, rkn_domain, .. } = &check.verdict { ( cdn_provider_subnets .values() .flatten() .map(|n| n.cidr.to_string()) .collect(), cdn_provider_subnets.keys().map(|p| p.to_string()).collect(), rkn_domain.clone(), ) } else { (vec![], vec![], None) }; let (resolved_ips, cdn_networks) = match target { Target::Asn(_) => (vec![], vec![]), _ => ( check .ips .iter() .map(|i| i.to_string()) .collect::>(), cdn_networks, ), }; let id = sqlx::query_scalar( "INSERT INTO queries ( query, source_ip, source_country_code, source_city_geo_name_id, target_country_code, target_asn, target_provider, resolved_ips, cdn_networks, cdn_providers, rkn_domain ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11) RETURNING id", ) .bind(target.to_query()) .bind(addr.ip.to_string()) .bind( checker .geo_ip(addr.ip) .await .map(|i| i.country_code) .ok() .flatten(), ) .bind(check.geo.city_geo_name_id.map(|id| id as i32)) .bind(check.geo.country_code.clone()) .bind(check.geo.asn.clone()) .bind(check.geo.organisation.clone()) .bind(resolved_ips) .bind(cdn_networks) .bind(cdn_providers) .bind(rkn_domain) .fetch_one(db) .await?; Ok(id) } #[rocket::async_trait] impl<'r> FromRequest<'r> for Agency { type Error = Option; async fn from_request(request: &'r Request<'_>) -> Outcome { let pool = request.guard::<&State>().await.unwrap(); let mut db = try_outcome!( pool.acquire() .await .map_err(|e| Some(e)) .or_forward(Status::InternalServerError) ); let token = request.headers().get_one("Authorization"); let token = try_outcome!( token .and_then(|t| t.split_once(" ")) .map(|(_, tok)| tok.to_string()) .or_forward(Status::Unauthorized) ); let agency = try_outcome!( sqlx::query!("SELECT id, name FROM reporters WHERE token = $1", token) .fetch_optional(&mut *db) .await .map_err(|e| Some(e)) .or_forward(Status::InternalServerError) ); agency .map(|r| Agency { id: r.id, name: r.name, }) .or_forward(Status::Unauthorized) } } #[derive(Serialize, Debug, sqlx::FromRow)] pub struct WhitelistedEntry { domain: Option, rank: Option, last_ok: Option, } pub async fn check_whitelist( domain: &str, db: &mut sqlx::PgConnection, ) -> Result, sqlx::Error> { if domain.chars().filter(|c| *c == '.').count() > 4 { return Ok(None); } sqlx::query_as!( WhitelistedEntry, "SELECT * FROM whitelist WHERE $1 = domain OR $1 LIKE CONCAT('%.', domain) ORDER BY LENGTH(domain) DESC LIMIT 1", domain ) .fetch_optional(db) .await .into() } #[derive(Debug, Serialize, sqlx::FromRow)] pub struct WhitelistHistogramBin { pub bin_id: Option, pub bin_min_rank: Option, pub bin_max_rank: Option, pub count: Option, } pub async fn collect_histogram( db: &mut sqlx::PgConnection, bins: i32, limit: i32, filter: bool, ) -> Result, sqlx::Error> { sqlx::query_as!( WhitelistHistogramBin, "WITH bins AS ( SELECT generate_series(0, $1 - 1) AS bin ) SELECT b.bin as bin_id, b.bin * $2 + 1 AS bin_min_rank, (b.bin + 1) * $2 AS bin_max_rank, COUNT(case when not $3 or w.domain not like '%.co.uk' then 1 end) AS count FROM bins b LEFT JOIN whitelist w ON FLOOR(w.rank / $2) = b.bin GROUP BY b.bin ORDER BY b.bin;", bins, limit / bins, filter ) .fetch_all(db) .await .into() }