use core::fmt; use std::ops::Deref; use axum::{ async_trait, extract::FromRequestParts, http::{header, request::Parts, HeaderName, StatusCode}, RequestPartsExt, }; use sha2::{Digest, Sha256}; use time::{format_description::FormatItem, macros::format_description, OffsetDateTime}; pub struct RemoteIp(String); #[async_trait] impl FromRequestParts for RemoteIp where S: Send + Sync, { type Rejection = (StatusCode, &'static str); async fn from_request_parts(parts: &mut Parts, _state: &S) -> Result { if let Some(forwarded_for) = parts .headers .get(HeaderName::from_static("x-forwarded-for")) { match forwarded_for.to_str() { Err(_e) => Err((StatusCode::INTERNAL_SERVER_ERROR, "server error")), Ok(remoteip) => Ok(Self(remoteip.to_owned())), } } else { Err((StatusCode::INTERNAL_SERVER_ERROR, "server error")) } } } impl Deref for RemoteIp { type Target = str; fn deref(&self) -> &Self::Target { &self.0 } } const FMT: &[FormatItem<'_>] = format_description!("[year]-[month]-[day]"); pub struct SessionId([u8; 4]); #[async_trait] impl FromRequestParts for SessionId where S: Send + Sync, { type Rejection = (StatusCode, &'static str); async fn from_request_parts(parts: &mut Parts, _state: &S) -> Result { let remote = parts .extract::() .await .map(|rip| rip.0) .unwrap_or(String::from("127.0.0.1")); let agent = parts .headers .get(header::USER_AGENT) .and_then(|ua| ua.to_str().ok()) .unwrap_or("unknown"); let date = OffsetDateTime::now_local().unwrap_or_else(|_| OffsetDateTime::now_utc()); let string = format!("{remote} - {agent} - {}", date.format(FMT).unwrap()); let mut hasher = Sha256::new(); hasher.update(string.as_bytes()); let hash = hasher.finalize(); let slc: [u8; 4] = hash.as_slice()[0..4].try_into().unwrap(); Ok(Self(slc)) } } impl fmt::Display for SessionId { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { for byte in self.0.as_slice() { write!(f, "{byte:02X}")?; } Ok(()) } } pub struct Referer(String); #[async_trait] impl FromRequestParts for Referer where S: Send + Sync, { type Rejection = (StatusCode, &'static str); async fn from_request_parts(parts: &mut Parts, _state: &S) -> Result { if let Some(refer) = parts.headers.get(header::REFERER) { match refer.to_str() { Err(_e) => Err((StatusCode::INTERNAL_SERVER_ERROR, "server error")), Ok(refer) => match url::Url::parse(refer) { Err(_e) => Err(( StatusCode::INTERNAL_SERVER_ERROR, "failed to parse referer url", )), Ok(url) => match url.host_str() { None => Err((StatusCode::INTERNAL_SERVER_ERROR, "referer url had no host")), Some(str) => Ok(Referer(str.to_owned())), }, }, } } else { Err((StatusCode::INTERNAL_SERVER_ERROR, "server error")) } } } impl Deref for Referer { type Target = str; fn deref(&self) -> &Self::Target { &self.0 } } impl fmt::Display for Referer { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { write!(f, "{}", self.0) } }