From 79b22e350160fe810f174a13b960ab5eebb2a9b9 Mon Sep 17 00:00:00 2001 From: lifegpc Date: Wed, 28 Sep 2022 09:22:33 +0000 Subject: [PATCH] Update --- src/db/sqlite/db.rs | 68 +++++++++++++++++++++++++++++++++++-- src/db/token.rs | 2 +- src/db/traits.rs | 16 +++++++++ src/server/auth/token.rs | 21 +++++++++++- src/server/context.rs | 2 +- src/server/unittest/auth.rs | 14 +++++++- 6 files changed, 116 insertions(+), 7 deletions(-) diff --git a/src/db/sqlite/db.rs b/src/db/sqlite/db.rs index cae245f..7d99d55 100644 --- a/src/db/sqlite/db.rs +++ b/src/db/sqlite/db.rs @@ -5,6 +5,7 @@ use super::super::{ use super::super::{Token, User}; use super::SqliteError; use bytes::BytesMut; +use chrono::{DateTime, Utc}; use futures_util::lock::Mutex; use rusqlite::{Connection, OpenFlags, OptionalExtension, Transaction}; use std::collections::HashMap; @@ -56,12 +57,11 @@ lang TEXT, translated TEXT );"; const TOKEN_TABLE: &'static str = "CREATE TABLE token ( -id INT, +id INTEGER PRIMARY KEY AUTOINCREMENT, user_id INT, token TEXT, created_at DATETIME, -expired_at DATETIME, -PRIMARY KEY (id) +expired_at DATETIME );"; const USERS_TABLE: &'static str = "CREATE TABLE users ( id INT, @@ -103,6 +103,20 @@ impl PixivDownloaderSqlite { Ok(()) } + fn _add_token( + tx: &Transaction, + user_id: u64, + token: &[u8; 64], + created_at: &DateTime, + expired_at: &DateTime, + ) -> Result<(), SqliteError> { + tx.execute( + "INSERT INTO token (user_id, token, created_at, expired_at) VALUES (?, ?, ?, ?);", + (user_id, token, created_at, expired_at), + )?; + Ok(()) + } + fn _add_user( tx: &Transaction, name: &str, @@ -224,6 +238,30 @@ impl PixivDownloaderSqlite { .optional()?) } + #[cfg(feature = "server")] + async fn _get_token_by_user_id_and_token( + &self, + user_id: u64, + token: &[u8; 64], + ) -> Result, SqliteError> { + let con = self.db.lock().await; + Ok(con + .query_row( + "SELECT * FROM token WHERE user_id = ? AND token = ?;", + (user_id, token), + |row| { + Ok(Token { + id: row.get(0)?, + user_id: row.get(1)?, + token: row.get(2)?, + created_at: row.get(3)?, + expired_at: row.get(4)?, + }) + }, + ) + .optional()?) + } + #[cfg(feature = "server")] async fn _get_user(&self, id: u64) -> Result, SqliteError> { let con = self.db.lock().await; @@ -401,6 +439,30 @@ impl PixivDownloaderDb for PixivDownloaderSqlite { Ok(self.get_user(0).await?.expect("Root user not found:")) } + #[cfg(feature = "server")] + async fn add_token( + &self, + user_id: u64, + token: &[u8; 64], + created_at: &DateTime, + expired_at: &DateTime, + ) -> Result, PixivDownloaderDbError> { + if self + ._get_token_by_user_id_and_token(user_id, token) + .await? + .is_some() + { + return Ok(None); + } + { + let mut db = self.db.lock().await; + let mut tx = db.transaction()?; + Self::_add_token(&mut tx, user_id, token, created_at, expired_at)?; + tx.commit()?; + } + Ok(self._get_token_by_user_id_and_token(user_id, token).await?) + } + #[cfg(feature = "server")] async fn add_user( &self, diff --git a/src/db/token.rs b/src/db/token.rs index 3abf7be..5ce4fad 100644 --- a/src/db/token.rs +++ b/src/db/token.rs @@ -7,7 +7,7 @@ pub struct Token { /// The user ID of the token pub user_id: u64, /// The token - pub token: String, + pub token: [u8; 64], /// The token's creation time pub created_at: DateTime, /// The token's expiration time diff --git a/src/db/traits.rs b/src/db/traits.rs index 54ebe2a..3252954 100644 --- a/src/db/traits.rs +++ b/src/db/traits.rs @@ -2,6 +2,7 @@ use super::PixivDownloaderDbConfig; use super::PixivDownloaderDbError; #[cfg(feature = "server")] use super::{Token, User}; +use chrono::{DateTime, Utc}; #[async_trait] pub trait PixivDownloaderDb { @@ -24,6 +25,21 @@ pub trait PixivDownloaderDb { password: &[u8], ) -> Result; #[cfg(feature = "server")] + /// Add a new token + /// * `user_id` - The user ID + /// * `token` - The token + /// * `created_at` - The token's expiration time + /// * `expired_at` - The token's creation time + /// # Note + /// if a token with the same user_id already exists, must return None + async fn add_token( + &self, + user_id: u64, + token: &[u8; 64], + created_at: &DateTime, + expired_at: &DateTime, + ) -> Result, PixivDownloaderDbError>; + #[cfg(feature = "server")] /// Add a new user to database. /// * `name` - User name /// * `username` - Unique user name diff --git a/src/server/auth/token.rs b/src/server/auth/token.rs index 79935d5..513bb3a 100644 --- a/src/server/auth/token.rs +++ b/src/server/auth/token.rs @@ -73,7 +73,26 @@ impl AuthTokenContext { if pass != &hashed_password { return Err((9, gettext("Wrong password.")).into()); } - Ok(json::object! {}) + let mut token = [0; 64]; + openssl::rand::rand_bytes(&mut token) + .try_err3(-1003, gettext("Failed to generate token:"))?; + let created_at = chrono::Utc::now(); + let expired_at = created_at + chrono::Duration::days(30); + let token = loop { + if let Some(token) = self + .ctx + .db + .add_token(user.id, &token, &created_at, &expired_at) + .await + .try_err3(-1001, gettext("Failed to operate the database:"))? + { + break token; + } + }; + let b64token = base64::encode(&token.token); + Ok( + json::object! { "id": token.id, "user_id": token.user_id, "token": b64token, "created_at": token.created_at.timestamp(), "expired_at": token.expired_at.timestamp() }, + ) } }, None => { diff --git a/src/server/context.rs b/src/server/context.rs index df259aa..70c4576 100644 --- a/src/server/context.rs +++ b/src/server/context.rs @@ -103,7 +103,7 @@ impl ServerContext { par.insert(k, v); } let mut sha = openssl::sha::Sha512::new(); - sha.update(token.token.as_bytes()); + sha.update(&token.token); for (k, v) in par { for v in v { sha.update(k.as_bytes()); diff --git a/src/server/unittest/auth.rs b/src/server/unittest/auth.rs index 2566400..81132e2 100644 --- a/src/server/unittest/auth.rs +++ b/src/server/unittest/auth.rs @@ -31,7 +31,7 @@ pub async fn test(ctx: &UnitTestContext) -> Result "test", "name" => "test", - "password" => b64_password, + "password" => b64_password.as_str(), }, ) .await? @@ -46,5 +46,17 @@ pub async fn test(ctx: &UnitTestContext) -> Result