From 31604d286629e1d38e6a0620a8c6c5e25295c00c Mon Sep 17 00:00:00 2001
From: DanGould <d@ngould.dev>
Date: Fri, 6 Sep 2024 12:43:27 -0400
Subject: [PATCH] Pad plaintext client messages

This way an ohttp relay sees uniform ciphertexts
---
 payjoin/src/v2.rs | 27 ++++++++++++++-------------
 1 file changed, 14 insertions(+), 13 deletions(-)

diff --git a/payjoin/src/v2.rs b/payjoin/src/v2.rs
index e1482d23..773affd2 100644
--- a/payjoin/src/v2.rs
+++ b/payjoin/src/v2.rs
@@ -3,13 +3,16 @@ use std::{error, fmt};
 
 use bitcoin::base64::prelude::BASE64_URL_SAFE_NO_PAD;
 use bitcoin::base64::Engine;
+use bitcoin::key::constants::UNCOMPRESSED_PUBLIC_KEY_SIZE;
 use hpke::aead::ChaCha20Poly1305;
 use hpke::kdf::HkdfSha256;
 use hpke::kem::SecpK256HkdfSha256;
 use hpke::rand_core::OsRng;
 use hpke::{Deserializable, OpModeR, OpModeS, Serializable};
 
-pub const PADDED_MESSAGE_BYTES: usize = 7168; // 7KB
+pub const PADDED_MESSAGE_BYTES: usize = 7168;
+pub const PADDED_PLAINTEXT_LENGTH: usize = PADDED_MESSAGE_BYTES - UNCOMPRESSED_PUBLIC_KEY_SIZE * 2;
+
 pub const INFO_A: &[u8] = b"Payjoin v2 Message A";
 pub const INFO_B: &[u8] = b"Payjoin v2 Message B";
 
@@ -106,7 +109,7 @@ impl<'de> serde::Deserialize<'de> for HpkePublicKey {
 
 #[cfg(feature = "send")]
 pub fn encrypt_message_a_hpke(
-    plaintext: Vec<u8>,
+    mut plaintext: Vec<u8>,
     pj_s_sk_e: &HpkeSecretKey,
     pj_r_pk_s: &HpkePublicKey,
 ) -> Result<Vec<u8>, HpkeError> {
@@ -119,11 +122,11 @@ pub fn encrypt_message_a_hpke(
             &mut OsRng,
         )?;
     let aad = pk.to_bytes().to_vec();
-    let ciphertext: Vec<u8> = encryption_context.seal(&plaintext, &aad)?;
+    let plaintext = pad_plaintext(&mut plaintext)?;
+    let ciphertext = encryption_context.seal(plaintext, &aad)?;
     let mut message_a = encapsulated_key.to_bytes().to_vec();
     message_a.extend(&aad);
     message_a.extend(&ciphertext);
-    //TODO let message_a = pad(&mut message_a).expect("TODO: handle error");
     Ok(message_a.to_vec())
 }
 
@@ -148,7 +151,7 @@ pub fn decrypt_message_a_hpke(
 
 #[cfg(feature = "receive")]
 pub fn encrypt_message_b_hpke(
-    plaintext: Vec<u8>,
+    mut plaintext: Vec<u8>,
     pj_r_s: (HpkeSecretKey, HpkePublicKey),
     pj_s_pk_re: &HpkePublicKey,
 ) -> Result<Vec<u8>, HpkeError> {
@@ -161,11 +164,11 @@ pub fn encrypt_message_b_hpke(
             &mut OsRng,
         )?;
     let aad = pk.to_bytes().to_vec();
-    let ciphertext = encryption_context.seal(&plaintext, &aad)?;
+    let plaintext = pad_plaintext(&mut plaintext)?;
+    let ciphertext = encryption_context.seal(plaintext, &aad)?;
     let mut message_b = encapsulated_key.to_bytes().to_vec();
     message_b.extend(&aad);
     message_b.extend(&ciphertext);
-    //let message_b = pad(&mut message_b).expect("TODO: handle error");
     Ok(message_b.to_vec())
 }
 
@@ -185,13 +188,11 @@ pub fn decrypt_message_b_hpke(message_b: &[u8], s: HpkeSecretKey) -> Result<Vec<
     Ok(plaintext)
 }
 
-fn pad(msg: &mut Vec<u8>) -> Result<&[u8], HpkeError> {
-    if msg.len() > PADDED_MESSAGE_BYTES {
+fn pad_plaintext(msg: &mut Vec<u8>) -> Result<&[u8], HpkeError> {
+    if msg.len() > PADDED_PLAINTEXT_LENGTH {
         return Err(HpkeError::PayloadTooLarge);
     }
-    while msg.len() < PADDED_MESSAGE_BYTES {
-        msg.push(0);
-    }
+    msg.resize(PADDED_PLAINTEXT_LENGTH, 0);
     Ok(msg)
 }
 
@@ -216,7 +217,7 @@ impl fmt::Display for HpkeError {
             Hpke(e) => e.fmt(f),
             InvalidKeyLength => write!(f, "Invalid Length"),
             PayloadTooLarge =>
-                write!(f, "Payload too large, max size is {} bytes", PADDED_MESSAGE_BYTES),
+                write!(f, "Plaintext too large, max size is {} bytes", PADDED_PLAINTEXT_LENGTH),
             PayloadTooShort => write!(f, "Payload too small"),
         }
     }