diff --git a/src/main.rs b/src/main.rs index b7aa0c8..5f0b390 100644 --- a/src/main.rs +++ b/src/main.rs @@ -44,7 +44,6 @@ mod mycelium; mod ratelimit; mod settings; mod simple_dns; -mod templating_engine; mod view; use crate::geoip::{ @@ -54,11 +53,14 @@ use crate::geoip::{ LocationResult, }; use crate::idna::IdnaName; +use crate::mycelium::MycEngine; use crate::simple_dns::DnsLookupResult; use crate::settings::*; use crate::view::View; use crate::ipinfo::{AddressCast,AddressInfo,AddressScope}; +type TemplatingEngine = MycEngine; + #[derive(Deserialize, Serialize, Clone)] pub struct SettingsQuery { format: Option, @@ -95,7 +97,7 @@ pub struct DigResult { } struct ServiceSharedState { - templating_engine: templating_engine::Engine, + templating_engine: TemplatingEngine, dns_resolvers: HashMap,TokioAsyncResolver>, dns_resolver_aliases: HashMap,Arc>, asn_db: geoip::MMDBCarrier, @@ -201,10 +203,10 @@ async fn main() { } }; - let templating_engine = templating_engine::Engine{ - tera: tera, - template_config: template_extra_config, - }; + let templating_engine = TemplatingEngine::new( + tera, + template_extra_config, + ); // Static file directory diff --git a/src/templating_engine.rs b/src/mycelium/engine.rs similarity index 60% rename from src/templating_engine.rs rename to src/mycelium/engine.rs index d32da16..59c4862 100644 --- a/src/templating_engine.rs +++ b/src/mycelium/engine.rs @@ -1,7 +1,3 @@ -/* - * This is the echoip-slatecave templating engine. - * It wraps around tera in is specialized for echoip-slatecave. - */ use axum::{ body::{Bytes,Full}, @@ -12,38 +8,55 @@ use axum::{ response::IntoResponse, response::Response, }; -use axum_extra::extract::cookie::Cookie; -use axum_extra::extract::cookie; use tera::Tera; use toml::Table; -use crate::view::View; +use std::marker::PhantomData; + use crate::mycelium::MycView; use crate::mycelium::MycFormat; use crate::mycelium::MycFormatFamily; use crate::mycelium::MycQuerySettings; -use crate::settings::QuerySettings; /* The engine itself */ #[derive(Clone)] -pub struct Engine { +pub struct MycEngine +where V: MycView, S: MycQuerySettings, F: MycFormat { pub tera: Tera, pub template_config: Option, + + phantom_view: PhantomData, + phantom_settings: PhantomData, + phantom_format: PhantomData, } -impl Engine { +impl MycEngine +where V: MycView, S: MycQuerySettings, F: MycFormat { + + pub fn new(tera: Tera, template_config: Option
) -> Self { + Self { + tera: tera, + template_config: template_config, + phantom_view: PhantomData, + phantom_settings: PhantomData, + phantom_format: PhantomData, + } + } + + /// Takes settings and a view, converting it to a serveable response. pub async fn render_view( &self, - settings: &QuerySettings, - view: View, + settings: &S, + view: V, ) -> Response { let status_code = view.get_status_code(&settings); + let cookie_string = view.get_cookie_header(&settings); + let format = settings.get_format(); - let mut response = match settings.format.get_family() { + let mut response = match format.get_family() { MycFormatFamily::Template => { let template_name = view.get_template_name(); - let format = settings.get_format(); let mime_type = format.get_mime_type(); let mut context = tera::Context::new(); @@ -81,22 +94,29 @@ impl Engine { view.get_api_response(&settings) } }; - // Set status code - if response.status() == StatusCode::OK && status_code != StatusCode::OK { - *response.status_mut() = status_code; - } - // Set cookies - let cookie = Cookie::build("dns_resolver",settings.dns_resolver_id.to_string()) - .path("/") - .same_site(cookie::SameSite::Strict) - .finish(); - if let Ok(header_value) = HeaderValue::from_str(&cookie.to_string()) { - response.headers_mut().append( - SET_COOKIE, - header_value, - ); + + // Everything went well and nobody did the following work for us. + if response.status() == StatusCode::OK { + + // Set status code + if status_code != StatusCode::OK { + *response.status_mut() = status_code; + } + + // Set cookie header + if let Some(cookie_string) = cookie_string { + if let Ok(header_value) = HeaderValue::from_str(&cookie_string) { + response.headers_mut().append( + SET_COOKIE, + header_value, + ); + } + } + } + // return response response } } + diff --git a/src/mycelium/mod.rs b/src/mycelium/mod.rs index 146113a..138d482 100644 --- a/src/mycelium/mod.rs +++ b/src/mycelium/mod.rs @@ -1,7 +1,9 @@ +mod engine; mod format; mod query_settings; mod view; +pub use self::engine::MycEngine; pub use self::format::HtmlTextJsonFormat; pub use self::format::MycFormat; pub use self::format::MycFormatFamily; diff --git a/src/view.rs b/src/view.rs index ebba2f4..5346ed3 100644 --- a/src/view.rs +++ b/src/view.rs @@ -3,6 +3,8 @@ use axum::http::status::StatusCode; use axum::Json; use axum::response::IntoResponse; use axum::response::Response; +use axum_extra::extract::cookie::Cookie; +use axum_extra::extract::cookie; use crate::DigResult; use crate::IpResult; @@ -47,6 +49,16 @@ impl MycView for View { } } + fn get_cookie_header(&self, settings: &QuerySettings) -> Option { + Some( + Cookie::build("dns_resolver",settings.dns_resolver_id.to_string()) + .path("/") + .same_site(cookie::SameSite::Strict) + .finish() + .to_string() + ) + } + fn get_api_response(self, settings: &QuerySettings) -> Response { match self { Self::Dig{result, ..} => {