Rework the way 2FA is enabled/disabled (fixes #3309) (#3959)

* Rework the way 2FA is enabled/disabled (fixes #3309)

* postgres format

* change algo to sha1 for better compat

* review comments

* review

* clippy

---------

Co-authored-by: Dessalines <dessalines@users.noreply.github.com>

Cargo.lock

@@ -2627,6 +2627,7 @@ dependencies = [
  "serial_test",
  "sitemap-rs",
  "tokio",
+ "totp-rs",
  "tracing",
  "url",
  "uuid",
@@ -2937,7 +2938,6 @@ dependencies = [
  "strum",
  "strum_macros",
  "tokio",
- "totp-rs",
  "tracing",
  "tracing-error",
  "ts-rs",

crates/api/Cargo.toml

@@ -34,6 +34,7 @@ chrono = { workspace = true }
 url = { workspace = true }
 wav = "1.0.0"
 sitemap-rs = "0.2.0"
+totp-rs = { version = "5.0.2", features = ["gen_secret", "otpauth"] }
 
 [dev-dependencies]
 serial_test = { workspace = true }

crates/api/src/lib.rs

@@ -2,11 +2,13 @@ use base64::{engine::general_purpose::STANDARD_NO_PAD as base64, Engine};
 use captcha::Captcha;
 use lemmy_api_common::utils::local_site_to_slur_regex;
 use lemmy_db_schema::source::local_site::LocalSite;
+use lemmy_db_views::structs::LocalUserView;
 use lemmy_utils::{
-  error::{LemmyError, LemmyErrorExt, LemmyErrorType},
+  error::{LemmyError, LemmyErrorExt, LemmyErrorType, LemmyResult},
   utils::slurs::check_slurs,
 };
 use std::io::Cursor;
+use totp_rs::{Secret, TOTP};
 
 pub mod comment;
 pub mod comment_report;
@@ -67,11 +69,63 @@ pub(crate) fn check_report_reason(reason: &str, local_site: &LocalSite) -> Resul
   }
 }
 
+pub(crate) fn check_totp_2fa_valid(
+  local_user_view: &LocalUserView,
+  totp_token: &Option<String>,
+  site_name: &str,
+) -> LemmyResult<()> {
+  // Throw an error if their token is missing
+  let token = totp_token
+    .as_deref()
+    .ok_or(LemmyErrorType::MissingTotpToken)?;
+  let secret = local_user_view
+    .local_user
+    .totp_2fa_secret
+    .as_deref()
+    .ok_or(LemmyErrorType::MissingTotpSecret)?;
+
+  let totp = build_totp_2fa(site_name, &local_user_view.person.name, secret)?;
+
+  let check_passed = totp.check_current(token)?;
+  if !check_passed {
+    return Err(LemmyErrorType::IncorrectTotpToken.into());
+  }
+
+  Ok(())
+}
+
+pub(crate) fn generate_totp_2fa_secret() -> String {
+  Secret::generate_secret().to_string()
+}
+
+pub(crate) fn build_totp_2fa(
+  site_name: &str,
+  username: &str,
+  secret: &str,
+) -> Result<TOTP, LemmyError> {
+  let sec = Secret::Raw(secret.as_bytes().to_vec());
+  let sec_bytes = sec
+    .to_bytes()
+    .map_err(|_| LemmyErrorType::CouldntParseTotpSecret)?;
+
+  TOTP::new(
+    totp_rs::Algorithm::SHA1,
+    6,
+    1,
+    30,
+    sec_bytes,
+    Some(site_name.to_string()),
+    username.to_string(),
+  )
+  .with_lemmy_type(LemmyErrorType::CouldntGenerateTotp)
+}
+
 #[cfg(test)]
 mod tests {
   #![allow(clippy::unwrap_used)]
   #![allow(clippy::indexing_slicing)]
 
+  use super::*;
   use lemmy_api_common::utils::check_validator_time;
   use lemmy_db_schema::{
     source::{
@@ -134,4 +188,11 @@ mod tests {
     let num_deleted = Person::delete(pool, inserted_person.id).await.unwrap();
     assert_eq!(1, num_deleted);
   }
+
+  #[test]
+  fn test_build_totp() {
+    let generated_secret = generate_totp_2fa_secret();
+    let totp = build_totp_2fa("lemmy", "my_name", &generated_secret);
+    assert!(totp.is_ok());
+  }
 }

crates/api/src/local_user/generate_totp_secret.rs

@@ -0,0 +1,47 @@
+use crate::{build_totp_2fa, generate_totp_2fa_secret};
+use activitypub_federation::config::Data;
+use actix_web::web::Json;
+use lemmy_api_common::{
+  context::LemmyContext,
+  person::GenerateTotpSecretResponse,
+  sensitive::Sensitive,
+};
+use lemmy_db_schema::{
+  source::local_user::{LocalUser, LocalUserUpdateForm},
+  traits::Crud,
+};
+use lemmy_db_views::structs::{LocalUserView, SiteView};
+use lemmy_utils::error::{LemmyError, LemmyErrorType};
+
+/// Generate a new secret for two-factor-authentication. Afterwards you need to call [toggle_totp]
+/// to enable it. This can only be called if 2FA is currently disabled.
+#[tracing::instrument(skip(context))]
+pub async fn generate_totp_secret(
+  local_user_view: LocalUserView,
+  context: Data<LemmyContext>,
+) -> Result<Json<GenerateTotpSecretResponse>, LemmyError> {
+  let site_view = SiteView::read_local(&mut context.pool()).await?;
+
+  if local_user_view.local_user.totp_2fa_enabled {
+    return Err(LemmyErrorType::TotpAlreadyEnabled)?;
+  }
+
+  let secret = generate_totp_2fa_secret();
+  let secret_url =
+    build_totp_2fa(&site_view.site.name, &local_user_view.person.name, &secret)?.get_url();
+
+  let local_user_form = LocalUserUpdateForm {
+    totp_2fa_secret: Some(Some(secret)),
+    ..Default::default()
+  };
+  LocalUser::update(
+    &mut context.pool(),
+    local_user_view.local_user.id,
+    &local_user_form,
+  )
+  .await?;
+
+  Ok(Json(GenerateTotpSecretResponse {
+    totp_secret_url: Sensitive::new(secret_url),
+  }))
+}

crates/api/src/local_user/login.rs

@@ -1,3 +1,4 @@
+use crate::check_totp_2fa_valid;
 use actix_web::web::{Data, Json};
 use bcrypt::verify;
 use lemmy_api_common::{
@@ -9,7 +10,6 @@ use lemmy_db_views::structs::{LocalUserView, SiteView};
 use lemmy_utils::{
   claims::Claims,
   error::{LemmyError, LemmyErrorExt, LemmyErrorType},
-  utils::validation::check_totp_2fa_valid,
 };
 
 #[tracing::instrument(skip(context))]
@@ -53,13 +53,10 @@ pub async fn login(
   check_registration_application(&local_user_view, &site_view.local_site, &mut context.pool())
     .await?;
 
-  // Check the totp
+  // Check the totp if enabled
-  check_totp_2fa_valid(
+  if local_user_view.local_user.totp_2fa_enabled {
-    &local_user_view.local_user.totp_2fa_secret,
+    check_totp_2fa_valid(&local_user_view, &data.totp_2fa_token, &site_view.site.name)?;
-    &data.totp_2fa_token,
+  }
-    &site_view.site.name,
-    &local_user_view.person.name,
-  )?;
 
   // Return the jwt
   Ok(Json(LoginResponse {

crates/api/src/local_user/mod.rs

@@ -3,6 +3,7 @@ pub mod ban_person;
 pub mod block;
 pub mod change_password;
 pub mod change_password_after_reset;
+pub mod generate_totp_secret;
 pub mod get_captcha;
 pub mod list_banned;
 pub mod login;
@@ -10,4 +11,5 @@ pub mod notifications;
 pub mod report_count;
 pub mod reset_password;
 pub mod save_settings;
+pub mod update_totp;
 pub mod verify_email;

crates/api/src/local_user/save_settings.rs

@@ -17,13 +17,7 @@ use lemmy_db_views::structs::SiteView;
 use lemmy_utils::{
   claims::Claims,
   error::{LemmyError, LemmyErrorExt, LemmyErrorType},
-  utils::validation::{
+  utils::validation::{is_valid_bio_field, is_valid_display_name, is_valid_matrix_id},
-    build_totp_2fa,
-    generate_totp_2fa_secret,
-    is_valid_bio_field,
-    is_valid_display_name,
-    is_valid_matrix_id,
-  },
 };
 
 #[tracing::instrument(skip(context))]
@@ -105,20 +99,6 @@ pub async fn save_user_settings(
     LocalUserLanguage::update(&mut context.pool(), discussion_languages, local_user_id).await?;
   }
 
-  // If generate_totp is Some(false), this will clear it out from the database.
-  let (totp_2fa_secret, totp_2fa_url) = if let Some(generate) = data.generate_totp_2fa {
-    if generate {
-      let secret = generate_totp_2fa_secret();
-      let url =
-        build_totp_2fa(&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 {
     email,
     show_avatars: data.show_avatars,
@@ -134,8 +114,6 @@ pub async fn save_user_settings(
     default_listing_type,
     theme,
     interface_language: data.interface_language.clone(),
-    totp_2fa_secret,
-    totp_2fa_url,
     open_links_in_new_tab: data.open_links_in_new_tab,
     infinite_scroll_enabled: data.infinite_scroll_enabled,
     ..Default::default()

crates/api/src/local_user/update_totp.rs

@@ -0,0 +1,54 @@
+use crate::check_totp_2fa_valid;
+use actix_web::web::{Data, Json};
+use lemmy_api_common::{
+  context::LemmyContext,
+  person::{UpdateTotp, UpdateTotpResponse},
+};
+use lemmy_db_schema::{
+  source::local_user::{LocalUser, LocalUserUpdateForm},
+  traits::Crud,
+};
+use lemmy_db_views::structs::{LocalUserView, SiteView};
+use lemmy_utils::error::LemmyError;
+
+/// Enable or disable two-factor-authentication. The current setting is determined from
+/// [LocalUser.totp_2fa_enabled].
+///
+/// To enable, you need to first call [generate_totp_secret] and then pass a valid token to this
+/// function.
+///
+/// Disabling is only possible if 2FA was previously enabled. Again it is necessary to pass a valid
+/// token.
+#[tracing::instrument(skip(context))]
+pub async fn update_totp(
+  data: Json<UpdateTotp>,
+  local_user_view: LocalUserView,
+  context: Data<LemmyContext>,
+) -> Result<Json<UpdateTotpResponse>, LemmyError> {
+  let site_view = SiteView::read_local(&mut context.pool()).await?;
+
+  check_totp_2fa_valid(
+    &local_user_view,
+    &Some(data.totp_token.clone()),
+    &site_view.site.name,
+  )?;
+
+  // toggle the 2fa setting
+  let local_user_form = LocalUserUpdateForm {
+    totp_2fa_enabled: Some(data.enabled),
+    // if totp is enabled, leave unchanged. otherwise clear secret
+    totp_2fa_secret: if data.enabled { None } else { Some(None) },
+    ..Default::default()
+  };
+
+  LocalUser::update(
+    &mut context.pool(),
+    local_user_view.local_user.id,
+    &local_user_form,
+  )
+  .await?;
+
+  Ok(Json(UpdateTotpResponse {
+    enabled: data.enabled,
+  }))
+}

crates/api_common/src/person.rs

@@ -119,10 +119,6 @@ pub struct SaveUserSettings {
   pub show_new_post_notifs: Option<bool>,
   /// A list of languages you are able to see discussion in.
   pub discussion_languages: Option<Vec<LanguageId>>,
-  /// Generates a TOTP / 2-factor authentication token.
-  ///
-  /// None leaves it as is, true will generate or regenerate it, false clears it out.
-  pub generate_totp_2fa: Option<bool>,
   pub auth: Sensitive<String>,
   /// Open links in a new tab
   pub open_links_in_new_tab: Option<bool>,
@@ -443,3 +439,25 @@ pub struct VerifyEmail {
 #[cfg_attr(feature = "full", ts(export))]
 /// A response to verifying your email.
 pub struct VerifyEmailResponse {}
+
+#[derive(Debug, Serialize, Deserialize, Clone)]
+#[cfg_attr(feature = "full", derive(TS))]
+#[cfg_attr(feature = "full", ts(export))]
+pub struct GenerateTotpSecretResponse {
+  pub totp_secret_url: Sensitive<String>,
+}
+
+#[derive(Debug, Serialize, Deserialize, Clone)]
+#[cfg_attr(feature = "full", derive(TS))]
+#[cfg_attr(feature = "full", ts(export))]
+pub struct UpdateTotp {
+  pub totp_token: String,
+  pub enabled: bool,
+}
+
+#[derive(Debug, Serialize, Deserialize, Clone)]
+#[cfg_attr(feature = "full", derive(TS))]
+#[cfg_attr(feature = "full", ts(export))]
+pub struct UpdateTotpResponse {
+  pub enabled: bool,
+}

crates/db_schema/src/schema.rs

@@ -436,13 +436,13 @@ diesel::table! {
         email_verified -> Bool,
         accepted_application -> Bool,
         totp_2fa_secret -> Nullable<Text>,
-        totp_2fa_url -> Nullable<Text>,
         open_links_in_new_tab -> Bool,
         blur_nsfw -> Bool,
         auto_expand -> Bool,
         infinite_scroll_enabled -> Bool,
         admin -> Bool,
         post_listing_mode -> PostListingModeEnum,
+        totp_2fa_enabled -> Bool,
     }
 }
 

crates/db_schema/src/source/local_user.rs

@@ -51,8 +51,6 @@ pub struct LocalUser {
   pub accepted_application: bool,
   #[serde(skip)]
   pub totp_2fa_secret: Option<String>,
-  /// A URL to add their 2-factor auth.
-  pub totp_2fa_url: Option<String>,
   /// Open links in a new tab.
   pub open_links_in_new_tab: bool,
   pub blur_nsfw: bool,
@@ -62,6 +60,7 @@ pub struct LocalUser {
   /// Whether the person is an admin.
   pub admin: bool,
   pub post_listing_mode: PostListingMode,
+  pub totp_2fa_enabled: bool,
 }
 
 #[derive(Clone, TypedBuilder)]
@@ -88,13 +87,13 @@ pub struct LocalUserInsertForm {
   pub email_verified: Option<bool>,
   pub accepted_application: Option<bool>,
   pub totp_2fa_secret: Option<Option<String>>,
-  pub totp_2fa_url: Option<Option<String>>,
   pub open_links_in_new_tab: Option<bool>,
   pub blur_nsfw: Option<bool>,
   pub auto_expand: Option<bool>,
   pub infinite_scroll_enabled: Option<bool>,
   pub admin: Option<bool>,
   pub post_listing_mode: Option<PostListingMode>,
+  pub totp_2fa_enabled: Option<bool>,
 }
 
 #[derive(Clone, Default)]
@@ -117,11 +116,11 @@ pub struct LocalUserUpdateForm {
   pub email_verified: Option<bool>,
   pub accepted_application: Option<bool>,
   pub totp_2fa_secret: Option<Option<String>>,
-  pub totp_2fa_url: Option<Option<String>>,
   pub open_links_in_new_tab: Option<bool>,
   pub blur_nsfw: Option<bool>,
   pub auto_expand: Option<bool>,
   pub infinite_scroll_enabled: Option<bool>,
   pub admin: Option<bool>,
   pub post_listing_mode: Option<PostListingMode>,
+  pub totp_2fa_enabled: Option<bool>,
 }

crates/db_views/src/registration_application_view.rs

@@ -262,12 +262,12 @@ mod tests {
         email_verified: inserted_sara_local_user.email_verified,
         accepted_application: inserted_sara_local_user.accepted_application,
         totp_2fa_secret: inserted_sara_local_user.totp_2fa_secret,
-        totp_2fa_url: inserted_sara_local_user.totp_2fa_url,
         password_encrypted: inserted_sara_local_user.password_encrypted,
         open_links_in_new_tab: inserted_sara_local_user.open_links_in_new_tab,
         infinite_scroll_enabled: inserted_sara_local_user.infinite_scroll_enabled,
         admin: false,
         post_listing_mode: inserted_sara_local_user.post_listing_mode,
+        totp_2fa_enabled: inserted_sara_local_user.totp_2fa_enabled,
       },
       creator: Person {
         id: inserted_sara_person.id,

crates/utils/Cargo.toml

@@ -47,7 +47,6 @@ smart-default = "0.7.1"
 jsonwebtoken = "8.3.0"
 lettre = { version = "0.10.4", features = ["tokio1", "tokio1-native-tls"] }
 markdown-it = "0.5.1"
-totp-rs = { version = "5.0.2", features = ["gen_secret", "otpauth"] }
 ts-rs = { workspace = true, optional = true }
 enum-map = "2.6"
 

crates/utils/src/error.rs

@@ -159,8 +159,11 @@ pub enum LemmyErrorType {
   InvalidBodyField,
   BioLengthOverflow,
   MissingTotpToken,
+  MissingTotpSecret,
   IncorrectTotpToken,
   CouldntParseTotpSecret,
+  CouldntGenerateTotp,
+  TotpAlreadyEnabled,
   CouldntLikeComment,
   CouldntSaveComment,
   CouldntCreateReport,
@@ -192,7 +195,6 @@ pub enum LemmyErrorType {
   InvalidUrl,
   EmailSendFailed,
   Slurs,
-  CouldntGenerateTotp,
   CouldntFindObject,
   RegistrationDenied(String),
   FederationDisabled,

crates/utils/src/utils/validation.rs

@@ -1,8 +1,7 @@
-use crate::error::{LemmyError, LemmyErrorExt, LemmyErrorType, LemmyResult};
+use crate::error::{LemmyErrorExt, LemmyErrorType, LemmyResult};
 use itertools::Itertools;
 use once_cell::sync::Lazy;
 use regex::{Regex, RegexBuilder};
-use totp_rs::{Secret, TOTP};
 use url::Url;
 
 static VALID_ACTOR_NAME_REGEX: Lazy<Regex> =
@@ -238,54 +237,6 @@ pub fn clean_url_params(url: &Url) -> Url {
   url_out
 }
 
-pub fn check_totp_2fa_valid(
-  totp_secret: &Option<String>,
-  totp_token: &Option<String>,
-  site_name: &str,
-  username: &str,
-) -> LemmyResult<()> {
-  // Check only if they have a totp_secret in the DB
-  if let Some(totp_secret) = totp_secret {
-    // Throw an error if their token is missing
-    let token = totp_token
-      .as_deref()
-      .ok_or(LemmyErrorType::MissingTotpToken)?;
-
-    let totp = build_totp_2fa(site_name, username, totp_secret)?;
-
-    let check_passed = totp.check_current(token)?;
-    if !check_passed {
-      Err(LemmyErrorType::IncorrectTotpToken.into())
-    } else {
-      Ok(())
-    }
-  } else {
-    Ok(())
-  }
-}
-
-pub fn generate_totp_2fa_secret() -> String {
-  Secret::generate_secret().to_string()
-}
-
-pub fn build_totp_2fa(site_name: &str, username: &str, secret: &str) -> Result<TOTP, LemmyError> {
-  let sec = Secret::Raw(secret.as_bytes().to_vec());
-  let sec_bytes = sec
-    .to_bytes()
-    .map_err(|_| LemmyErrorType::CouldntParseTotpSecret)?;
-
-  TOTP::new(
-    totp_rs::Algorithm::SHA256,
-    6,
-    1,
-    30,
-    sec_bytes,
-    Some(site_name.to_string()),
-    username.to_string(),
-  )
-  .with_lemmy_type(LemmyErrorType::CouldntGenerateTotp)
-}
-
 pub fn check_site_visibility_valid(
   current_private_instance: bool,
   current_federation_enabled: bool,
@@ -319,7 +270,6 @@ mod tests {
   #![allow(clippy::unwrap_used)]
   #![allow(clippy::indexing_slicing)]
 
-  use super::build_totp_2fa;
   use crate::{
     error::LemmyErrorType,
     utils::validation::{
@@ -327,7 +277,6 @@ mod tests {
       check_site_visibility_valid,
       check_url_scheme,
       clean_url_params,
-      generate_totp_2fa_secret,
       is_valid_actor_name,
       is_valid_bio_field,
       is_valid_display_name,
@@ -400,13 +349,6 @@ mod tests {
     assert!(is_valid_matrix_id("@dess:matrix.org t").is_err());
   }
 
-  #[test]
-  fn test_build_totp() {
-    let generated_secret = generate_totp_2fa_secret();
-    let totp = build_totp_2fa("lemmy", "my_name", &generated_secret);
-    assert!(totp.is_ok());
-  }
-
   #[test]
   fn test_valid_site_name() {
     let valid_names = [

migrations/2023-09-11-110040_rework-2fa-setup/down.sql

@@ -0,0 +1,6 @@
+ALTER TABLE local_user
+    ADD COLUMN totp_2fa_url text;
+
+ALTER TABLE local_user
+    DROP COLUMN totp_2fa_enabled;
+

migrations/2023-09-11-110040_rework-2fa-setup/up.sql

@@ -0,0 +1,6 @@
+ALTER TABLE local_user
+    DROP COLUMN totp_2fa_url;
+
+ALTER TABLE local_user
+    ADD COLUMN totp_2fa_enabled boolean NOT NULL DEFAULT FALSE;
+

src/api_routes_http.rs

@@ -20,6 +20,7 @@ use lemmy_api::{
     block::block_person,
     change_password::change_password,
     change_password_after_reset::change_password_after_reset,
+    generate_totp_secret::generate_totp_secret,
     get_captcha::get_captcha,
     list_banned::list_banned_users,
     login::login,
@@ -34,6 +35,7 @@ use lemmy_api::{
     report_count::report_count,
     reset_password::reset_password,
     save_settings::save_user_settings,
+    update_totp::update_totp,
     verify_email::verify_email,
   },
   post::{
@@ -287,7 +289,9 @@ pub fn config(cfg: &mut web::ServiceConfig, rate_limit: &RateLimitCell) {
           .route("/report_count", web::get().to(report_count))
           .route("/unread_count", web::get().to(unread_count))
           .route("/verify_email", web::post().to(verify_email))
-          .route("/leave_admin", web::post().to(leave_admin)),
+          .route("/leave_admin", web::post().to(leave_admin))
+          .route("/totp/generate", web::post().to(generate_totp_secret))
+          .route("/totp/update", web::post().to(update_totp)),
       )
       // Admin Actions
       .service(