Fix integer overflow during OTP code calculation (#443)

Digits value must be in this range -> 0..=9

Author: replydev

src/arguments/add.rs

@@ -1,4 +1,4 @@
-use clap::Args;
+use clap::{value_parser, Args};
 use color_eyre::eyre::{eyre, ErrReport};
 
 use zeroize::Zeroize;
@@ -39,7 +39,8 @@ pub struct AddArgs {
         short,
         long,
         default_value_t = 6,
-        default_value_if("type", "STEAM", "5")
+        default_value_if("type", "STEAM", "5"),
+        value_parser=value_parser!(u64).range(0..=9)
     )]
     pub digits: u64,
 

src/arguments/edit.rs

@@ -1,4 +1,4 @@
-use clap::Args;
+use clap::{value_parser, Args};
 use color_eyre::eyre::eyre;
 
 use crate::otp::{otp_algorithm::OTPAlgorithm, otp_element::OTPDatabase};
@@ -24,7 +24,7 @@ pub struct EditArgs {
     pub algorithm: Option<OTPAlgorithm>,
 
     /// Code digits
-    #[arg(short, long)]
+    #[arg(short, long, value_parser=value_parser!(u64).range(0..=9))]
     pub digits: Option<u64>,
 
     /// Code period

src/arguments/list.rs

@@ -58,11 +58,11 @@ impl<'a> TryFrom<&'a OTPElement> for JsonOtpList<'a> {
 
 impl SubcommandExecutor for ListArgs {
     fn run_command(self, otp_database: OTPDatabase) -> color_eyre::Result<OTPDatabase> {
-        let elements_iterator = otp_database.elements.iter().enumerate();
-
         if self.format.unwrap_or_default().json {
-            let json_elements = elements_iterator
-                .map(|(_, element)| element.try_into())
+            let json_elements = otp_database
+                .elements
+                .iter()
+                .map(|element| element.try_into())
                 .collect::<Result<Vec<JsonOtpList>>>()?;
 
             let stringified = serde_json::to_string_pretty(&json_elements)
@@ -73,19 +73,23 @@ impl SubcommandExecutor for ListArgs {
                 "{0: <6} {1: <30} {2: <70} {3: <10}",
                 "Index", "Issuer", "Label", "OTP"
             );
-            elements_iterator.for_each(|(index, e)| {
-                println!(
-                    "{0: <6} {1: <30} {2: <70} {3: <10}",
-                    index,
-                    if e.issuer.is_empty() {
-                        "<No issuer>"
-                    } else {
-                        e.issuer.as_str()
-                    },
-                    &e.label,
-                    e.get_otp_code().unwrap_or("ERROR".to_string())
-                )
-            });
+            otp_database
+                .elements
+                .iter()
+                .enumerate()
+                .for_each(|(index, e)| {
+                    println!(
+                        "{0: <6} {1: <30} {2: <70} {3: <10}",
+                        index,
+                        if e.issuer.is_empty() {
+                            "<No issuer>"
+                        } else {
+                            e.issuer.as_str()
+                        },
+                        &e.label,
+                        e.get_otp_code().unwrap_or("ERROR".to_string())
+                    )
+                });
         }
 
         Ok(otp_database)

src/interface/event.rs

@@ -2,7 +2,7 @@ use std::sync::mpsc;
 use std::thread;
 use std::time::{Duration, Instant};
 
-use crossterm::event::{self, Event as CrosstermEvent, KeyEvent, KeyEventKind, MouseEvent};
+use crossterm::event::{self, Event as CrosstermEvent, KeyEvent, KeyEventKind};
 
 use crate::interface::app::AppResult;
 
@@ -14,15 +14,15 @@ pub enum Event {
     /// Key press.
     Key(KeyEvent),
     /// Mouse click/scroll.
-    Mouse(MouseEvent),
+    Mouse(()),
     /// Terminal resize.
-    Resize(u16, u16),
+    Resize((), ()),
     /// Focus gained
     FocusGained(),
     /// Focus lost
     FocusLost(),
     /// Paste text
-    Paste(String),
+    Paste(()),
 }
 
 /// Terminal event handler.
@@ -59,11 +59,11 @@ impl EventHandler {
                                 Ok(())
                             }
                         }
-                        CrosstermEvent::Mouse(e) => sender.send(Event::Mouse(e)),
-                        CrosstermEvent::Resize(w, h) => sender.send(Event::Resize(w, h)),
+                        CrosstermEvent::Mouse(_e) => sender.send(Event::Mouse(())),
+                        CrosstermEvent::Resize(_w, _h) => sender.send(Event::Resize((), ())),
                         CrosstermEvent::FocusGained => sender.send(Event::FocusGained()),
                         CrosstermEvent::FocusLost => sender.send(Event::FocusLost()),
-                        CrosstermEvent::Paste(e) => sender.send(Event::Paste(e)),
+                        CrosstermEvent::Paste(_e) => sender.send(Event::Paste(())),
                     }
                     .expect("failed to send terminal event")
                 }

src/otp/algorithms/hotp_maker.rs

