Persist resover choice using a cookie

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"

Cargo.toml

@@ -8,6 +8,7 @@ authors = ["Slatian <baschdel@disroot.org>"]
 
 [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"

src/main.rs

@@ -384,11 +384,22 @@ async fn settings_query_middleware<B>(
 	Query(query): Query<SettingsQuery>,
 	Extension(config): Extension<config::EchoIpServiceConfig>,
 	Extension(derived_config): Extension<DerivedConfiguration>,
+	cookie_header: Option<TypedHeader<headers::Cookie>>,
 	user_agent_header: Option<TypedHeader<headers::UserAgent>>,
 	mut req: Request<B>,
 	next: Next<B>
 ) -> 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<B>(
 			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
 }

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<Selectable>,
-	//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
 	}
 }

templates/base.html

@@ -20,10 +20,6 @@
 	{% if extra.stylesheet %}<link rel="stylesheet" href="{{extra.stylesheet}}" type="text/css" />{% endif %}
 	{% if extra.favicon %}<link rel="icon" href="{{extra.favicon}}" type="{{extra.favicon_mimetype|default(value="image/png")}}" \>{% endif %}
 	<!-- View: {{view}} -->
-	{% set used_dns_resolver = "default" %}
-	{% if data.result.used_dns_resolver %}
-		{% set used_dns_resolver = data.result.used_dns_resolver %}
-	{% endif%}
 </head>
 <body>
 	<header>
@@ -36,7 +32,7 @@
 					value="{% if view == "dig" %}{{ data.query }}{% elif view == "ip" %}{{ data.result.address }}{% elif view == "asn"%}AS{{ data.asn }}{% endif %}"/>
 				<select name="dns">
 					{% for r in dns_resolvers %}
-						<option value="{{ r.id }}" {% if r.id == used_dns_resolver %}selected{%endif%}>{{ r.name }}</option>
+						<option value="{{ r.id }}" {% if r.id == dns_resolver_id %}selected{%endif%}>{{ r.name }}</option>
 					{% endfor %}
 				</select>
 				<input type="submit" value="Query"/>