From 6fa434e89c4231d3ce209fefbdb8aeae0246b512 Mon Sep 17 00:00:00 2001
From: ryan <ryan@distrust.co>
Date: Sun, 14 Apr 2024 20:27:00 -0400
Subject: [PATCH] keyfork-shard: shorten length and pad inside encrypted block

---
 crates/keyfork-shard/src/lib.rs | 68 +++++++++++++++++----------------
 1 file changed, 35 insertions(+), 33 deletions(-)

diff --git a/crates/keyfork-shard/src/lib.rs b/crates/keyfork-shard/src/lib.rs
index 338aaea..e81c57a 100644
--- a/crates/keyfork-shard/src/lib.rs
+++ b/crates/keyfork-shard/src/lib.rs
@@ -24,8 +24,10 @@ use sha2::Sha256;
 use sharks::{Share, Sharks};
 use x25519_dalek::{EphemeralSecret, PublicKey};
 
-// 256 bit share encrypted is 49 bytes, couple more bytes before we reach max size
-const ENC_LEN: u8 = 4 * 16;
+// 32-byte share, 1-byte index, 1-byte threshold, 1-byte version == 36 bytes
+// Encrypted, is 52 bytes
+const PLAINTEXT_LENGTH: u8 = 36;
+const ENCRYPTED_LENGTH: u8 = PLAINTEXT_LENGTH + 16;
 
 #[cfg(feature = "openpgp")]
 pub mod openpgp;
@@ -263,41 +265,34 @@ pub trait Format {
         payload.insert(0, HUNK_VERSION);
         payload.insert(1, threshold);
         assert!(
-            payload.len() <= ENC_LEN as usize,
-            "invalid share length (too long, max {ENC_LEN} bytes)"
+            payload.len() < PLAINTEXT_LENGTH as usize,
+            "invalid share length (too long, max {PLAINTEXT_LENGTH} bytes)"
         );
 
-        // encrypt data
-        let payload_bytes = shared_key.encrypt(nonce, payload.as_slice())?;
-
-        // convert data to a static-size payload
-        // NOTE: Padding length is less than u8::MAX because ENC_LEN < u8::MAX
+        // convert plaintext to static-size payload
         #[allow(clippy::assertions_on_constants)]
         {
-            assert!(ENC_LEN < u8::MAX, "padding byte can be u8");
+            assert!(PLAINTEXT_LENGTH < u8::MAX, "length byte can be u8");
         }
-        #[allow(clippy::cast_possible_truncation)]
-        let mut out_bytes = [payload_bytes.len() as u8; ENC_LEN as usize];
-        assert!(
-            payload_bytes.len() < out_bytes.len(),
-            "encrypted payload larger than acceptable limit"
-        );
-        out_bytes[..payload_bytes.len()].clone_from_slice(&payload_bytes);
 
-        // NOTE: This previously used a single repeated value as the padding byte, but resulted in
-        // difficulty when entering in prompts manually, as one's place could be lost due to
-        // repeated keywords. This is resolved below by having sequentially increasing numbers up to
-        // but not including the last byte.
+        // NOTE: Previous versions of Keyfork Shard would modify the padding bytes to avoid
+        // duplicate mnemonic words. This version does not include that, and instead uses a
+        // repeated length byte.
         #[allow(clippy::cast_possible_truncation)]
-        for (i, byte) in (out_bytes[payload_bytes.len()..(ENC_LEN as usize - 1)])
-            .iter_mut()
-            .enumerate()
-        {
-            *byte = (i % u8::MAX as usize) as u8;
-        }
+        let mut plaintext_bytes = [u8::try_from(payload.len()).expect(bug!(
+            "previously asserted length must be < {PLAINTEXT_LENGTH}",
+            PLAINTEXT_LENGTH = PLAINTEXT_LENGTH
+        )); PLAINTEXT_LENGTH as usize];
+        plaintext_bytes[..payload.len()].clone_from_slice(&payload);
+
+        // encrypt data
+        let encrypted_bytes = shared_key.encrypt(nonce, plaintext_bytes.as_slice())?;
+
+        assert_eq!(encrypted_bytes.len(), ENCRYPTED_LENGTH as usize);
 
         // safety: size of out_bytes is constant and always % 4 == 0
-        let payload_mnemonic = unsafe { Mnemonic::from_raw_bytes(&out_bytes) };
+        let payload_mnemonic = unsafe { Mnemonic::from_raw_bytes(&encrypted_bytes) };
+        dbg!(payload_mnemonic.words().len());
 
         #[cfg(feature = "qrcode")]
         {
@@ -399,7 +394,7 @@ pub struct InvalidData;
 /// 1 byte: Version
 /// 1 byte: Threshold
 /// Data: &[u8]
-pub(crate) const HUNK_VERSION: u8 = 1;
+pub(crate) const HUNK_VERSION: u8 = 2;
 pub(crate) const HUNK_OFFSET: usize = 2;
 
 const QRCODE_PROMPT: &str = "Press enter, then present QR code to camera.";
@@ -468,7 +463,7 @@ pub fn remote_decrypt(w: &mut impl Write) -> Result<(), Box<dyn std::error::Erro
             if let Ok(Some(hex)) =
                 keyfork_qrcode::scan_camera(std::time::Duration::from_secs(QRCODE_TIMEOUT), 0)
             {
-                let decoded_data = smex::decode(&hex)?;
+                let decoded_data = smex::decode(hex)?;
                 let _ = pubkey_data.insert(decoded_data[..32].try_into().map_err(|_| InvalidData)?);
                 let _ = payload_data.insert(decoded_data[32..].to_vec());
             } else {
@@ -480,7 +475,7 @@ pub fn remote_decrypt(w: &mut impl Write) -> Result<(), Box<dyn std::error::Erro
             (Some(pubkey), Some(payload)) => (pubkey, payload),
             _ => {
                 let validator = MnemonicSetValidator {
-                    word_lengths: [24, 48],
+                    word_lengths: [24, 39],
                 };
 
                 let [pubkey_mnemonic, payload_mnemonic] = pm
@@ -498,6 +493,12 @@ pub fn remote_decrypt(w: &mut impl Write) -> Result<(), Box<dyn std::error::Erro
             }
         };
 
+        assert_eq!(
+            payload.len(),
+            ENCRYPTED_LENGTH as usize,
+            bug!("invalid payload data")
+        );
+
         let shared_secret = our_key.diffie_hellman(&PublicKey::from(pubkey)).to_bytes();
         let hkdf = Hkdf::<Sha256>::new(None, &shared_secret);
 
@@ -509,7 +510,7 @@ pub fn remote_decrypt(w: &mut impl Write) -> Result<(), Box<dyn std::error::Erro
         hkdf.expand(b"nonce", &mut nonce_data)?;
         let nonce = Nonce::<U12>::from_slice(&nonce_data);
 
-        let payload = shared_key.decrypt(nonce, &payload[..payload[payload.len() - 1] as usize])?;
+        let payload = shared_key.decrypt(nonce, payload.as_slice())?;
         assert_eq!(HUNK_VERSION, payload[0], "Incompatible hunk version");
 
         match &mut iter_count {
@@ -524,7 +525,8 @@ pub fn remote_decrypt(w: &mut impl Write) -> Result<(), Box<dyn std::error::Erro
             }
         }
 
-        shares.push(payload[HUNK_OFFSET..].to_vec());
+        let payload_len = payload.last().expect(bug!("payload should not be empty"));
+        shares.push(payload[HUNK_OFFSET..usize::from(*payload_len)].to_vec());
     }
 
     let shares = shares