From 628c98a2fe199b54998d11718594bd0ce834e0d6 Mon Sep 17 00:00:00 2001
From: ryan <ryan@distrust.co>
Date: Sun, 14 Apr 2024 23:15:35 -0400
Subject: [PATCH] keyfork-shard: add in bug messages

---
 crates/keyfork-shard/src/lib.rs | 25 ++++++++++++++++++++-----
 1 file changed, 20 insertions(+), 5 deletions(-)

diff --git a/crates/keyfork-shard/src/lib.rs b/crates/keyfork-shard/src/lib.rs
index cb5c669..a672951 100644
--- a/crates/keyfork-shard/src/lib.rs
+++ b/crates/keyfork-shard/src/lib.rs
@@ -10,6 +10,7 @@ use aes_gcm::{
     aead::{consts::U12, Aead},
     Aes256Gcm, KeyInit, Nonce,
 };
+use base64::prelude::{Engine, BASE64_STANDARD};
 use hkdf::Hkdf;
 use keyfork_bug::{bug, POISONED_MUTEX};
 use keyfork_mnemonic_util::{English, Mnemonic};
@@ -23,7 +24,6 @@ use keyfork_prompt::{
 use sha2::Sha256;
 use sharks::{Share, Sharks};
 use x25519_dalek::{EphemeralSecret, PublicKey};
-use base64::prelude::{BASE64_STANDARD, Engine};
 
 // 32-byte share, 1-byte index, 1-byte threshold, 1-byte version == 36 bytes
 // Encrypted, is 52 bytes
@@ -212,7 +212,9 @@ pub trait Format {
             if let Ok(Some(qrcode_content)) =
                 keyfork_qrcode::scan_camera(std::time::Duration::from_secs(30), 0)
             {
-                let decoded_data = BASE64_STANDARD.decode(qrcode_content).unwrap();
+                let decoded_data = BASE64_STANDARD
+                    .decode(qrcode_content)
+                    .expect(bug!("qrcode should contain base64 encoded data"));
                 pubkey_data = Some(decoded_data.try_into().map_err(|_| InvalidData)?)
             } else {
                 prompt
@@ -300,7 +302,10 @@ pub trait Format {
             use keyfork_qrcode::{qrencode, ErrorCorrection};
             let mut qrcode_data = our_pubkey_mnemonic.to_bytes();
             qrcode_data.extend(payload_mnemonic.as_bytes());
-            if let Ok(qrcode) = qrencode(&BASE64_STANDARD.encode(qrcode_data), ErrorCorrection::Highest) {
+            if let Ok(qrcode) = qrencode(
+                &BASE64_STANDARD.encode(qrcode_data),
+                ErrorCorrection::Highest,
+            ) {
                 prompt
                     .lock()
                     .expect(bug!(POISONED_MUTEX))
@@ -433,7 +438,10 @@ pub fn remote_decrypt(w: &mut impl Write) -> Result<(), Box<dyn std::error::Erro
         {
             use keyfork_qrcode::{qrencode, ErrorCorrection};
             let qrcode_data = key_mnemonic.to_bytes();
-            if let Ok(qrcode) = qrencode(&BASE64_STANDARD.encode(qrcode_data), ErrorCorrection::Highest) {
+            if let Ok(qrcode) = qrencode(
+                &BASE64_STANDARD.encode(qrcode_data),
+                ErrorCorrection::Highest,
+            ) {
                 pm.prompt_message(PromptMessage::Text(format!(
                     concat!(
                         "A QR code will be displayed after this prompt. ",
@@ -464,7 +472,14 @@ pub fn remote_decrypt(w: &mut impl Write) -> Result<(), Box<dyn std::error::Erro
             if let Ok(Some(qrcode_content)) =
                 keyfork_qrcode::scan_camera(std::time::Duration::from_secs(QRCODE_TIMEOUT), 0)
             {
-                let decoded_data = BASE64_STANDARD.decode(qrcode_content).unwrap();
+                let decoded_data = BASE64_STANDARD
+                    .decode(qrcode_content)
+                    .expect(bug!("qrcode should contain base64 encoded data"));
+                assert_eq!(
+                    decoded_data.len(),
+                    ENCRYPTED_LENGTH as usize,
+                    bug!("invalid payload data")
+                );
                 let _ = pubkey_data.insert(decoded_data[..32].try_into().map_err(|_| InvalidData)?);
                 let _ = payload_data.insert(decoded_data[32..].to_vec());
             } else {