src/controllers/login_controller.rs

@@ -11,7 +11,7 @@ use crate::controllers::base_controller::{FatalErrorPage, redirect_user, redirec
 use crate::data::login_redirect::LoginRedirect;
 use crate::data::remote_ip::RemoteIP;
 use crate::data::session_identity::{SessionIdentity, SessionStatus};
-use crate::data::user::{FactorID, TwoFactor, User};
+use crate::data::user::{FactorID, TwoFactor, TwoFactorType, User};
 
 struct BaseLoginPage<'a> {
     danger: Option<String>,
@@ -235,7 +235,6 @@ pub struct ChooseSecondFactorQuery {
     redirect: LoginRedirect,
 }
 
-
 /// Let the user select the factor to use to authenticate
 pub async fn choose_2fa_method(id: Identity, query: web::Query<ChooseSecondFactorQuery>,
                                users: web::Data<Addr<UsersActor>>) -> impl Responder {
@@ -269,10 +268,18 @@ pub struct LoginWithOTPQuery {
     id: FactorID,
 }
 
+#[derive(serde::Deserialize)]
+pub struct LoginWithOTPForm {
+    code: String,
+}
+
 
 /// Login with OTP
 pub async fn login_with_otp(id: Identity, query: web::Query<LoginWithOTPQuery>,
+                            form: Option<web::Form<LoginWithOTPForm>>,
                             users: web::Data<Addr<UsersActor>>) -> impl Responder {
+    let mut danger = None;
+
     if !SessionIdentity(&id).need_2fa_auth() {
         return redirect_user_for_login(query.redirect.get());
     }
@@ -286,9 +293,26 @@ pub async fn login_with_otp(id: Identity, query: web::Query<LoginWithOTPQuery>,
             .body(FatalErrorPage { message: "Factor not found!" }.render().unwrap())
     };
 
+    let key = match &factor.kind {
+        TwoFactorType::TOTP(key) => key,
+        _ => {
+            return HttpResponse::Ok()
+                .body(FatalErrorPage { message: "Factor is not a TOTP key!" }.render().unwrap());
+        }
+    };
+
+    if let Some(form) = form {
+        if !key.check_code(&form.code).unwrap_or(false) {
+            danger = Some("Specified code is invalid!".to_string());
+        } else {
+            SessionIdentity(&id).set_status(SessionStatus::SignedIn);
+            return redirect_user(query.redirect.get());
+        }
+    }
+
     HttpResponse::Ok().body(LoginWithOTPTemplate {
         _p: BaseLoginPage {
-            danger: None,
+            danger,
             success: None,
             page_title: "Two-Factor Auth",
             app_name: APP_NAME,

src/data/user.rs

@@ -12,7 +12,8 @@ pub struct FactorID(pub String);
 
 #[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
 pub enum TwoFactorType {
-    TOTP(TotpKey)
+    TOTP(TotpKey),
+    _OTHER,
 }
 
 #[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
@@ -25,14 +26,16 @@ pub struct TwoFactor {
 impl TwoFactor {
     pub fn type_str(&self) -> &'static str {
         match self.kind {
-            TwoFactorType::TOTP(_) => "Authenticator app"
+            TwoFactorType::TOTP(_) => "Authenticator app",
+            _ => unimplemented!()
         }
     }
 
     pub fn login_url(&self, redirect_uri: &LoginRedirect) -> String {
         match self.kind {
             TwoFactorType::TOTP(_) => format!("/2fa_otp?id={}&redirect={}",
-                                              self.id.0, redirect_uri.get_encoded())
+                                              self.id.0, redirect_uri.get_encoded()),
+            _ => unimplemented!()
         }
     }
 }