From 3df3caa235456fb3fbd9f75cafc5e13e29f8784c Mon Sep 17 00:00:00 2001
From: ryan <ryan@distrust.co>
Date: Mon, 8 Jan 2024 14:00:31 -0500
Subject: [PATCH] keyfork-shard: use hkdf for remote shard keys

---
 Cargo.lock                   | 11 +++++++++++
 keyfork-shard/Cargo.toml     |  2 ++
 keyfork-shard/src/lib.rs     |  7 ++++++-
 keyfork-shard/src/openpgp.rs | 12 +++++++++++-
 4 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 1aab32c..c7653b6 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -906,6 +906,15 @@ version = "0.1.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "5491a308e0214554f07a81d8944abe45f552871c12e3c3c6e7e5d354039a6c4c"
 
+[[package]]
+name = "hkdf"
+version = "0.12.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7b5f8eb2ad728638ea2c7d47a21db23b7b58a72ed6a38256b8a1849f15fbbdf7"
+dependencies = [
+ "hmac",
+]
+
 [[package]]
 name = "hmac"
 version = "0.12.1"
@@ -1152,12 +1161,14 @@ dependencies = [
  "anyhow",
  "card-backend",
  "card-backend-pcsc",
+ "hkdf",
  "keyfork-derive-openpgp",
  "keyfork-mnemonic-util",
  "keyfork-prompt",
  "openpgp-card",
  "openpgp-card-sequoia",
  "sequoia-openpgp",
+ "sha2",
  "sharks",
  "smex",
  "thiserror",
diff --git a/keyfork-shard/Cargo.toml b/keyfork-shard/Cargo.toml
index 36bab88..1636ebf 100644
--- a/keyfork-shard/Cargo.toml
+++ b/keyfork-shard/Cargo.toml
@@ -31,3 +31,5 @@ openpgp-card-sequoia = { version = "0.2.0", optional = true }
 openpgp-card = { version = "0.4.0", optional = true }
 sequoia-openpgp = { version = "1.16.1", optional = true }
 keyfork-prompt = { version = "0.1.0", path = "../keyfork-prompt", optional = true }
+hkdf = { version = "0.12.4", features = ["std"] }
+sha2 = "0.10.8"
diff --git a/keyfork-shard/src/lib.rs b/keyfork-shard/src/lib.rs
index 416224c..bb51102 100644
--- a/keyfork-shard/src/lib.rs
+++ b/keyfork-shard/src/lib.rs
@@ -7,6 +7,8 @@ use aes_gcm::{
     aead::{Aead, AeadCore, OsRng},
     Aes256Gcm, KeyInit,
 };
+use hkdf::Hkdf;
+use sha2::Sha256;
 use keyfork_mnemonic_util::{Mnemonic, Wordlist};
 use keyfork_prompt::{qrencode, Message as PromptMessage, PromptManager};
 use sharks::{Share, Sharks};
@@ -89,8 +91,11 @@ pub fn remote_decrypt(w: &mut impl Write) -> Result<(), Box<dyn std::error::Erro
         let shared_secret = our_key
             .diffie_hellman(&PublicKey::from(their_key))
             .to_bytes();
+        let hkdf = Hkdf::<Sha256>::new(None, &shared_secret);
+        let mut hkdf_output = [0u8; 256 / 8];
+        hkdf.expand(&[], &mut hkdf_output)?;
         let shared_key =
-            Aes256Gcm::new_from_slice(&shared_secret)?;
+            Aes256Gcm::new_from_slice(&hkdf_output)?;
 
         let payload = Mnemonic::from_str(&payload_mnemonic)?.entropy();
         let payload =
diff --git a/keyfork-shard/src/openpgp.rs b/keyfork-shard/src/openpgp.rs
index 1bba1bb..daaba77 100644
--- a/keyfork-shard/src/openpgp.rs
+++ b/keyfork-shard/src/openpgp.rs
@@ -10,6 +10,8 @@ use aes_gcm::{
     aes::cipher::InvalidLength,
     Aes256Gcm, Error as AesError, KeyInit, Nonce,
 };
+use hkdf::{Hkdf, InvalidLength as HkdfInvalidLength};
+use sha2::Sha256;
 use keyfork_derive_openpgp::derive_util::{
     request::{DerivationAlgorithm, DerivationRequest},
     DerivationPath,
@@ -65,6 +67,9 @@ pub enum Error {
     #[error("Invalid length of AES key: {0}")]
     AesLength(#[from] InvalidLength),
 
+    #[error("Invalid KDF length: {0}")]
+    HkdfLength(#[from] HkdfInvalidLength),
+
     #[error("Derived secret hash {0} != expected {1}")]
     InvalidSecret(Fingerprint, Fingerprint),
 
@@ -442,7 +447,12 @@ pub fn decrypt(
         "invalid share length (too long, max {ENC_LEN} bytes)"
     );
 
-    let shared_key = Aes256Gcm::new_from_slice(&shared_secret)?;
+    let hkdf = Hkdf::<Sha256>::new(None, &shared_secret);
+    let mut hkdf_output = [0u8; 256 / 8];
+    hkdf.expand(&[], &mut hkdf_output)?;
+    let shared_key =
+        Aes256Gcm::new_from_slice(&hkdf_output)?;
+
     let bytes = shared_key.encrypt(their_nonce, share.as_slice())?;
     shared_key.decrypt(their_nonce, &bytes[..])?;