Adding TOTP / 2FA to lemmy (#2741)

* Combine prod and dev docker setups using build-arg

- Fixes #2603

* Dont use cache for release build.

* Adding 2FA / TOTP support.

- Fixes #2363

* Changed name to totp_2fa for clarity.

* Switch to sha256 for totp.

Cargo.lock

@@ -599,6 +599,12 @@ dependencies = [
  "rustc-demangle",
 ]
 
+[[package]]
+name = "base32"
+version = "0.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "23ce669cd6c8588f79e15cf450314f9638f967fc5770ff1c7c1deb0925ea7cfa"
+
 [[package]]
 name = "base64"
 version = "0.13.1"
@@ -932,6 +938,12 @@ dependencies = [
  "tracing-subscriber",
 ]
 
+[[package]]
+name = "constant_time_eq"
+version = "0.2.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f3ad85c1f65dc7b37604eb0e89748faf0b9653065f2a8ef69f96a687ec1e9279"
+
 [[package]]
 name = "convert_case"
 version = "0.4.0"
@@ -2639,6 +2651,7 @@ dependencies = [
  "strum",
  "strum_macros",
  "tokio",
+ "totp-rs",
  "tracing",
  "tracing-error",
  "typed-builder",
@@ -5033,6 +5046,22 @@ dependencies = [
  "syn 1.0.103",
 ]
 
+[[package]]
+name = "totp-rs"
+version = "4.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9fdd21080b6cf581e0c8fe849626ad627b42af1a0f71ce980244f2d6b1a47836"
+dependencies = [
+ "base32",
+ "constant_time_eq",
+ "hmac",
+ "rand 0.8.5",
+ "sha1",
+ "sha2",
+ "url",
+ "urlencoding",
+]
+
 [[package]]
 name = "tower"
 version = "0.4.13"
@@ -5387,6 +5416,12 @@ dependencies = [
  "serde",
 ]
 
+[[package]]
+name = "urlencoding"
+version = "2.1.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e8db7427f936968176eaa7cdf81b7f98b980b18495ec28f1b5791ac3bfe3eea9"
+
 [[package]]
 name = "utf-8"
 version = "0.7.6"

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_2fa_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_2fa_valid(
+      &local_user_view.local_user.totp_2fa_secret,
+      &data.totp_2fa_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_2fa,
+    generate_totp_2fa_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_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::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_2fa_secret(totp_2fa_secret)
+      .totp_2fa_url(totp_2fa_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_2fa_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_2fa: Option<bool>,
   pub auth: Sensitive<String>,
 }
 

crates/db_schema/src/schema.rs

@@ -170,6 +170,8 @@ table! {
         show_new_post_notifs -> Bool,
         email_verified -> Bool,
         accepted_application -> Bool,
+        totp_2fa_secret -> Nullable<Text>,
+        totp_2fa_url -> Nullable<Text>,
     }
 }
 

crates/db_schema/src/source/local_user.rs

@@ -27,6 +27,9 @@ pub struct LocalUser {
   pub show_new_post_notifs: bool,
   pub email_verified: bool,
   pub accepted_application: bool,
+  #[serde(skip)]
+  pub totp_2fa_secret: Option<String>,
+  pub totp_2fa_url: Option<String>,
 }
 
 #[derive(Clone, TypedBuilder)]
@@ -52,6 +55,8 @@ pub struct LocalUserInsertForm {
   pub show_new_post_notifs: Option<bool>,
   pub email_verified: Option<bool>,
   pub accepted_application: Option<bool>,
+  pub totp_2fa_secret: Option<Option<String>>,
+  pub totp_2fa_url: Option<Option<String>>,
 }
 
 #[derive(Clone, TypedBuilder)]
@@ -74,4 +79,6 @@ pub struct LocalUserUpdateForm {
   pub show_new_post_notifs: Option<bool>,
   pub email_verified: Option<bool>,
   pub accepted_application: Option<bool>,
+  pub totp_2fa_secret: Option<Option<String>>,
+  pub totp_2fa_url: Option<Option<String>>,
 }

crates/db_views/src/registration_application_view.rs

@@ -284,6 +284,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_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,
       },
       creator: Person {

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"] }
 
 [dev-dependencies]
 reqwest = { workspace = true }

crates/utils/src/utils/validation.rs

@@ -1,6 +1,8 @@
+use crate::error::LemmyError;
 use itertools::Itertools;
 use once_cell::sync::Lazy;
 use regex::Regex;
+use totp_rs::{Secret, TOTP};
 use url::Url;
 
 static VALID_ACTOR_NAME_REGEX: Lazy<Regex> =
@@ -56,10 +58,58 @@ 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,
+) -> Result<(), LemmyError> {
+  // 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_else(|| LemmyError::from_message("missing_totp_token"))?;
+
+    let totp = build_totp_2fa(site_name, username, totp_secret)?;
+
+    let check_passed = totp.check_current(token)?;
+    if !check_passed {
+      return Err(LemmyError::from_message("incorrect_totp token"));
+    }
+  }
+
+  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(|_| LemmyError::from_message("Couldnt parse totp secret"))?;
+
+  TOTP::new(
+    totp_rs::Algorithm::SHA256,
+    6,
+    1,
+    30,
+    sec_bytes,
+    Some(site_name.to_string()),
+    username.to_string(),
+  )
+  .map_err(|e| LemmyError::from_error_message(e, "Couldnt generate TOTP"))
+}
+
 #[cfg(test)]
 mod tests {
+  use super::build_totp_2fa;
   use crate::utils::validation::{
     clean_url_params,
+    generate_totp_2fa_secret,
     is_valid_actor_name,
     is_valid_display_name,
     is_valid_matrix_id,
@@ -128,4 +178,11 @@ mod tests {
     assert!(!is_valid_matrix_id(" @dess:matrix.org"));
     assert!(!is_valid_matrix_id("@dess:matrix.org t"));
   }
+
+  #[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());
+  }
 }

migrations/2023-02-16-194139_add_totp_secret/down.sql

@@ -0,0 +1,2 @@
+alter table local_user drop column totp_2fa_secret;
+alter table local_user drop column totp_2fa_url;

migrations/2023-02-16-194139_add_totp_secret/up.sql

@@ -0,0 +1,2 @@
+alter table local_user add column totp_2fa_secret text;
+alter table local_user add column totp_2fa_url text;