fix: token and cookie session id miss match. (#16786)

* refactor: adjust response of login handler. * update tests. * fix: session id miss match * feat: logging of logout request. * fix
databendlabs · Nov 8, 2024 · a645b0d · a645b0d
1 parent 9e56fe4
commit a645b0d
Show file tree

Hide file tree

Showing 16 changed files with 234 additions and 58 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -260,6 +260,7 @@ cidr = { version = "0.2.2" }
 clap = { version = "4.4.2", features = ["derive"] }
 comfy-table = "7"
 convert_case = "0.6.0"
+cookie = "0.18.1"
 crc32fast = "1.3.2"
 criterion = "0.5"
 cron = "0.12.0"

diff --git a/src/query/service/src/auth.rs b/src/query/service/src/auth.rs
@@ -34,6 +34,14 @@ pub struct AuthMgr {
     jwt_auth: Option<JwtAuthenticator>,
 }
 
+#[derive(Debug)]
+pub enum CredentialType {
+    DatabendToken,
+    Jwt,
+    Password,
+    NoNeed,
+}
+
 #[derive(Clone)]
 pub enum Credential {
     DatabendToken {
@@ -51,6 +59,17 @@ pub enum Credential {
     NoNeed,
 }
 
+impl Credential {
+    pub fn type_name(&self) -> CredentialType {
+        match self {
+            Credential::DatabendToken { .. } => CredentialType::DatabendToken,
+            Credential::Jwt { .. } => CredentialType::Jwt,
+            Credential::Password { .. } => CredentialType::Password,
+            Credential::NoNeed => CredentialType::NoNeed,
+        }
+    }
+}
+
 impl AuthMgr {
     pub fn init(cfg: &InnerConfig) -> Result<()> {
         GlobalInstance::set(AuthMgr::create(cfg));

diff --git a/src/query/service/src/servers/http/middleware/session.rs b/src/query/service/src/servers/http/middleware/session.rs
@@ -372,7 +372,7 @@ impl<E> HTTPSessionEndpoint<E> {
             (None, None) => {
                 if cookie_enabled {
                     let id = Uuid::new_v4().to_string();
-                    info!("new cookie session id: {}", id);
+                    info!("new session id: {}", id);
                     req.cookie()
                         .add(Cookie::new_with_str(COOKIE_SESSION_ID, &id));
                     Some(id)

diff --git a/src/query/service/src/servers/http/v1/session/client_session_manager.rs b/src/query/service/src/servers/http/v1/session/client_session_manager.rs
@@ -173,6 +173,7 @@ impl ClientSessionManager {
     pub async fn new_token_pair(
         &self,
         session: &Arc<Session>,
+        client_session_id: String,
         old_refresh_token: Option<String>,
         old_session_token: Option<String>,
     ) -> Result<(String, TokenPair)> {
@@ -181,22 +182,12 @@ impl ClientSessionManager {
         let tenant_name = tenant.tenant_name().to_string();
         let user = session.get_current_user()?.name;
         let auth_role = session.privilege_mgr().get_auth_role();
-        let client_session_id = if let Some(old) = &old_refresh_token {
-            let (claim, _) = SessionClaim::decode(old)?;
-            assert_eq!(tenant_name, claim.tenant);
-            assert_eq!(user, claim.user);
-            assert_eq!(auth_role, claim.auth_role);
-            claim.session_id
-        } else {
-            uuid::Uuid::new_v4().to_string()
-        };
-
         let client_session_api = UserApiProvider::instance().client_session_api(&tenant);
 
         // new refresh token
         let now = unix_ts();
         let mut claim = SessionClaim::new(
-            None,
+            client_session_id.clone(),
             &tenant_name,
             &user,
             &auth_role,

diff --git a/src/query/service/src/servers/http/v1/session/login_handler.rs b/src/query/service/src/servers/http/v1/session/login_handler.rs
@@ -36,15 +36,19 @@ struct LoginRequest {
     pub settings: Option<BTreeMap<String, String>>,
 }
 
+#[derive(Serialize, Clone, Debug)]
+pub(crate) struct TokensInfo {
+    pub(crate) session_token_ttl_in_secs: u64,
+    pub(crate) session_token: String,
+    pub(crate) refresh_token: String,
+}
+
 #[derive(Serialize, Debug, Clone)]
 pub struct LoginResponse {
     version: String,
     session_id: String,
-    session_token_ttl_in_secs: u64,
-
-    /// for now, only use session token when authed by user-password
-    session_token: Option<String>,
-    refresh_token: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    tokens: Option<TokensInfo>,
 }
 
 /// Although theses can be checked for each /v1/query for now,
@@ -89,18 +93,20 @@ pub async fn login_handler(
     Json(req): Json<LoginRequest>,
     Query(query): Query<LoginQuery>,
 ) -> PoemResult<impl IntoResponse> {
+    let session_id = ctx
+        .client_session_id
+        .as_ref()
+        .expect("login_handler expect session id in ctx")
+        .clone();
     let version = DATABEND_SEMVER.to_string();
     check_login(ctx, &req)
         .await
         .map_err(HttpErrorCode::bad_request)?;
     let id_only = || {
-        let session_id = uuid::Uuid::new_v4().to_string();
         Ok(Json(LoginResponse {
             version: version.clone(),
-            session_id,
-            session_token_ttl_in_secs: 0,
-            session_token: None,
-            refresh_token: None,
+            session_id: session_id.clone(),
+            tokens: None,
         }))
     };
 
@@ -111,16 +117,17 @@ pub async fn login_handler(
                 id_only()
             } else {
                 let (session_id, token_pair) = ClientSessionManager::instance()
-                    .new_token_pair(&ctx.session, None, None)
+                    .new_token_pair(&ctx.session, session_id, None, None)
                     .await
                     .map_err(HttpErrorCode::server_error)?;
                 Ok(Json(LoginResponse {
                     version,
                     session_id,
-
-                    session_token_ttl_in_secs: SESSION_TOKEN_TTL.as_secs(),
-                    session_token: Some(token_pair.session.clone()),
-                    refresh_token: Some(token_pair.refresh.clone()),
+                    tokens: Some(TokensInfo {
+                        session_token_ttl_in_secs: SESSION_TOKEN_TTL.as_secs(),
+                        session_token: token_pair.session.clone(),
+                        refresh_token: token_pair.refresh.clone(),
+                    }),
                 }))
             }
         }

diff --git a/src/query/service/src/servers/http/v1/session/logout_handler.rs b/src/query/service/src/servers/http/v1/session/logout_handler.rs
@@ -13,6 +13,7 @@
 // limitations under the License.
 
 use jwt_simple::prelude::Serialize;
+use log::info;
 use poem::error::Result as PoemResult;
 use poem::web::Json;
 use poem::IntoResponse;
@@ -32,6 +33,12 @@ pub struct LogoutResponse {
 #[async_backtrace::framed]
 pub async fn logout_handler(ctx: &HttpQueryContext) -> PoemResult<impl IntoResponse> {
     if let Some(id) = &ctx.client_session_id {
+        info!(
+            "logout with user={}, client session id={}, credential type={:?}",
+            ctx.user_name,
+            id,
+            ctx.credential.type_name()
+        );
         ClientSessionManager::instance()
             .drop_client_session_state(&ctx.session.get_current_tenant(), &ctx.user_name, id)
             .await

diff --git a/src/query/service/src/servers/http/v1/session/refresh_handler.rs b/src/query/service/src/servers/http/v1/session/refresh_handler.rs
@@ -22,6 +22,7 @@ use crate::auth::Credential;
 use crate::servers::http::error::HttpErrorCode;
 use crate::servers::http::v1::session::client_session_manager::ClientSessionManager;
 use crate::servers::http::v1::session::consts::SESSION_TOKEN_TTL;
+use crate::servers::http::v1::session::login_handler::TokensInfo;
 use crate::servers::http::v1::HttpQueryContext;
 
 #[derive(Deserialize, Clone)]
@@ -32,9 +33,7 @@ struct RefreshRequest {
 
 #[derive(Serialize, Debug, Clone)]
 pub struct RefreshResponse {
-    session_token: Option<String>,
-    refresh_token: Option<String>,
-    session_token_ttl_in_secs: u64,
+    tokens: TokensInfo,
 }
 
 #[poem::handler]
@@ -43,17 +42,29 @@ pub async fn refresh_handler(
     ctx: &HttpQueryContext,
     Json(req): Json<RefreshRequest>,
 ) -> PoemResult<impl IntoResponse> {
+    let client_session_id = ctx
+        .client_session_id
+        .as_ref()
+        .expect("login_handler expect session id in ctx")
+        .clone();
     let mgr = ClientSessionManager::instance();
     match &ctx.credential {
         Credential::DatabendToken { token, .. } => {
             let (_, token_pair) = mgr
-                .new_token_pair(&ctx.session, Some(token.clone()), req.session_token)
+                .new_token_pair(
+                    &ctx.session,
+                    client_session_id,
+                    Some(token.clone()),
+                    req.session_token,
+                )
                 .await
                 .map_err(HttpErrorCode::server_error)?;
             Ok(Json(RefreshResponse {
-                session_token_ttl_in_secs: SESSION_TOKEN_TTL.as_secs(),
-                session_token: Some(token_pair.session.clone()),
-                refresh_token: Some(token_pair.refresh.clone()),
+                tokens: TokensInfo {
+                    session_token_ttl_in_secs: SESSION_TOKEN_TTL.as_secs(),
+                    session_token: token_pair.session.clone(),
+                    refresh_token: token_pair.refresh.clone(),
+                },
             }))
         }
         _ => {

diff --git a/src/query/service/src/servers/http/v1/session/token.rs b/src/query/service/src/servers/http/v1/session/token.rs
@@ -50,7 +50,7 @@ pub fn unix_ts() -> Duration {
 
 impl SessionClaim {
     pub fn new(
-        session_id: Option<String>,
+        session_id: String,
         tenant: &str,
         user: &str,
         auth_role: &Option<String>,
@@ -60,7 +60,7 @@ impl SessionClaim {
             tenant: tenant.to_string(),
             user: user.to_string(),
             auth_role: auth_role.clone(),
-            session_id: session_id.unwrap_or(uuid::Uuid::new_v4().to_string()),
+            session_id,
             nonce: generate_secure_nonce(),
             expire_at_in_secs: (unix_ts() + ttl).as_secs(),
         }

diff --git a/tests/sqllogictests/Cargo.toml b/tests/sqllogictests/Cargo.toml
@@ -16,6 +16,7 @@ name = "databend-sqllogictests"
 [dependencies]
 async-trait = { workspace = true }
 clap = { workspace = true }
+cookie = { workspace = true }
 databend-common-base = { workspace = true }
 databend-common-exception = { workspace = true }
 env_logger = { workspace = true }
@@ -25,13 +26,14 @@ mysql_async = { workspace = true }
 mysql_common = { workspace = true }
 rand = { workspace = true }
 regex = { workspace = true }
-reqwest = { workspace = true }
+reqwest = { workspace = true, features = ["cookies"] }
 serde = { workspace = true }
 serde_json = { workspace = true }
 sqllogictest = { workspace = true }
 sqlparser = { workspace = true }
 thiserror = { workspace = true }
 tokio = { workspace = true }
+url = { workspace = true }
 walkdir = { workspace = true }
 
 [lints]