redirect back to original page on login

src/main.rs

@@ -32,7 +32,7 @@ use axum::http::header::{CONTENT_TYPE, ETAG, LOCATION, SET_COOKIE};
 use axum::http::{HeaderValue, Request, StatusCode};
 use axum::response::IntoResponse;
 use axum::{extract::State, routing::get, serve, Router};
-use axum_extra::headers::Cookie;
+use axum_extra::headers::{Cookie, Referer};
 use axum_extra::TypedHeader;
 use demostf_build::Asset;
 pub use error::Error;
@@ -49,6 +49,7 @@ use opentelemetry::KeyValue;
 use opentelemetry_otlp::{SpanExporter, WithExportConfig, WithTonicConfig};
 use opentelemetry_sdk::trace::SdkTracerProvider;
 use opentelemetry_sdk::Resource;
+use reqwest::Url;
 use secretfile::load;
 use sqlx::PgPool;
 use std::convert::Infallible;
@@ -67,6 +68,7 @@ use tonic::transport::{ClientTlsConfig, Identity};
 use tower_http::trace::TraceLayer;
 use tracing::{error, info, info_span, instrument};
 use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt, EnvFilter, Layer};
+use urlencoding::Encoded;
 
 pub type Result<T, E = Error> = std::result::Result<T, E>;
 
@@ -77,6 +79,7 @@ struct App {
     maps: String,
     sync: String,
     map_list: MapList,
+    url: String,
     pub session_store: MemoryStore,
 }
 
@@ -160,6 +163,7 @@ async fn main() -> Result<()> {
         api: config.site.api,
         maps: config.site.maps,
         sync: config.site.sync,
+        url: config.site.url,
         map_list,
         session_store: session_store.clone(),
     });
@@ -195,6 +199,7 @@ async fn main() -> Result<()> {
         .route("/about", get(about))
         .route("/api", get(api))
         .route("/login/callback", get(login_callback))
+        .route("/login/callback/{return_to}", get(login_callback))
         .route("/login", get(login))
         .route("/logout", get(logout))
         .route("/upload", get(upload))
@@ -358,9 +363,17 @@ async fn demo(
 async fn login_callback(
     State(app): State<Arc<App>>,
     RawQuery(query): RawQuery,
+    return_to: Option<Path<String>>,
 ) -> Result<impl IntoResponse> {
+    let referrer_openid = return_to
+        .as_ref()
+        .map(|return_to| format!("/login/callback/{}", Encoded(return_to.as_str())))
+        .map(|callback| SteamOpenId::new(&app.url, &callback).ok())
+        .flatten();
+    let openid = referrer_openid.as_ref().unwrap_or(&app.openid);
+
     let query = query.as_deref().unwrap_or_default();
-    let steam_id = app.openid.verify(query).await.map_err(|e| {
+    let steam_id = openid.verify(query).await.map_err(|e| {
         error!("{e:?}");
         Error::SteamAuth
     })?;
@@ -376,6 +389,14 @@ async fn login_callback(
         .store_session(session)
         .await?
         .unwrap_or_default();
+
+    let return_to = return_to
+        .as_ref()
+        .map(|Path(return_to)| return_to)
+        .map(|return_to| HeaderValue::from_str(&return_to).ok())
+        .flatten()
+        .unwrap_or(HeaderValue::from_static("/"));
+
     Ok((
         StatusCode::FOUND,
         [
@@ -387,20 +408,35 @@ async fn login_callback(
                 ))
                 .expect("invalid cookie"),
             ),
-            (LOCATION, HeaderValue::from_static("/")),
+            (LOCATION, return_to),
         ],
     ))
 }
 
 #[instrument(skip(app))]
 #[axum::debug_handler]
-async fn login(State(app): State<Arc<App>>) -> impl IntoResponse {
+async fn login(
+    State(app): State<Arc<App>>,
+    referrer: Option<TypedHeader<Referer>>,
+) -> impl IntoResponse {
+    let referrer = referrer
+        .map(|TypedHeader(referrer)| referrer.to_string())
+        .map(|referrer| Url::parse(&referrer).ok())
+        .flatten();
+
+    let referrer_openid = referrer
+        .as_ref()
+        .map(|referrer| format!("/login/callback/{}", Encoded(referrer.path())))
+        .map(|callback| SteamOpenId::new(&app.url, &callback).ok())
+        .flatten();
+    let openid = referrer_openid.as_ref().unwrap_or(&app.openid);
     (
         StatusCode::FOUND,
         [(
             LOCATION,
-            HeaderValue::from_str(app.openid.get_redirect_url()).unwrap(),
+            HeaderValue::from_str(openid.get_redirect_url()).unwrap(),
         )],
+        "",
     )
 }