Persist resover choice using a cookie

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
 	}
 }