feat: support token validation for buckal cli

benjamin-747 · benjamin-747 · commit b51f554a1ef1 · 2025-11-04T11:56:09.000+08:00
diff --git a/jupiter/src/storage/base_storage.rs b/jupiter/src/storage/base_storage.rs
@@ -92,7 +92,7 @@ pub trait StorageConnector {
     }
 }
 
-#[derive(Clone)]
+#[derive(Debug, Clone)]
 pub struct BaseStorage {
     pub connection: Arc<DatabaseConnection>,
 }
diff --git a/jupiter/src/storage/user_storage.rs b/jupiter/src/storage/user_storage.rs
@@ -10,7 +10,7 @@ use common::{errors::MegaError, utils::generate_id};
 
 use crate::storage::base_storage::{BaseStorage, StorageConnector};
 
-#[derive(Clone)]
+#[derive(Debug, Clone)]
 pub struct UserStorage {
     pub base: BaseStorage,
 }
@@ -23,6 +23,11 @@ impl Deref for UserStorage {
 }
 
 impl UserStorage {
+    pub fn mock() -> Self {
+        let mock = BaseStorage::mock();
+        Self { base: mock }
+    }
+
     pub async fn save_ssh_key(
         &self,
         username: String,
diff --git a/mono/src/api/oauth/campsite_store.rs b/mono/src/api/oauth/campsite_store.rs
@@ -3,6 +3,7 @@ use std::sync::Arc;
 use anyhow::Context;
 use async_trait::async_trait;
 use http::header::COOKIE;
+use jupiter::storage::user_storage::UserStorage;
 use reqwest::Client;
 use reqwest::Url;
 use tower_sessions::{
@@ -20,6 +21,7 @@ pub struct CampsiteApiStore {
     client: Arc<Client>,
     // cookie_store: Arc<Jar>,
     api_base_url: String,
+    user_storage: UserStorage,
 }
 
 #[async_trait]
@@ -42,14 +44,15 @@ impl SessionStore for CampsiteApiStore {
 }
 
 impl CampsiteApiStore {
-    pub fn new(api_base_url: String) -> Self {
+    pub fn new(api_base_url: String, user_storage: UserStorage) -> Self {
         let client = Client::builder()
             .no_proxy()
             .build()
             .expect("Failed to build client");
         Self {
             client: Arc::new(client),
             api_base_url,
+            user_storage,
         }
     }
 
@@ -82,4 +85,15 @@ impl CampsiteApiStore {
             Ok(None)
         }
     }
+
+    pub async fn load_user_from_token(&self, token: String) -> anyhow::Result<Option<LoginUser>> {
+        if let Some(username) = self.user_storage.find_user_by_token(&token).await? {
+            let user = LoginUser {
+                username,
+                ..Default::default()
+            };
+            return Ok(Some(user));
+        }
+        Ok(None)
+    }
 }
diff --git a/mono/src/api/oauth/mod.rs b/mono/src/api/oauth/mod.rs
@@ -6,7 +6,11 @@ use axum::{
     response::{IntoResponse, Redirect, Response},
     routing::get,
 };
-use axum_extra::{TypedHeader, headers, typed_header::TypedHeaderRejectionReason};
+use axum_extra::{
+    TypedHeader,
+    headers::{self, Authorization, authorization::Bearer},
+    typed_header::TypedHeaderRejectionReason,
+};
 use chrono::{Duration, Utc};
 use http::request::Parts;
 use oauth2::{
@@ -203,6 +207,25 @@ where
     async fn from_request_parts(parts: &mut Parts, state: &S) -> Result<Self, Self::Rejection> {
         let store = CampsiteApiStore::from_ref(state);
 
+        // Try Bearer Token for CLI authentication
+        if let Ok(TypedHeader(Authorization(bearer))) =
+            parts.extract::<TypedHeader<Authorization<Bearer>>>().await
+        {
+            let token = bearer.token().to_string();
+            match store.load_user_from_token(token).await {
+                Ok(Some(user)) => return Ok(user),
+                Ok(None) => {
+                    tracing::error!("Invalid or expired bearer token");
+                    return Err(AuthRedirect);
+                }
+                Err(e) => {
+                    tracing::error!("Error validating bearer token: {:?}", e);
+                    return Err(AuthRedirect);
+                }
+            }
+        }
+
+        // Fall back to Cookie Session (UI)
         let cookies = parts
             .extract::<TypedHeader<headers::Cookie>>()
             .await
@@ -217,15 +240,16 @@ where
         let session_cookie = cookies.get(CAMPSITE_API_COOKIE).ok_or(AuthRedirect)?;
 
         // Load user from external API
-        let user = store
-            .load_user_from_api(session_cookie.to_string())
-            .await
-            .map_err(|e| {
-                tracing::error!("load_user_from_api error: {:?}", e);
-                AuthRedirect
-            })?
-            .ok_or(AuthRedirect)?;
-
-        Ok(user)
+        match store.load_user_from_api(session_cookie.to_string()).await {
+            Ok(Some(user)) => Ok(user),
+            Ok(None) => {
+                tracing::error!("Invalid or expired session cookie");
+                Err(AuthRedirect)
+            }
+            Err(e) => {
+                tracing::error!("Error loading user from cookie session: {:?}", e);
+                Err(AuthRedirect)
+            }
+        }
     }
 }
diff --git a/mono/src/api/oauth/model.rs b/mono/src/api/oauth/model.rs
@@ -1,4 +1,3 @@
-use chrono::{DateTime, NaiveDateTime, Utc};
 use serde::{Deserialize, Serialize};
 
 #[derive(Serialize, Deserialize, Clone, Debug)]
@@ -30,7 +29,6 @@ impl From<GitHubUserJson> for LoginUser {
             email: value.email.unwrap_or_default(),
             avatar_url: value.avatar_url,
             campsite_user_id: String::new(),
-            created_at: Utc::now().naive_utc(),
         }
     }
 }
@@ -41,7 +39,6 @@ pub struct CampsiteUserJson {
     pub id: String,
     pub avatar_url: String,
     pub email: Option<String>,
-    pub created_at: DateTime<Utc>,
 }
 
 impl From<CampsiteUserJson> for LoginUser {
@@ -51,16 +48,14 @@ impl From<CampsiteUserJson> for LoginUser {
             email: value.email.unwrap_or_default(),
             avatar_url: value.avatar_url,
             campsite_user_id: value.id,
-            created_at: value.created_at.naive_utc(),
         }
     }
 }
 
