From 591ae101448a19555fce5bca4f5caed0ff82457c Mon Sep 17 00:00:00 2001
From: vpl <vpl@vpl.me>
Date: Mon, 26 Aug 2019 20:26:54 +0200
Subject: [PATCH] Get token from single u64

---
 src/api/core/two_factor/email.rs | 37 +++++++++++++++++++++-----------
 src/config.rs                    | 12 +++++++++--
 2 files changed, 35 insertions(+), 14 deletions(-)

diff --git a/src/api/core/two_factor/email.rs b/src/api/core/two_factor/email.rs
index c96ed94..1986f81 100644
--- a/src/api/core/two_factor/email.rs
+++ b/src/api/core/two_factor/email.rs
@@ -59,7 +59,7 @@ fn send_email_login(data: JsonUpcase<SendEmailLoginData>, conn: DbConn) -> Empty
     let type_ = TwoFactorType::Email as i32;
     let mut twofactor = TwoFactor::find_by_user_and_type(&user.uuid, type_, &conn)?;
 
-    let generated_token = generate_token(CONFIG.email_token_size());
+    let generated_token = generate_token(CONFIG.email_token_size())?;
     let mut twofactor_data = EmailTokenData::from_json(&twofactor.data)?;
     twofactor_data.set_token(generated_token);
     twofactor.data = twofactor_data.to_json();
@@ -101,14 +101,20 @@ struct SendEmailData {
     MasterPasswordHash: String,
 }
 
-fn generate_token(token_size: u64) -> String {
-    crypto::get_random(vec![0; token_size as usize])
-        .iter()
-        .map(|byte| { (byte % 10)})
-        .map(|num| {
-            char::from_digit(num as u32, 10).unwrap()
-        })
-        .collect()
+
+fn generate_token(token_size: u32) -> Result<String, Error> {
+    if token_size > 19 {
+        err!("Generating token failed")
+    }
+
+    // 8 bytes to create an u64 for up to 19 token digits
+    let bytes = crypto::get_random(vec![0; 8]);
+    let mut bytes_array = [0u8; 8];
+    bytes_array.copy_from_slice(&bytes);
+
+    let number = u64::from_be_bytes(bytes_array) % 10u64.pow(token_size);
+    let token = format!("{:0size$}", number, size = token_size as usize);
+    Ok(token)
 }
 
 /// Send a verification email to the specified email address to check whether it exists/belongs to user.
@@ -131,7 +137,7 @@ fn send_email(data: JsonUpcase<SendEmailData>, headers: Headers, conn: DbConn) -
         tf.delete(&conn)?;
     }
 
-    let generated_token = generate_token(CONFIG.email_token_size());
+    let generated_token = generate_token(CONFIG.email_token_size())?;
     let twofactor_data = EmailTokenData::new(data.Email, generated_token);
 
     // Uses EmailVerificationChallenge as type to show that it's not verified yet.
@@ -321,8 +327,15 @@ mod tests {
 
     #[test]
     fn test_token() {
-        let result = generate_token(100);
+        let result = generate_token(19).unwrap();
 
-        assert_eq!(result.chars().count(), 100);
+        assert_eq!(result.chars().count(), 19);
+    }
+
+    #[test]
+    fn test_token_too_large() {
+        let result = generate_token(20);
+
+        assert!(result.is_err(), "too large token should give an error");
     }
 }
diff --git a/src/config.rs b/src/config.rs
index 31754e0..49f41be 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -322,8 +322,8 @@ make_config! {
     email_2fa: _enable_email_2fa {
         /// Enabled |> Disabling will prevent users from setting up new email 2FA and using existing email 2FA configured
         _enable_email_2fa:      bool,   true,   def,      true;
-        /// Token number length |> Length of the numbers in an email token
-        email_token_size:       u64,    true,   def,      6;
+        /// Token number length |> Length of the numbers in an email token. Minimum of 6. Maximum is 19.
+        email_token_size:       u32,    true,   def,      6;
         /// Token expiration time |> Maximum time in seconds a token is valid. The time the user has to open email client and copy token.
         email_expiration_time:  u64,    true,   def,      600;
         /// Maximum attempts |> Maximum attempts before an email token is reset and a new email will need to be sent
@@ -378,6 +378,14 @@ fn validate_config(cfg: &ConfigItems) -> Result<(), Error> {
         err!("Both `SMTP_USERNAME` and `SMTP_PASSWORD` need to be set to enable email authentication")
     }
 
+    if cfg.email_token_size < 6 {
+        err!("`EMAIL_TOKEN_SIZE` has a minimum size of 6")
+    }
+
+    if cfg.email_token_size > 19 {
+        err!("`EMAIL_TOKEN_SIZE` has a maximum size of 19")
+    }
+
     Ok(())
 }