mobile update

backend/src/lib.rs

@@ -1,10 +1,14 @@
 use axum::{
     routing::{get, post, put, delete},
     Router, middleware as axum_middleware,
+    Extension,
 };
 use sea_orm::{sqlx, Database, DatabaseConnection, DbErr};
 use sea_orm_migration::prelude::*;
 use std::net::SocketAddr;
+use std::collections::HashMap;
+use std::sync::{Arc, Mutex};
+use std::time::Instant;
 use utoipa::OpenApi;
 use utoipa_swagger_ui::SwaggerUi;
 use tower_sessions::{Expiry, SessionManagerLayer, cookie::SameSite};
@@ -14,6 +18,8 @@ use time::Duration;
 use tower_http::cors::CorsLayer;
 use axum::http::{Method, HeaderValue};
 
+pub type MobileTokenStore = Arc<Mutex<HashMap<String, (i32, Instant)>>>;
+
 pub mod models;
 pub mod services;
 pub mod migration;
@@ -168,9 +174,13 @@ pub async fn create_app(db: DatabaseConnection) -> Result<Router, DbErr> {
         .layer(auth_layer.clone())
         .with_state(db.clone());
 
+    let mobile_token_store: MobileTokenStore = Arc::new(Mutex::new(HashMap::new()));
+
     let oauth_routes = Router::new()
         .route("/auth/google", get(routes::oauth::google_auth))
         .route("/auth/google/callback", get(routes::oauth::google_callback))
+        .route("/auth/mobile-callback", get(routes::oauth::mobile_callback))
+        .layer(Extension(mobile_token_store))
         .layer(auth_layer.clone())
         .with_state(db.clone());
 

backend/src/routes/oauth.rs

@@ -1,17 +1,20 @@
 use axum::{
     extract::{Query, State},
     http::StatusCode,
-    response::Redirect,
+    response::{Html, IntoResponse, Redirect, Response},
+    Extension,
     Json,
 };
 use axum_login::AuthSession;
-use sea_orm::DatabaseConnection;
+use sea_orm::{DatabaseConnection, EntityTrait};
 use serde::{Deserialize, Serialize};
 use tower_sessions::Session;
 use utoipa::ToSchema;
 
 use crate::auth::AuthBackend;
+use crate::models::User;
 use crate::services::OAuthService;
+use crate::MobileTokenStore;
 
 const CSRF_TOKEN_KEY: &str = "oauth_csrf_token";
 const FRONTEND_URL_KEY: &str = "oauth_frontend_url";
@@ -19,6 +22,7 @@ const FRONTEND_URL_KEY: &str = "oauth_frontend_url";
 #[derive(Debug, Deserialize, ToSchema)]
 pub struct GoogleAuthQuery {
     pub redirect_url: Option<String>,
+    pub mobile: Option<bool>,
 }
 
 #[derive(Debug, Deserialize)]
@@ -62,6 +66,13 @@ pub async fn google_auth(
             .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
     }
 
+    if query.mobile.unwrap_or(false) {
+        session
+            .insert("oauth_mobile", true)
+            .await
+            .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
+    }
+
     Ok(Json(OAuthUrlResponse { url: auth_url }))
 }
 
@@ -78,8 +89,9 @@ pub async fn google_callback(
     mut auth_session: AuthSession<AuthBackend>,
     session: Session,
     State(db): State<DatabaseConnection>,
+    Extension(token_store): Extension<MobileTokenStore>,
     Query(query): Query<GoogleCallbackQuery>,
-) -> Result<Redirect, StatusCode> {
+) -> Result<Response, StatusCode> {
     let stored_csrf: Option<String> = session
         .get(CSRF_TOKEN_KEY)
         .await
@@ -90,8 +102,15 @@ pub async fn google_callback(
         .await
         .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
 
+    let is_mobile: bool = session
+        .get("oauth_mobile")
+        .await
+        .unwrap_or(None)
+        .unwrap_or(false);
+
     session.remove::<String>(CSRF_TOKEN_KEY).await.ok();
     session.remove::<String>(FRONTEND_URL_KEY).await.ok();
+    session.remove::<bool>("oauth_mobile").await.ok();
 
     if stored_csrf.as_deref() != Some(&query.state) {
         return Err(StatusCode::UNAUTHORIZED);
@@ -134,7 +153,61 @@ pub async fn google_callback(
         }
     }
 
-    let redirect_url = frontend_url.unwrap_or_else(|| "http://localhost:3000".to_string());
+    if is_mobile {
+        let token = uuid::Uuid::new_v4().to_string();
+        {
+            let mut store = token_store.lock().unwrap();
+            store.insert(token.clone(), (user.id, std::time::Instant::now()));
+        }
+        let deep_link = format!("com.arrelin.family-budget-android://auth?token={}", token);
+        let html = format!(
+            r#"<!DOCTYPE html><html><head><meta http-equiv="refresh" content="0;url={0}"></head>
+            <body><script>window.location="{0}"</script></body></html>"#,
+            deep_link
+        );
+        return Ok(Html(html).into_response());
+    }
 
-    Ok(Redirect::temporary(&redirect_url))
+    let redirect_url = frontend_url.unwrap_or_else(|| "http://localhost:3000".to_string());
+    Ok(Redirect::temporary(&redirect_url).into_response())
+}
+
+#[derive(Debug, Deserialize)]
+pub struct MobileCallbackQuery {
+    pub token: String,
+}
+
+pub async fn mobile_callback(
+    mut auth_session: AuthSession<AuthBackend>,
+    State(db): State<DatabaseConnection>,
+    Extension(token_store): Extension<MobileTokenStore>,
+    Query(query): Query<MobileCallbackQuery>,
+) -> Result<Json<serde_json::Value>, StatusCode> {
+    let user_id = {
+        let mut store = token_store.lock().unwrap();
+        match store.get(&query.token) {
+            Some((uid, created_at)) if created_at.elapsed().as_secs() < 300 => {
+                let uid = *uid;
+                store.remove(&query.token);
+                uid
+            }
+            _ => {
+                store.remove(&query.token);
+                return Err(StatusCode::UNAUTHORIZED);
+            }
+        }
+    };
+
+    let user = User::find_by_id(user_id)
+        .one(&db)
+        .await
+        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?
+        .ok_or(StatusCode::UNAUTHORIZED)?;
+
+    auth_session
+        .login(&user)
+        .await
+        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
+
+    Ok(Json(serde_json::json!({"success": true})))
 }