@@ -86,18 +86,16 @@ mod tests {
     #[test]
     fn test_hotp() {
         assert_eq!(
-            format_code(generate_hotp::<Sha1>("BASE32SECRET3232", 0).unwrap(), 6),
-            "260182"
-        );
-        assert_eq!(
-            format_code(generate_hotp::<Sha1>("BASE32SECRET3232", 1).unwrap(), 6),
-            "055283"
+            455260182,
+            generate_hotp::<Sha1>("BASE32SECRET3232", 0).unwrap()
         );
     }
 
-    fn format_code(value: u32, digits: u32) -> String {
-        // Get the formatted code
-        let s = (value % 10_u32.pow(digits)).to_string();
-        "0".repeat(digits as usize - s.chars().count()) + s.as_str()
+    #[test]
+    fn test_hotp_2() {
+        assert_eq!(
+            1617055283,
+            generate_hotp::<Sha1>("BASE32SECRET3232", 1).unwrap()
+        );
     }
 }

src/otp/algorithms/steam_otp_maker.rs

@@ -8,10 +8,7 @@ use super::totp_maker::totp;
 const STEAM_ALPHABET: &str = "23456789BCDFGHJKMNPQRTVWXY";
 
 pub fn steam(secret: &str, algorithm: OTPAlgorithm, digits: usize) -> Result<String, OtpError> {
-    match totp(secret, algorithm) {
-        Ok(v) => Ok(to_steam_string(v as usize, digits)),
-        Err(e) => Err(e),
-    }
+    totp(secret, algorithm).map(|v| to_steam_string(v as usize, digits))
 }
 
 fn to_steam_string(mut code: usize, digits: usize) -> String {

src/otp/algorithms/totp_maker.rs

@@ -30,17 +30,8 @@ mod tests {
     #[test]
     fn test_totp() {
         assert_eq!(
-            format_code(
-                generate_totp("BASE32SECRET3232", OTPAlgorithm::Sha1, 0, 30, 0).unwrap(),
-                6
-            ),
-            "260182"
+            455260182,
+            generate_totp("BASE32SECRET3232", OTPAlgorithm::Sha1, 0, 30, 0).unwrap()
         );
     }
-
-    fn format_code(value: u32, digits: u32) -> String {
-        // Get the formatted code
-        let s = (value % 10_u32.pow(digits)).to_string();
-        "0".repeat(digits as usize - s.chars().count()) + s.as_str()
-    }
 }

src/otp/otp_element.rs

@@ -171,13 +171,13 @@ impl OTPElement {
             OTPType::Totp => {
                 let code = totp(&self.secret, self.algorithm)?;
 
-                Ok(self.format_code(code))
+                Ok(self.format_code(code)?)
             }
             OTPType::Hotp => match self.counter {
                 Some(counter) => {
                     let code = hotp(&self.secret, self.algorithm, counter)?;
 
-                    Ok(self.format_code(code))
+                    Ok(self.format_code(code)?)
                 }
                 None => Err(OtpError::MissingCounter),
             },
@@ -204,10 +204,13 @@ impl OTPElement {
         }
     }
 
-    pub fn format_code(&self, value: u32) -> String {
+    fn format_code(&self, value: u32) -> Result<String, OtpError> {
         // Get the formatted code
-        let s = (value % 10_u32.pow(self.digits as u32)).to_string();
-        "0".repeat(self.digits as usize - s.chars().count()) + s.as_str()
+        let exponential = 10_u32
+            .checked_pow(self.digits as u32)
+            .ok_or(OtpError::InvalidDigits)?;
+        let s = (value % exponential).to_string();
+        Ok("0".repeat(self.digits as usize - s.chars().count()) + s.as_str())
     }
 
     pub fn valid_secret(&self) -> bool {
@@ -231,6 +234,7 @@ mod test {
     use crate::otp::otp_element::OTPType::Totp;
 
     use crate::otp::from_otp_uri::FromOtpUri;
+    use crate::otp::otp_error::OtpError;
 
     #[test]
     fn test_serialization_otp_uri_full_element() {
@@ -287,4 +291,28 @@ mod test {
         let otp_uri = "otpauth://totp/2Ponies%40Github%20No.1?secret=JBSWY3DPEHPK3PXP&algorithm=SHA1&digits=6&period=30&lock=false&issuer=test";
         assert_eq!(true, OTPElement::from_otp_uri(otp_uri).is_ok())
     }
+
+    #[test]
+    fn test_invalid_digits_should_not_overflow() {
+        // Arrange
+        let invalid_digits_value = 10;
+
+        let element = OTPElement {
+            secret: "xr5gh44x7bprcqgrdtulafeevt5rxqlbh5wvked22re43dh2d4mapv5g".to_uppercase(),
+            issuer: String::from("IssuerText"),
+            label: String::from("LabelText"),
+            digits: invalid_digits_value,
+            type_: Totp,
+            algorithm: Sha1,
+            period: 30,
+            counter: None,
+            pin: None,
+        };
+
+        // Act
+        let result = element.get_otp_code();
+
+        // Assert
+        assert_eq!(Err(OtpError::InvalidDigits), result)
+    }
 }

src/otp/otp_error.rs

@@ -9,6 +9,7 @@ pub enum OtpError {
     MissingCounter,                    // Missing counter for HOTP codes
     InvalidOffset,                     // Invalid offset
     InvalidDigest,                     // Invalid digest
+    InvalidDigits,                     // Invalid Digits value (too high or low)
 }
 
 impl Display for OtpError {
@@ -22,6 +23,7 @@ impl Display for OtpError {
             OtpError::InvalidDigest => f.write_str("Invalid digest"),
             OtpError::InvalidOffset => f.write_str("Invalid offset"),
             OtpError::ShortSecret => f.write_str("Secret length less than 16 bytes"),
+            OtpError::InvalidDigits => f.write_str("Digits value too high or low"),
         }
     }
 }