Browse Source

moved cache from static global

feature/settings-cleanup
John Doe 2 years ago
parent
commit
48bdacdebc
  1. 5
      server/src/api/user.rs
  2. 13
      server/src/lib.rs
  3. 5
      server/src/main.rs
  4. 125
      server/src/twofactor.rs
  5. 9
      server/src/websocket/chat_server.rs

5
server/src/api/user.rs

@ -30,7 +30,6 @@ use crate::{
blocking, captcha_espeak_wav_base64,
hcaptcha::hcaptcha_verify,
is_within_message_char_limit,
twofactor::{check_2fa, generate_2fa},
websocket::{
messages::{CaptchaItem, CheckCaptcha, JoinUserRoom, SendAllMessage, SendUserRoomMessage},
UserOperation,
@ -153,7 +152,7 @@ impl Perform for Login {
//handle 2fa
if user.has_2fa {
match &data.code_2fa {
Some(code) => match check_2fa(&user, code) {
Some(code) => match context.code_cache_2fa().check_2fa(&user, code) {
Ok(matches) => {
if matches {
return Ok(LoginResponse {
@ -166,7 +165,7 @@ impl Perform for Login {
Err(e) => return Err(e),
},
None => {
match generate_2fa(user) {
match context.code_cache_2fa().generate_2fa(user) {
Ok(_k) => (),
Err(e) => return Err(e),
};

13
server/src/lib.rs

@ -33,6 +33,7 @@ pub mod websocket;
use crate::{
request::{retry, RecvError},
twofactor::CodeCacheHandler,
websocket::chat_server::ChatServer,
};
use actix::Addr;
@ -43,7 +44,7 @@ use log::error;
use percent_encoding::{utf8_percent_encode, NON_ALPHANUMERIC};
use reqwest::Client;
use serde::Deserialize;
use std::process::Command;
use std::{process::Command, sync::Arc};
pub type DbPool = diesel::r2d2::Pool<diesel::r2d2::ConnectionManager<diesel::PgConnection>>;
@ -52,6 +53,7 @@ pub struct LemmyContext {
pub chat_server: Addr<ChatServer>,
pub client: Client,
pub activity_queue: QueueHandle,
pub cache_handler: Arc<CodeCacheHandler>,
}
impl LemmyContext {
@ -60,12 +62,14 @@ impl LemmyContext {
chat_server: Addr<ChatServer>,
client: Client,
activity_queue: QueueHandle,
cache_handler: Arc<CodeCacheHandler>,
) -> LemmyContext {
LemmyContext {
pool,
chat_server,
client,
activity_queue,
cache_handler,
}
}
pub fn pool(&self) -> &DbPool {
@ -80,9 +84,12 @@ impl LemmyContext {
pub fn activity_queue(&self) -> &QueueHandle {
&self.activity_queue
}
pub fn code_cache_2fa(&self) -> &CodeCacheHandler {
&self.cache_handler
}
}
impl Clone for LemmyContext {
/*impl Clone for LemmyContext {
fn clone(&self) -> Self {
LemmyContext {
pool: self.pool.clone(),
@ -91,7 +98,7 @@ impl Clone for LemmyContext {
activity_queue: self.activity_queue.clone(),
}
}
}
}*/
#[derive(Deserialize, Debug)]
pub struct IframelyResponse {

5
server/src/main.rs

@ -20,7 +20,7 @@ use diesel::{
use lemmy_db::get_database_url_from_env;
use lemmy_rate_limit::{rate_limiter::RateLimiter, RateLimit};
use lemmy_server::{
apub::activity_queue::create_activity_queue, blocking, routes::*,
apub::activity_queue::create_activity_queue, blocking, routes::*, twofactor::CodeCacheHandler,
websocket::chat_server::ChatServer, LemmyContext,
};
use lemmy_utils::{settings::Settings, LemmyError, CACHE_CONTROL_REGEX};
@ -72,11 +72,13 @@ async fn main() -> Result<(), LemmyError> {
);
let activity_queue = create_activity_queue();
let cache_handler = Arc::new(CodeCacheHandler::new());
let chat_server = ChatServer::startup(
pool.clone(),
rate_limiter.clone(),
Client::default(),
activity_queue.clone(),
cache_handler.clone(),
)
.start();
@ -87,6 +89,7 @@ async fn main() -> Result<(), LemmyError> {
chat_server.to_owned(),
Client::default(),
activity_queue.to_owned(),
cache_handler.clone(),
);
let settings = Settings::get();
let rate_limiter = rate_limiter.clone();

125
server/src/twofactor.rs

@ -5,79 +5,86 @@ use lemmy_utils::{send_email, LemmyError};
use std::{sync::Mutex, time::Duration};
use chrono::prelude::*;
use lazy_static::lazy_static;
use rand::seq::SliceRandom;
use ttl_cache::TtlCache;
lazy_static! { //no real good way to synchronously pass the cache between all login requests without an ugly static with inner mutability
static ref CODE_CACHE: Mutex<TtlCache<String, User_>> = Mutex::new(TtlCache::new(100)); //making me choose a max number of items
}
const ALLOWED_CODE_CHARS: &[u8] = b"0123456789";
const CODE_LENGTH: usize = 8;
pub fn generate_2fa(user: User_) -> Result<(), LemmyError> {
if !user.has_2fa || user.email.is_none() {
return Err(APIError::err("user_has_no_2fa").into());
pub struct CodeCacheHandler {
cache: Mutex<TtlCache<String, User_>>,
}
impl CodeCacheHandler {
pub fn new() -> CodeCacheHandler {
CodeCacheHandler {
cache: Mutex::new(TtlCache::new(100)),
}
}
let mut genned_code;
{
//inner scope to unlock code cache once we write the code
let mut code_cache = match CODE_CACHE.lock() {
Ok(k) => k,
Err(_e) => return Err(APIError::err("internal_error").into()),
};
let mut rng = rand::thread_rng();
loop {
genned_code = String::from("");
for _ in 0..=CODE_LENGTH {
genned_code.push(
ALLOWED_CODE_CHARS
.choose(&mut rng)
.map(|&c| c as char)
.unwrap(),
);
}
if code_cache.get(genned_code.as_str()).is_none() {
//break if this is a unique 2fa code
break;
}
pub fn generate_2fa(&self, user: User_) -> Result<(), LemmyError> {
if !user.has_2fa || user.email.is_none() {
return Err(APIError::err("user_has_no_2fa").into());
}
code_cache.insert(genned_code.clone(), user.clone(), Duration::from_secs(3600));
//code is valid for one hour
}
let mut genned_code;
{
//inner scope to unlock code cache once we write the code
let mut code_cache = match self.cache.lock() {
Ok(k) => k,
Err(_e) => return Err(APIError::err("internal_error").into()),
};
let mut rng = rand::thread_rng();
loop {
genned_code = String::from("");
for _ in 0..=CODE_LENGTH {
genned_code.push(
ALLOWED_CODE_CHARS
.choose(&mut rng)
.map(|&c| c as char)
.unwrap(),
);
}
if code_cache.get(genned_code.as_str()).is_none() {
//break if this is a unique 2fa code
break;
}
}
let subject = &format!("ChapoChat: Attempted login for {}", &user.name);
let time = Utc::now().format("%Y-%m-%d %H:%M:%S");
let html = &format!("<h1>Attempted login for {}</h1><br><p>At {} UTC a login was attempted on your account.
Because your account is setup with two-factor authentication, you must enter a code to successfully login. This code will expire within one hour.</p>
<h3>Your login code is {}</h3>", user.name, time, genned_code);
//println!("Sending 2fa email with code {}", genned_code);
match send_email(subject, user.email.unwrap().as_str(), &user.name, html) {
Ok(_k) => (),
Err(e) => println!("Failed to send email: {}", e),
}
Ok(())
}
code_cache.insert(genned_code.clone(), user.clone(), Duration::from_secs(3600));
//code is valid for one hour
}
pub fn check_2fa(user: &User_, code: &String) -> Result<bool, LemmyError> {
if !user.has_2fa || user.email.is_none() {
return Err(APIError::err("user_has_no_2fa").into());
let subject = &format!("ChapoChat: Attempted login for {}", &user.name);
let time = Utc::now().format("%Y-%m-%d %H:%M:%S");
let html = &format!("<h1>Attempted login for {}</h1><br><p>At {} UTC a login was attempted on your account.
Because your account is setup with two-factor authentication, you must enter a code to successfully login. This code will expire within one hour.</p>
<h3>Your login code is {}</h3>", user.name, time, genned_code);
println!("Sending 2fa email with code {}", genned_code);
match send_email(subject, user.email.unwrap().as_str(), &user.name, html) {
Ok(_k) => (),
Err(e) => println!("Failed to send email: {}", e),
}
Ok(())
}
let mut code_cache = match CODE_CACHE.lock() {
Ok(k) => k,
Err(_e) => return Err(APIError::err("internal_error").into()),
};
match code_cache.get(code) {
Some(cached_user) => {
if cached_user.id == user.id {
code_cache.remove(code);
return Ok(true); //code exists and the user did request it
pub fn check_2fa(&self, user: &User_, code: &String) -> Result<bool, LemmyError> {
if !user.has_2fa || user.email.is_none() {
return Err(APIError::err("user_has_no_2fa").into());
}
let mut code_cache = match self.cache.lock() {
Ok(k) => k,
Err(_e) => return Err(APIError::err("internal_error").into()),
};
println!("Entered code {}", code);
match code_cache.get(code) {
Some(cached_user) => {
if cached_user.id == user.id {
code_cache.remove(code);
return Ok(true); //code exists and the user did request it
}
Ok(false) //the code exists, but the user wasn't the one who requested it
}
Ok(false) //the code exists, but the user wasn't the one who requested it
None => Ok(false), //no matching code
}
None => Ok(false), //no matching code
}
}

9
server/src/websocket/chat_server.rs

@ -1,4 +1,5 @@
use crate::{
twofactor::CodeCacheHandler,
websocket::{
handlers::{do_user_operation, to_json_string, Args},
messages::*,
@ -26,6 +27,7 @@ use serde_json::Value;
use std::{
collections::{HashMap, HashSet},
str::FromStr,
sync::Arc,
};
/// `ChatServer` manages chat rooms and responsible for coordinating chat
@ -55,6 +57,9 @@ pub struct ChatServer {
/// A list of the current captchas
pub(super) captchas: Vec<CaptchaItem>,
//A time-sensitive list of two-factor auth codes
pub cache_handler: Arc<CodeCacheHandler>,
/// An HTTP Client
client: Client,
@ -75,6 +80,7 @@ impl ChatServer {
rate_limiter: RateLimit,
client: Client,
activity_queue: QueueHandle,
cache_handler: Arc<CodeCacheHandler>,
) -> ChatServer {
ChatServer {
sessions: HashMap::new(),
@ -87,6 +93,7 @@ impl ChatServer {
captchas: Vec::new(),
client,
activity_queue,
cache_handler,
}
}
@ -341,6 +348,7 @@ impl ChatServer {
let addr = ctx.address();
let pool = self.pool.clone();
let rate_limiter = self.rate_limiter.clone();
let cache_handler = self.cache_handler.clone();
let ip: IPAddr = match self.sessions.get(&msg.id) {
Some(info) => info.ip.to_owned(),
@ -364,6 +372,7 @@ impl ChatServer {
chat_server: addr,
client,
activity_queue,
cache_handler,
};
let args = Args {
context,

Loading…
Cancel
Save