From 74f81cb02a208dcd2a22d985d1a1bb24dfa67ac4 Mon Sep 17 00:00:00 2001 From: Meng Zhang Date: Tue, 5 Dec 2023 23:47:10 +0800 Subject: [PATCH] refactor: Use DateTime for sqlite datetime fields (#946) * refactor: use DateTime for RefreshToken.expires_at * refactor: set other date time fields to be DateTime * [autofix.ci] apply automated fixes * [autofix.ci] apply automated fixes (attempt 2/3) --------- Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com> --- Cargo.lock | 1 + ee/tabby-webserver/Cargo.toml | 2 +- ee/tabby-webserver/src/schema/auth.rs | 18 ++++++++------ ee/tabby-webserver/src/service/auth.rs | 15 +++++------- ee/tabby-webserver/src/service/db.rs | 34 ++++++++++++-------------- 5 files changed, 34 insertions(+), 36 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index c9154a3b17bc..7f452dab408a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3868,6 +3868,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "549b9d036d571d42e6e85d1c1425e2ac83491075078ca9a15be021c56b1641f2" dependencies = [ "bitflags 2.4.0", + "chrono", "fallible-iterator", "fallible-streaming-iterator", "hashlink", diff --git a/ee/tabby-webserver/Cargo.toml b/ee/tabby-webserver/Cargo.toml index b3da0bd6689d..a948643d80a4 100644 --- a/ee/tabby-webserver/Cargo.toml +++ b/ee/tabby-webserver/Cargo.toml @@ -20,7 +20,7 @@ juniper-axum = { path = "../../crates/juniper-axum" } lazy_static = "1.4.0" mime_guess = "2.0.4" pin-project = "1.1.3" -rusqlite = { version = "0.29.0", features = ["bundled"] } +rusqlite = { version = "0.29.0", features = ["bundled", "chrono"] } # `async-tokio-rusqlite` is only available from 1.1.0-alpha.2, will bump up version when it's stable rusqlite_migration = { version = "1.1.0-alpha.2", features = ["async-tokio-rusqlite"] } rust-embed = "8.0.0" diff --git a/ee/tabby-webserver/src/schema/auth.rs b/ee/tabby-webserver/src/schema/auth.rs index 7331ca3f67ae..91ff0f3c4fa3 100644 --- a/ee/tabby-webserver/src/schema/auth.rs +++ b/ee/tabby-webserver/src/schema/auth.rs @@ -2,6 +2,7 @@ use std::fmt::Debug; use anyhow::Result; use async_trait::async_trait; +use chrono::{DateTime, Utc}; use jsonwebtoken as jwt; use juniper::{FieldError, GraphQLObject, IntoFieldError, ScalarValue}; use lazy_static::lazy_static; @@ -20,7 +21,6 @@ lazy_static! { jwt_token_secret().as_bytes() ); static ref JWT_DEFAULT_EXP: u64 = 30 * 60; // 30 minutes - static ref JWT_REFRESH_PERIOD: i64 = 7 * 24 * 60 * 60; // 7 days } pub fn generate_jwt(claims: Claims) -> jwt::errors::Result { @@ -39,9 +39,8 @@ fn jwt_token_secret() -> String { std::env::var("TABBY_WEBSERVER_JWT_TOKEN_SECRET").unwrap_or("default_secret".to_string()) } -pub fn generate_refresh_token(utc_ts: i64) -> (String, i64) { - let token = Uuid::new_v4().to_string().replace('-', ""); - (token, utc_ts + *JWT_REFRESH_PERIOD) +pub fn generate_refresh_token() -> String { + Uuid::new_v4().to_string().replace('-', "") } #[derive(Debug, GraphQLObject)] @@ -162,11 +161,15 @@ impl IntoFieldError for RefreshTokenError { pub struct RefreshTokenResponse { pub access_token: String, pub refresh_token: String, - pub refresh_expires_at: f64, + pub refresh_expires_at: DateTime, } impl RefreshTokenResponse { - pub fn new(access_token: String, refresh_token: String, refresh_expires_at: f64) -> Self { + pub fn new( + access_token: String, + refresh_token: String, + refresh_expires_at: DateTime, + ) -> Self { Self { access_token, refresh_token, @@ -292,8 +295,7 @@ mod tests { #[test] fn test_generate_refresh_token() { - let (token, exp) = generate_refresh_token(100); + let token = generate_refresh_token(); assert_eq!(token.len(), 32); - assert_eq!(exp, 100 + *JWT_REFRESH_PERIOD); } } diff --git a/ee/tabby-webserver/src/service/auth.rs b/ee/tabby-webserver/src/service/auth.rs index 833fe50e7901..5a7cb3189c7b 100644 --- a/ee/tabby-webserver/src/service/auth.rs +++ b/ee/tabby-webserver/src/service/auth.rs @@ -146,9 +146,8 @@ impl AuthenticationService for DbConn { .await?; let user = self.get_user(id).await?.unwrap(); - let (refresh_token, expires_at) = generate_refresh_token(chrono::Utc::now().timestamp()); - self.create_refresh_token(id, &refresh_token, expires_at) - .await?; + let refresh_token = generate_refresh_token(); + self.create_refresh_token(id, &refresh_token).await?; let Ok(access_token) = generate_jwt(Claims::new(UserInfo::new( user.email.clone(), @@ -177,9 +176,8 @@ impl AuthenticationService for DbConn { return Err(TokenAuthError::InvalidPassword); } - let (refresh_token, expires_at) = generate_refresh_token(chrono::Utc::now().timestamp()); - self.create_refresh_token(user.id, &refresh_token, expires_at) - .await?; + let refresh_token = generate_refresh_token(); + self.create_refresh_token(user.id, &refresh_token).await?; let Ok(access_token) = generate_jwt(Claims::new(UserInfo::new( user.email.clone(), @@ -206,7 +204,7 @@ impl AuthenticationService for DbConn { return Err(RefreshTokenError::UserNotFound); }; - let (new_token, _) = generate_refresh_token(chrono::Utc::now().timestamp()); + let new_token = generate_refresh_token(); self.replace_refresh_token(&token, &new_token).await?; // refresh token update is done, generate new access token based on user info @@ -217,8 +215,7 @@ impl AuthenticationService for DbConn { return Err(RefreshTokenError::Unknown); }; - let resp = - RefreshTokenResponse::new(access_token, new_token, refresh_token.expires_at as f64); + let resp = RefreshTokenResponse::new(access_token, new_token, refresh_token.expires_at); Ok(resp) } diff --git a/ee/tabby-webserver/src/service/db.rs b/ee/tabby-webserver/src/service/db.rs index 7941576ef444..4250e6315d4c 100644 --- a/ee/tabby-webserver/src/service/db.rs +++ b/ee/tabby-webserver/src/service/db.rs @@ -1,6 +1,7 @@ use std::{path::PathBuf, sync::Arc}; use anyhow::{anyhow, Result}; +use chrono::{DateTime, Utc}; use lazy_static::lazy_static; use rusqlite::{params, OptionalExtension, Row}; use rusqlite_migration::{AsyncMigrations, M}; @@ -57,7 +58,7 @@ lazy_static! { id INTEGER PRIMARY KEY AUTOINCREMENT, user_id INTEGER NOT NULL, token VARCHAR(255) NOT NULL COLLATE NOCASE, - expires_at INTEGER NOT NULL, + expires_at TIMESTAMP NOT NULL, created_at TIMESTAMP DEFAULT (DATETIME('now')), CONSTRAINT `idx_token` UNIQUE (`token`) ); @@ -69,8 +70,8 @@ lazy_static! { #[allow(unused)] pub struct User { - created_at: String, - updated_at: String, + created_at: DateTime, + updated_at: DateTime, pub id: i32, pub email: String, @@ -328,11 +329,11 @@ impl DbConn { #[allow(unused)] pub struct RefreshToken { id: u32, - created_at: String, + created_at: DateTime, pub user_id: i32, pub token: String, - pub expires_at: i64, + pub expires_at: DateTime, } impl RefreshToken { @@ -352,26 +353,21 @@ impl RefreshToken { } pub fn is_expired(&self) -> bool { - let now = chrono::Utc::now().timestamp(); + let now = chrono::Utc::now(); self.expires_at < now } } /// db read/write operations for `refresh_tokens` table impl DbConn { - pub async fn create_refresh_token( - &self, - user_id: i32, - token: &str, - expires_at: i64, - ) -> Result<()> { + pub async fn create_refresh_token(&self, user_id: i32, token: &str) -> Result<()> { let token = token.to_string(); let res = self .conn .call(move |c| { c.execute( - r#"INSERT INTO refresh_tokens (user_id, token, expires_at) VALUES (?, ?, ?)"#, - params![user_id, token, expires_at], + r#"INSERT INTO refresh_tokens (user_id, token, expires_at) VALUES (?, ?, datetime('now', '+7 days'))"#, + params![user_id, token], ) }) .await?; @@ -436,6 +432,8 @@ impl DbConn { #[cfg(test)] mod tests { + use std::ops::Add; + use super::*; use crate::schema::auth::AuthenticationService; @@ -527,20 +525,21 @@ mod tests { async fn test_create_refresh_token() { let conn = DbConn::new_in_memory().await.unwrap(); - conn.create_refresh_token(1, "test", 100).await.unwrap(); + conn.create_refresh_token(1, "test").await.unwrap(); let token = conn.get_refresh_token("test").await.unwrap().unwrap(); assert_eq!(token.user_id, 1); assert_eq!(token.token, "test"); - assert_eq!(token.expires_at, 100); + assert!(token.expires_at > Utc::now().add(chrono::Duration::days(6))); + assert!(token.expires_at < Utc::now().add(chrono::Duration::days(7))); } #[tokio::test] async fn test_replace_refresh_token() { let conn = DbConn::new_in_memory().await.unwrap(); - conn.create_refresh_token(1, "test", 100).await.unwrap(); + conn.create_refresh_token(1, "test").await.unwrap(); conn.replace_refresh_token("test", "test2").await.unwrap(); let token = conn.get_refresh_token("test").await.unwrap(); @@ -549,6 +548,5 @@ mod tests { let token = conn.get_refresh_token("test2").await.unwrap().unwrap(); assert_eq!(token.user_id, 1); assert_eq!(token.token, "test2"); - assert_eq!(token.expires_at, 100); } }