Add HpkeKeyPair abstraction

payjoin · Sep 26, 2024 · 3b9769b · 3b9769b
1 parent 76000be
commit 3b9769b
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 16 deletions.
diff --git a/payjoin/src/receive/v2/mod.rs b/payjoin/src/receive/v2/mod.rs
@@ -17,7 +17,7 @@ use super::{
 };
 use crate::psbt::PsbtExt;
 use crate::receive::optional_parameters::Params;
-use crate::v2::{HpkePublicKey, HpkeSecretKey, OhttpEncapsulationError};
+use crate::v2::{HpkeKeyPair, HpkePublicKey, OhttpEncapsulationError};
 use crate::{OhttpKeys, PjUriBuilder, Request};
 
 pub(crate) mod error;
@@ -33,7 +33,7 @@ struct SessionContext {
     ohttp_keys: OhttpKeys,
     expiry: SystemTime,
     ohttp_relay: url::Url,
-    s: (HpkeSecretKey, HpkePublicKey),
+    s: HpkeKeyPair,
     e: Option<HpkePublicKey>,
 }
 
@@ -85,15 +85,15 @@ impl SessionInitializer {
                 ohttp_relay,
                 expiry: SystemTime::now()
                     + expire_after.unwrap_or(TWENTY_FOUR_HOURS_DEFAULT_EXPIRY),
-                s: crate::v2::gen_keypair(),
+                s: HpkeKeyPair::gen_keypair(),
                 e: None,
             },
         }
     }
 
     pub fn extract_req(&mut self) -> Result<(Request, ohttp::ClientResponse), Error> {
         let url = self.context.ohttp_relay.clone();
-        let subdirectory = subdir_path_from_pubkey(&self.context.s.1);
+        let subdirectory = subdir_path_from_pubkey(self.context.s.public_key());
         let (body, ctx) = crate::v2::ohttp_encapsulate(
             &mut self.context.ohttp_keys,
             "POST",
@@ -192,7 +192,8 @@ impl ActiveSession {
     }
 
     fn extract_proposal_from_v2(&mut self, response: Vec<u8>) -> Result<UncheckedProposal, Error> {
-        let (payload_bytes, e) = crate::v2::decrypt_message_a(&response, self.context.s.0.clone())?;
+        let (payload_bytes, e) =
+            crate::v2::decrypt_message_a(&response, self.context.s.secret_key().clone())?;
         self.context.e = Some(e);
         let payload = String::from_utf8(payload_bytes).map_err(InternalRequestError::Utf8)?;
         Ok(self.unchecked_from_payload(payload)?)
@@ -252,7 +253,7 @@ impl ActiveSession {
     }
 
     /// The per-session public key to use as an identifier
-    pub fn id(&self) -> [u8; 33] { self.context.s.1.to_compressed_bytes() }
+    pub fn id(&self) -> [u8; 33] { self.context.s.public_key().to_compressed_bytes() }
 }
 
 /// The sender's original PSBT and optional parameters
@@ -532,11 +533,11 @@ impl PayjoinProposal {
             Some(e) => {
                 let payjoin_bytes = self.inner.payjoin_psbt.serialize();
                 log::debug!("THERE IS AN e: {:?}", e);
-                crate::v2::encrypt_message_b(payjoin_bytes, self.context.s.clone(), e)
+                crate::v2::encrypt_message_b(payjoin_bytes, &self.context.s, e)
             }
             None => Ok(self.extract_v1_req().as_bytes().to_vec()),
         }?;
-        let subdir_path = subdir_path_from_pubkey(&self.context.s.1);
+        let subdir_path = subdir_path_from_pubkey(self.context.s.public_key());
         let post_payjoin_target =
             self.context.directory.join(&subdir_path).map_err(|e| Error::Server(e.into()))?;
         log::debug!("Payjoin post target: {}", post_payjoin_target.as_str());
@@ -602,7 +603,7 @@ mod test {
                 ),
                 ohttp_relay: url::Url::parse("https://relay.com").unwrap(),
                 expiry: SystemTime::now() + Duration::from_secs(60),
-                s: crate::v2::gen_keypair(),
+                s: HpkeKeyPair::gen_keypair(),
                 e: None,
             },
         };

diff --git a/payjoin/src/send/mod.rs b/payjoin/src/send/mod.rs
@@ -242,7 +242,7 @@ impl<'a> RequestBuilder<'a> {
             sequence,
             min_fee_rate: self.min_fee_rate,
             #[cfg(feature = "v2")]
-            e: crate::v2::gen_keypair().0,
+            e: crate::v2::HpkeKeyPair::gen_keypair().secret_key().clone(),
         })
     }
 }

diff --git a/payjoin/src/v2.rs b/payjoin/src/v2.rs
@@ -9,6 +9,7 @@ use hpke::kdf::HkdfSha256;
 use hpke::kem::SecpK256HkdfSha256;
 use hpke::rand_core::OsRng;
 use hpke::{Deserializable, OpModeR, OpModeS, Serializable};
+use serde::{Deserialize, Serialize};
 
 pub const PADDED_MESSAGE_BYTES: usize = 7168;
 pub const PADDED_PLAINTEXT_A_LENGTH: usize =
@@ -23,9 +24,20 @@ pub type EncappedKey = <SecpK256HkdfSha256 as hpke::Kem>::EncappedKey;
 
 fn sk_to_pk(sk: &SecretKey) -> PublicKey { <SecpK256HkdfSha256 as hpke::Kem>::sk_to_pk(sk) }
 
-pub(crate) fn gen_keypair() -> (HpkeSecretKey, HpkePublicKey) {
-    let (sk, pk) = <SecpK256HkdfSha256 as hpke::Kem>::gen_keypair(&mut OsRng);
-    (HpkeSecretKey(sk), HpkePublicKey(pk))
+#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
+pub struct HpkeKeyPair(pub HpkeSecretKey, pub HpkePublicKey);
+
+impl From<HpkeKeyPair> for (HpkeSecretKey, HpkePublicKey) {
+    fn from(value: HpkeKeyPair) -> Self { (value.0, value.1) }
+}
+
+impl HpkeKeyPair {
+    pub fn gen_keypair() -> Self {
+        let (sk, pk) = <SecpK256HkdfSha256 as hpke::Kem>::gen_keypair(&mut OsRng);
+        Self(HpkeSecretKey(sk), HpkePublicKey(pk))
+    }
+    pub fn secret_key(&self) -> &HpkeSecretKey { &self.0 }
+    pub fn public_key(&self) -> &HpkePublicKey { &self.1 }
 }
 
 #[derive(Clone, PartialEq, Eq)]
@@ -164,13 +176,15 @@ pub fn decrypt_message_a(
 #[cfg(feature = "receive")]
 pub fn encrypt_message_b(
     mut plaintext: Vec<u8>,
-    receiver_keypair: (HpkeSecretKey, HpkePublicKey),
+    receiver_keypair: &HpkeKeyPair,
     sender_pk: &HpkePublicKey,
 ) -> Result<Vec<u8>, HpkeError> {
-    let pk = sk_to_pk(&receiver_keypair.0 .0);
     let (encapsulated_key, mut encryption_context) =
         hpke::setup_sender::<ChaCha20Poly1305, HkdfSha256, SecpK256HkdfSha256, _>(
-            &OpModeS::Auth((receiver_keypair.0 .0, pk.clone())),
+            &OpModeS::Auth((
+                receiver_keypair.secret_key().0.clone(),
+                receiver_keypair.public_key().0.clone(),
+            )),
             &sender_pk.0,
             INFO_B,
             &mut OsRng,