Adding TOTP / 2FA to lemmy

crates/api/src/local_user/login.rs

@@ -6,9 +6,13 @@ use lemmy_api_common::{
   person::{Login, LoginResponse},
   utils::{check_registration_application, check_user_valid},
 };
-use lemmy_db_schema::source::local_site::LocalSite;
-use lemmy_db_views::structs::LocalUserView;
-use lemmy_utils::{claims::Claims, error::LemmyError, ConnectionId};
+use lemmy_db_views::structs::{LocalUserView, SiteView};
+use lemmy_utils::{
+  claims::Claims,
+  error::LemmyError,
+  utils::validation::check_totp_valid,
+  ConnectionId,
+};
 
 #[async_trait::async_trait(?Send)]
 impl Perform for Login {
@@ -22,7 +26,7 @@ impl Perform for Login {
   ) -> Result<LoginResponse, LemmyError> {
     let data: &Login = self;
 
-    let local_site = LocalSite::read(context.pool()).await?;
+    let site_view = SiteView::read_local(context.pool()).await?;
 
     // Fetch that username / email
     let username_or_email = data.username_or_email.clone();
@@ -45,11 +49,20 @@ impl Perform for Login {
       local_user_view.person.deleted,
     )?;
 
-    if local_site.require_email_verification && !local_user_view.local_user.email_verified {
+    if site_view.local_site.require_email_verification && !local_user_view.local_user.email_verified
+    {
       return Err(LemmyError::from_message("email_not_verified"));
     }
 
-    check_registration_application(&local_user_view, &local_site, context.pool()).await?;
+    check_registration_application(&local_user_view, &site_view.local_site, context.pool()).await?;
+
+    // Check the totp
+    check_totp_valid(
+      &local_user_view.local_user.totp_secret,
+      &data.totp_token,
+      &site_view.site.name,
+      &local_user_view.person.name,
+    )?;
 
     // Return the jwt
     Ok(LoginResponse {

crates/api/src/local_user/save_settings.rs

@@ -8,17 +8,22 @@ use lemmy_api_common::{
 use lemmy_db_schema::{
   source::{
     actor_language::LocalUserLanguage,
-    local_site::LocalSite,
     local_user::{LocalUser, LocalUserUpdateForm},
     person::{Person, PersonUpdateForm},
   },
   traits::Crud,
   utils::{diesel_option_overwrite, diesel_option_overwrite_to_url},
 };
+use lemmy_db_views::structs::SiteView;
 use lemmy_utils::{
   claims::Claims,
   error::LemmyError,
-  utils::validation::{is_valid_display_name, is_valid_matrix_id},
+  utils::validation::{
+    build_totp,
+    generate_totp_secret,
+    is_valid_display_name,
+    is_valid_matrix_id,
+  },
   ConnectionId,
 };
 
@@ -35,14 +40,13 @@ impl Perform for SaveUserSettings {
     let data: &SaveUserSettings = self;
     let local_user_view =
       get_local_user_view_from_jwt(&data.auth, context.pool(), context.secret()).await?;
-    let local_site = LocalSite::read(context.pool()).await?;
+    let site_view = SiteView::read_local(context.pool()).await?;
 
     let avatar = diesel_option_overwrite_to_url(&data.avatar)?;
     let banner = diesel_option_overwrite_to_url(&data.banner)?;
     let bio = diesel_option_overwrite(&data.bio);
     let display_name = diesel_option_overwrite(&data.display_name);
     let matrix_user_id = diesel_option_overwrite(&data.matrix_user_id);
-    let bot_account = data.bot_account;
     let email_deref = data.email.as_deref().map(str::to_lowercase);
     let email = diesel_option_overwrite(&email_deref);
 
@@ -57,7 +61,7 @@ impl Perform for SaveUserSettings {
 
     // When the site requires email, make sure email is not Some(None). IE, an overwrite to a None value
     if let Some(email) = &email {
-      if email.is_none() && local_site.require_email_verification {
+      if email.is_none() && site_view.local_site.require_email_verification {
         return Err(LemmyError::from_message("email_required"));
       }
     }
@@ -71,7 +75,7 @@ impl Perform for SaveUserSettings {
     if let Some(Some(display_name)) = &display_name {
       if !is_valid_display_name(
         display_name.trim(),
-        local_site.actor_name_max_length as usize,
+        site_view.local_site.actor_name_max_length as usize,
       ) {
         return Err(LemmyError::from_message("invalid_username"));
       }
@@ -92,7 +96,7 @@ impl Perform for SaveUserSettings {
       .display_name(display_name)
       .bio(bio)
       .matrix_user_id(matrix_user_id)
-      .bot_account(bot_account)
+      .bot_account(data.bot_account)
       .avatar(avatar)
       .banner(banner)
       .build();
@@ -105,6 +109,20 @@ impl Perform for SaveUserSettings {
       LocalUserLanguage::update(context.pool(), discussion_languages, local_user_id).await?;
     }
 
+    // If generate_totp is Some(false), this will clear it out from the database.
+    let (totp_secret, totp_url) = if let Some(generate) = data.generate_totp {
+      if generate {
+        let secret = generate_totp_secret();
+        let url =
+          build_totp(&site_view.site.name, &local_user_view.person.name, &secret)?.get_url();
+        (Some(Some(secret)), Some(Some(url)))
+      } else {
+        (Some(None), Some(None))
+      }
+    } else {
+      (None, None)
+    };
+
     let local_user_form = LocalUserUpdateForm::builder()
       .email(email)
       .show_avatars(data.show_avatars)
@@ -118,6 +136,8 @@ impl Perform for SaveUserSettings {
       .default_listing_type(default_listing_type)
       .theme(data.theme.clone())
       .interface_language(data.interface_language.clone())
+      .totp_secret(totp_secret)
+      .totp_url(totp_url)
       .build();
 
     let local_user_res = LocalUser::update(context.pool(), local_user_id, &local_user_form).await;

crates/api_common/src/person.rs

@@ -17,6 +17,7 @@ use serde::{Deserialize, Serialize};
 pub struct Login {
   pub username_or_email: Sensitive<String>,
   pub password: Sensitive<String>,
+  pub totp_token: Option<String>,
 }
 
 #[derive(Debug, Serialize, Deserialize, Clone, Default)]
@@ -70,6 +71,8 @@ pub struct SaveUserSettings {
   pub show_read_posts: Option<bool>,
   pub show_new_post_notifs: Option<bool>,
   pub discussion_languages: Option<Vec<LanguageId>>,
+  /// None leaves it as is, true will generate or regenerate it, false clears it out
+  pub generate_totp: Option<bool>,
   pub auth: Sensitive<String>,
 }
 

crates/db_schema/src/impls/local_user.rs

@@ -37,6 +37,8 @@ mod safe_settings_type {
       show_read_posts,
       show_scores,
       theme,
+      totp_secret,
+      totp_url,
       validator_time,
     },
     source::local_user::LocalUser,
@@ -61,6 +63,8 @@ mod safe_settings_type {
     show_new_post_notifs,
     email_verified,
     accepted_application,
+    totp_secret,
+    totp_url,
   );
 
   impl ToSafeSettings for LocalUser {
@@ -86,6 +90,8 @@ mod safe_settings_type {
         show_new_post_notifs,
         email_verified,
         accepted_application,
+        totp_secret,
+        totp_url,
       )
     }
   }

crates/db_schema/src/schema.rs

@@ -168,6 +168,8 @@ table! {
         show_new_post_notifs -> Bool,
         email_verified -> Bool,
         accepted_application -> Bool,
+        totp_secret -> Nullable<Text>,
+        totp_url -> Nullable<Text>,
     }
 }
 

crates/db_schema/src/source/local_user.rs

@@ -26,6 +26,9 @@ pub struct LocalUser {
   pub show_new_post_notifs: bool,
   pub email_verified: bool,
   pub accepted_application: bool,
+  #[serde(skip)]
+  pub totp_secret: Option<String>,
+  pub totp_url: Option<String>,
 }
 
 /// A local user view that removes password encrypted
@@ -50,6 +53,9 @@ pub struct LocalUserSettings {
   pub show_new_post_notifs: bool,
   pub email_verified: bool,
   pub accepted_application: bool,
+  #[serde(skip)]
+  pub totp_secret: Option<String>,
+  pub totp_url: Option<String>,
 }
 
 #[derive(Clone, TypedBuilder)]
@@ -75,6 +81,8 @@ pub struct LocalUserInsertForm {
   pub show_new_post_notifs: Option<bool>,
   pub email_verified: Option<bool>,
   pub accepted_application: Option<bool>,
+  pub totp_secret: Option<Option<String>>,
+  pub totp_url: Option<Option<String>>,
 }
 
 #[derive(Clone, TypedBuilder)]
@@ -97,4 +105,6 @@ pub struct LocalUserUpdateForm {
   pub show_new_post_notifs: Option<bool>,
   pub email_verified: Option<bool>,
   pub accepted_application: Option<bool>,
+  pub totp_secret: Option<Option<String>>,
+  pub totp_url: Option<Option<String>>,
 }

crates/db_views/src/registration_application_view.rs

@@ -288,6 +288,8 @@ mod tests {
         show_new_post_notifs: inserted_sara_local_user.show_new_post_notifs,
         email_verified: inserted_sara_local_user.email_verified,
         accepted_application: inserted_sara_local_user.accepted_application,
+        totp_secret: inserted_sara_local_user.totp_secret,
+        totp_url: inserted_sara_local_user.totp_url,
       },
       creator: PersonSafe {
         id: inserted_sara_person.id,

crates/utils/Cargo.toml

@@ -46,6 +46,7 @@ smart-default = "0.6.0"
 jsonwebtoken = "8.1.1"
 lettre = "0.10.1"
 comrak = { version = "0.14.0", default-features = false }
+totp-rs = { version = "4.2.0", features = ["gen_secret", "otpauth"] }
 
 [build-dependencies]
 rosetta-build = "0.1.2"