From 2cf46e1a5f6eec563ae8dc57dd4b5dff56515fdb Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Daniel=20Garc=C3=ADa?=
 <dani-garcia@users.noreply.github.com>
Date: Sat, 26 May 2018 23:04:23 +0200
Subject: [PATCH] Make sure TOTP codes can be both Numbers or Strings, fixes
 #30

---
 src/api/core/two_factor.rs | 13 ++++++++-----
 src/api/mod.rs             |  9 ++++++++-
 2 files changed, 16 insertions(+), 6 deletions(-)

diff --git a/src/api/core/two_factor.rs b/src/api/core/two_factor.rs
index 7ebdaedf..b42ce340 100644
--- a/src/api/core/two_factor.rs
+++ b/src/api/core/two_factor.rs
@@ -6,7 +6,7 @@ use db::DbConn;
 
 use crypto;
 
-use api::{PasswordData, JsonResult};
+use api::{PasswordData, JsonResult, NumberOrString};
 use auth::Headers;
 
 #[get("/two-factor")]
@@ -98,12 +98,12 @@ fn generate_authenticator(data: Json<PasswordData>, headers: Headers) -> JsonRes
     })))
 }
 
-#[derive(Deserialize)]
+#[derive(Deserialize, Debug)]
 #[allow(non_snake_case)]
 struct EnableTwoFactorData {
     masterPasswordHash: String,
     key: String,
-    token: u64,
+    token: NumberOrString,
 }
 
 #[post("/two-factor/authenticator", data = "<data>")]
@@ -111,7 +111,10 @@ fn activate_authenticator(data: Json<EnableTwoFactorData>, headers: Headers, con
     let data: EnableTwoFactorData = data.into_inner();
     let password_hash = data.masterPasswordHash;
     let key = data.key;
-    let token = data.token;
+    let token = match data.token.to_i32() {
+        Some(n) => n as u64,
+        None => err!("Malformed token")
+    };
 
     if !headers.user.check_valid_password(&password_hash) {
         err!("Invalid password");
@@ -154,7 +157,7 @@ fn activate_authenticator(data: Json<EnableTwoFactorData>, headers: Headers, con
 struct DisableTwoFactorData {
     masterPasswordHash: String,
     #[serde(rename = "type")]
-    type_: u32,
+    type_: NumberOrString,
 }
 
 #[post("/two-factor/disable", data = "<data>")]
diff --git a/src/api/mod.rs b/src/api/mod.rs
index 7d9e63d5..ed0a82d3 100644
--- a/src/api/mod.rs
+++ b/src/api/mod.rs
@@ -22,7 +22,7 @@ struct PasswordData {
     masterPasswordHash: String
 }
 
-#[derive(Deserialize)]
+#[derive(Deserialize, Debug)]
 #[serde(untagged)]
 enum NumberOrString {
     Number(i32),
@@ -36,4 +36,11 @@ impl NumberOrString {
             NumberOrString::String(s) => s
         }
     }
+
+    fn to_i32(self) -> Option<i32> {
+        match self {
+            NumberOrString::Number(n) => Some(n),
+            NumberOrString::String(s) => s.parse().ok()
+        }  
+    }
 }