Fix: panic if polls are empty

2-5-perceivers · May 29, 2024 · cbf08ee · cbf08ee
1 parent 298de88
commit cbf08ee
Showing 1 changed file with 6 additions and 4 deletions.
diff --git a/polybius-lib/src/password_generation.rs b/polybius-lib/src/password_generation.rs
@@ -37,13 +37,15 @@ impl PasswordGeneration for PasswordData {
                     bits.push(PasswordBit::symbol_bit());
                 }
                 1 => {
-                    let number = self.numbers_poll.choose(&mut rng).unwrap();
-                    bits.push(PasswordBit::number_bit(number));
+                    if let Some(number) = self.numbers_poll.choose(&mut rng) {
+                        bits.push(PasswordBit::number_bit(number));
+                    }
                 }
                 // Increase the probability of choosing a string
                 _ => {
-                    let text = self.text_poll.choose(&mut rng).unwrap();
-                    bits.push(PasswordBit::string_bit(text));
+                    if let Some(text) = self.text_poll.choose(&mut rng) {
+                        bits.push(PasswordBit::string_bit(text));
+                    }
                 }
             }
         }