-#[derive(Serialize, Deserialize, Clone, Debug)]
+#[derive(Serialize, Deserialize, Clone, Debug, Default)]
 pub struct LoginUser {
     pub campsite_user_id: String,
     pub username: String,
     pub avatar_url: String,
     pub email: String,
-    pub created_at: NaiveDateTime,
 }
diff --git a/mono/src/server/http_server.rs b/mono/src/server/http_server.rs
@@ -110,7 +110,10 @@ pub async fn app(storage: Storage, host: String, port: u16) -> Router {
     let api_state = MonoApiServiceState {
         storage: storage.clone(),
         oauth_client: Some(oauth_client(oauth_config.clone()).unwrap()),
-        session_store: Some(CampsiteApiStore::new(oauth_config.campsite_api_domain)),
+        session_store: Some(CampsiteApiStore::new(
+            oauth_config.campsite_api_domain,
+            storage.user_storage(),
+        )),
         listen_addr: format!("http://{host}:{port}"),
         entity_store: EntityStore::new(),
     };
diff --git a/mono/tests/campsite_api_store_tests.rs b/mono/tests/campsite_api_store_tests.rs
@@ -3,6 +3,7 @@
 //! These tests cover the CampsiteApiStore's ability to load user information from an external API.
 
 use axum::Router;
+use jupiter::storage::user_storage::UserStorage;
 use mono::api::oauth::campsite_store::CampsiteApiStore;
 use serde_json::json;
 use std::net::SocketAddr;
@@ -53,7 +54,7 @@ async fn test_load_user_from_api_success() {
     let (addr, _handle) = create_mock_campsite_server().await;
     let api_url = format!("http://{}", addr);
 
-    let store = CampsiteApiStore::new(api_url);
+    let store = CampsiteApiStore::new(api_url, UserStorage::mock());
 
     // Test with a valid cookie
     let result = store
@@ -78,7 +79,8 @@ async fn test_load_user_from_api_invalid_cookie() {
 
     // Test with an invalid cookie that causes a 401 response
     // We'll simulate this by using a non-existent endpoint
-    let invalid_store = CampsiteApiStore::new(format!("http://{}/nonexistent", addr));
+    let invalid_store =
+        CampsiteApiStore::new(format!("http://{}/nonexistent", addr), UserStorage::mock());
 
     let result = invalid_store
         .load_user_from_api("invalid_session_cookie".to_string())
@@ -100,7 +102,7 @@ async fn test_load_user_from_api_server_error() {
     let (addr, _handle) = create_mock_campsite_server().await;
     let api_url = format!("http://{}", addr);
 
-    let store = CampsiteApiStore::new(format!("{}/v1/users/error", api_url));
+    let store = CampsiteApiStore::new(format!("{}/v1/users/error", api_url), UserStorage::mock());
 
     let result = store.load_user_from_api("any_cookie".to_string()).await;
 
@@ -118,7 +120,10 @@ async fn test_load_user_from_api_server_error() {
 #[tokio::test]
 async fn test_load_user_from_api_network_error() {
     // Test with an invalid URL that will cause a network error
-    let store = CampsiteApiStore::new("http://invalid.domain.localhost:12345".to_string());
+    let store = CampsiteApiStore::new(
+        "http://invalid.domain.localhost:12345".to_string(),
+        UserStorage::mock(),
+    );
 
     let result = store.load_user_from_api("any_cookie".to_string()).await;
 
diff --git a/mono/tests/login_user_extractor_tests.rs b/mono/tests/login_user_extractor_tests.rs
@@ -3,6 +3,7 @@
 //! These tests cover the LoginUser extractor's ability to extract user information from requests.
 //! Since the extractor relies on CampsiteApiStore, we focus on testing the underlying functionality.
 
+use jupiter::storage::user_storage::UserStorage;
 use mono::api::oauth::campsite_store::CampsiteApiStore;
 use serde_json::json;
 use std::net::SocketAddr;
@@ -44,7 +45,7 @@ async fn test_login_user_extractor_success() {
     let _api_url = format!("http://{}", addr);
 
     // Create a mock store
-    let store = CampsiteApiStore::new(format!("http://{}", addr));
+    let store = CampsiteApiStore::new(format!("http://{}", addr), UserStorage::mock());
 
     // Test the load_user_from_api method directly
     let result = store
@@ -68,7 +69,7 @@ async fn test_login_user_extractor_invalid_cookie() {
     let _api_url = format!("http://{}", addr);
 
     // Create a mock store with a non-existent endpoint to simulate an invalid cookie
-    let store = CampsiteApiStore::new(format!("http://{}/nonexistent", addr));
+    let store = CampsiteApiStore::new(format!("http://{}/nonexistent", addr), UserStorage::mock());
 
     // Test the load_user_from_api method directly
     let result = store
@@ -90,7 +91,7 @@ async fn test_login_user_extractor_missing_cookie() {
     let _api_url = format!("http://{}", addr);
 
     // Create a mock store
-    let store = CampsiteApiStore::new(format!("http://{}", addr));
+    let store = CampsiteApiStore::new(format!("http://{}", addr), UserStorage::mock());
 
     // Test with an empty cookie string to simulate missing cookie
     let result = store.load_user_from_api("".to_string()).await;
@@ -106,7 +107,10 @@ async fn test_login_user_extractor_missing_cookie() {
 #[tokio::test]
 async fn test_login_user_extractor_network_error() {
     // Test with an invalid URL that will cause a network error
-    let store = CampsiteApiStore::new("http://invalid.domain.localhost:12345".to_string());
+    let store = CampsiteApiStore::new(
+        "http://invalid.domain.localhost:12345".to_string(),
+        UserStorage::mock(),
+    );
 
     let result = store.load_user_from_api("any_cookie".to_string()).await;
 

Original file line number	Diff line number	Diff line change
`@@ -92,7 +92,7 @@ pub trait StorageConnector {`
`92`	`92`	`}`
`93`	`93`	`}`
`94`	`94`
`95`		`-#[derive(Clone)]`
	`95`	`+#[derive(Debug, Clone)]`
`96`	`96`	`pub struct BaseStorage {`
`97`	`97`	`pub connection: Arc<DatabaseConnection>,`
`98`	`98`	`}`
Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,3 @@`
`1`		`-use chrono::{DateTime, NaiveDateTime, Utc};`
`2`	`1`	`use serde::{Deserialize, Serialize};`
`3`	`2`
`4`	`3`	`#[derive(Serialize, Deserialize, Clone, Debug)]`
`@@ -30,7 +29,6 @@ impl From<GitHubUserJson> for LoginUser {`
`30`	`29`	`email: value.email.unwrap_or_default(),`
`31`	`30`	`avatar_url: value.avatar_url,`
`32`	`31`	`campsite_user_id: String::new(),`
`33`		`- created_at: Utc::now().naive_utc(),`
`34`	`32`	`}`
`35`	`33`	`}`
`36`	`34`	`}`
`@@ -41,7 +39,6 @@ pub struct CampsiteUserJson {`
`41`	`39`	`pub id: String,`
`42`	`40`	`pub avatar_url: String,`
`43`	`41`	`pub email: Option<String>,`
`44`		`- pub created_at: DateTime<Utc>,`
`45`	`42`	`}`
`46`	`43`
`47`	`44`	`impl From<CampsiteUserJson> for LoginUser {`
`@@ -51,16 +48,14 @@ impl From<CampsiteUserJson> for LoginUser {`
`51`	`48`	`email: value.email.unwrap_or_default(),`
`52`	`49`	`avatar_url: value.avatar_url,`
`53`	`50`	`campsite_user_id: value.id,`
`54`		`- created_at: value.created_at.naive_utc(),`
`55`	`51`	`}`
`56`	`52`	`}`
`57`	`53`	`}`
`58`	`54`
`59`		`-#[derive(Serialize, Deserialize, Clone, Debug)]`
	`55`	`+#[derive(Serialize, Deserialize, Clone, Debug, Default)]`
`60`	`56`	`pub struct LoginUser {`
`61`	`57`	`pub campsite_user_id: String,`
`62`	`58`	`pub username: String,`
`63`	`59`	`pub avatar_url: String,`
`64`	`60`	`pub email: String,`
`65`		`- pub created_at: NaiveDateTime,`
`66`	`61`	`}`