From d88b15ba02f490b95fc6d51724c93c0977cab1a3 Mon Sep 17 00:00:00 2001 From: Slatian Date: Sat, 5 Aug 2023 22:36:28 +0200 Subject: [PATCH] Persist resover choice using a cookie --- Cargo.lock | 68 ++++++++++++++++++++++++++++++++++++++++ Cargo.toml | 1 + src/main.rs | 14 ++++++++- src/templating_engine.rs | 17 +++++++++- templates/base.html | 6 +--- 5 files changed, 99 insertions(+), 7 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 1aae29c..45dbff6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -169,6 +169,28 @@ dependencies = [ "tower-service", ] +[[package]] +name = "axum-extra" +version = "0.7.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a93e433be9382c737320af3924f7d5fc6f89c155cf2bf88949d8f5126fab283f" +dependencies = [ + "axum", + "axum-core", + "bytes", + "cookie", + "futures-util", + "http", + "http-body", + "mime", + "pin-project-lite", + "serde", + "tokio", + "tower", + "tower-layer", + "tower-service", +] + [[package]] name = "axum-macros" version = "0.3.8" @@ -353,6 +375,17 @@ version = "1.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "acbf1af155f9b9ef647e42cdc158db4b64a1b61f743629225fde6f3e0be2a7c7" +[[package]] +name = "cookie" +version = "0.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7efb37c3e1ccb1ff97164ad95ac1606e8ccd35b3fa0a7d99a304c7f4a428cc24" +dependencies = [ + "percent-encoding", + "time", + "version_check", +] + [[package]] name = "core-foundation" version = "0.9.3" @@ -416,6 +449,12 @@ version = "2.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2e66c9d817f1720209181c316d28635c050fa304f9c79e47a520882661b7308" +[[package]] +name = "deranged" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7684a49fb1af197853ef7b2ee694bc1f5b4179556f1e5710e1760c5db6f5e929" + [[package]] name = "deunicode" version = "0.4.4" @@ -438,6 +477,7 @@ version = "0.1.0" dependencies = [ "axum", "axum-client-ip", + "axum-extra", "clap", "governor", "idna 0.3.0", @@ -1900,6 +1940,34 @@ dependencies = [ "once_cell", ] +[[package]] +name = "time" +version = "0.3.25" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b0fdd63d58b18d663fbdf70e049f00a22c8e42be082203be7f26589213cd75ea" +dependencies = [ + "deranged", + "itoa", + "serde", + "time-core", + "time-macros", +] + +[[package]] +name = "time-core" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7300fbefb4dadc1af235a9cef3737cea692a9d97e1b9cbcd4ebdae6f8868e6fb" + +[[package]] +name = "time-macros" +version = "0.2.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eb71511c991639bb078fd5bf97757e03914361c48100d52878b8e52b46fb92cd" +dependencies = [ + "time-core", +] + [[package]] name = "tinyvec" version = "1.6.0" diff --git a/Cargo.toml b/Cargo.toml index 570b1cc..aff3838 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -8,6 +8,7 @@ authors = ["Slatian "] [dependencies] axum = { version = "0.6", features = ["macros", "headers"] } +axum-extra = { version = "0.7", features = ["cookie"] } axum-client-ip = "0.4" clap = { version = "4", features = ["derive"] } governor = "0.5" diff --git a/src/main.rs b/src/main.rs index 8ae308f..d370a0c 100644 --- a/src/main.rs +++ b/src/main.rs @@ -384,11 +384,22 @@ async fn settings_query_middleware( Query(query): Query, Extension(config): Extension, Extension(derived_config): Extension, + cookie_header: Option>, user_agent_header: Option>, mut req: Request, next: Next ) -> Response { let mut format = query.format; + let mut dns_resolver_id = config.dns.default_resolver; + + if let Some(resolver_id) = query.dns { + dns_resolver_id = resolver_id; + } else if let Some(cookie_header) = cookie_header { + if let Some(resolver_id) = cookie_header.0.get("dns_resolver") { + dns_resolver_id = resolver_id.to_string(); + } + } + // Try to guess type from user agent if format.is_none() { if let Some(TypedHeader(user_agent)) = user_agent_header { @@ -407,8 +418,9 @@ async fn settings_query_middleware( format: format.unwrap_or(ResponseFormat::TextHtml), lang: query.lang.unwrap_or("en".to_string()), available_dns_resolvers: derived_config.dns_resolver_selectables, + dns_resolver: dns_resolver_id.clone(), }, - dns_resolver_id: query.dns.unwrap_or(config.dns.default_resolver), + dns_resolver_id: dns_resolver_id, }); next.run(req).await } diff --git a/src/templating_engine.rs b/src/templating_engine.rs index e64efaf..9bb54a9 100644 --- a/src/templating_engine.rs +++ b/src/templating_engine.rs @@ -4,12 +4,16 @@ */ use axum::{ + headers::HeaderValue, http::StatusCode, + http::header::SET_COOKIE, response::Html, response::IntoResponse, response::Response, response::Json, }; +use axum_extra::extract::cookie::Cookie; +use axum_extra::extract::cookie; use tera::Tera; use toml::Table; @@ -56,7 +60,7 @@ pub struct TemplateSettings { pub format: ResponseFormat, pub lang: String, pub available_dns_resolvers: Vec, - //pub dns_resolver: String, + pub dns_resolver: String, } #[derive(serde::Deserialize, serde::Serialize, Clone)] @@ -120,6 +124,7 @@ impl Engine { context.insert("format", &settings.format.to_string()); context.insert("language", &settings.lang); context.insert("dns_resolvers", &settings.available_dns_resolvers); + context.insert("dns_resolver_id", &settings.dns_resolver); context.insert("data", &view); context.insert("extra", &self.template_config); @@ -161,6 +166,16 @@ impl Engine { View::NotFound => *response.status_mut() = StatusCode::NOT_FOUND, _ => {}, } + let cookie = Cookie::build("dns_resolver",settings.dns_resolver.clone()) + .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, + ); + } response } } diff --git a/templates/base.html b/templates/base.html index 60d0491..7f858f8 100644 --- a/templates/base.html +++ b/templates/base.html @@ -20,10 +20,6 @@ {% if extra.stylesheet %}{% endif %} {% if extra.favicon %}{% endif %} - {% set used_dns_resolver = "default" %} - {% if data.result.used_dns_resolver %} - {% set used_dns_resolver = data.result.used_dns_resolver %} - {% endif%}
@@ -36,7 +32,7 @@ value="{% if view == "dig" %}{{ data.query }}{% elif view == "ip" %}{{ data.result.address }}{% elif view == "asn"%}AS{{ data.asn }}{% endif %}"/>