From 187512520aa53335fd4ecd3eefb3745a1a11288f Mon Sep 17 00:00:00 2001
From: Yuval Kogman <nothingmuch@woobling.org>
Date: Mon, 11 Nov 2024 00:07:25 +0100
Subject: [PATCH] store public key in fragment of pj URI

---
 payjoin/src/receive/v2/mod.rs |  1 +
 payjoin/src/uri/mod.rs        |  7 +++++++
 payjoin/src/uri/url_ext.rs    | 13 +++++++++++++
 payjoin/tests/integration.rs  | 10 ++++++----
 4 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/payjoin/src/receive/v2/mod.rs b/payjoin/src/receive/v2/mod.rs
index 2f897ba5..86de9380 100644
--- a/payjoin/src/receive/v2/mod.rs
+++ b/payjoin/src/receive/v2/mod.rs
@@ -182,6 +182,7 @@ impl Receiver {
         PjUriBuilder::new(
             self.context.address.clone(),
             self.pj_url(),
+            Some(self.context.s.public_key().clone()),
             Some(self.context.ohttp_keys.clone()),
             Some(self.context.expiry),
         )
diff --git a/payjoin/src/uri/mod.rs b/payjoin/src/uri/mod.rs
index 3914e58b..6281446c 100644
--- a/payjoin/src/uri/mod.rs
+++ b/payjoin/src/uri/mod.rs
@@ -5,6 +5,8 @@ use bitcoin::{Address, Amount};
 pub use error::PjParseError;
 use url::Url;
 
+#[cfg(feature = "v2")]
+use crate::hpke::HpkePublicKey;
 use crate::uri::error::InternalPjParseError;
 #[cfg(feature = "v2")]
 pub(crate) use crate::uri::url_ext::UrlExt;
@@ -114,12 +116,15 @@ impl PjUriBuilder {
     pub fn new(
         address: Address,
         origin: Url,
+        #[cfg(feature = "v2")] receiver_pubkey: Option<HpkePublicKey>,
         #[cfg(feature = "v2")] ohttp_keys: Option<OhttpKeys>,
         #[cfg(feature = "v2")] expiry: Option<std::time::SystemTime>,
     ) -> Self {
         #[allow(unused_mut)]
         let mut pj = origin;
         #[cfg(feature = "v2")]
+        pj.set_receiver_pubkey(receiver_pubkey);
+        #[cfg(feature = "v2")]
         pj.set_ohttp(ohttp_keys);
         #[cfg(feature = "v2")]
         pj.set_exp(expiry);
@@ -361,6 +366,8 @@ mod tests {
                     None,
                     #[cfg(feature = "v2")]
                     None,
+                    #[cfg(feature = "v2")]
+                    None,
                 )
                 .amount(amount)
                 .message("message".to_string())
diff --git a/payjoin/src/uri/url_ext.rs b/payjoin/src/uri/url_ext.rs
index 0153e792..73d191ef 100644
--- a/payjoin/src/uri/url_ext.rs
+++ b/payjoin/src/uri/url_ext.rs
@@ -1,11 +1,15 @@
 use std::str::FromStr;
 
+use bitcoin::base64::prelude::BASE64_URL_SAFE_NO_PAD;
+use bitcoin::base64::Engine;
 use url::Url;
 
+use crate::hpke::HpkePublicKey;
 use crate::OhttpKeys;
 
 /// Parse and set fragment parameters from `&pj=` URI parameter URLs
 pub(crate) trait UrlExt {
+    fn set_receiver_pubkey(&mut self, exp: Option<HpkePublicKey>);
     fn ohttp(&self) -> Option<OhttpKeys>;
     fn set_ohttp(&mut self, ohttp: Option<OhttpKeys>);
     fn exp(&self) -> Option<std::time::SystemTime>;
@@ -13,6 +17,15 @@ pub(crate) trait UrlExt {
 }
 
 impl UrlExt for Url {
+    /// Set the receiver's public key in the URL fragment
+    fn set_receiver_pubkey(&mut self, pubkey: Option<HpkePublicKey>) {
+        set_param(
+            self,
+            "rk=",
+            pubkey.map(|k| BASE64_URL_SAFE_NO_PAD.encode(k.to_compressed_bytes())),
+        )
+    }
+
     /// Retrieve the ohttp parameter from the URL fragment
     fn ohttp(&self) -> Option<OhttpKeys> {
         get_param(self, "ohttp=", |value| OhttpKeys::from_str(value).ok())
diff --git a/payjoin/tests/integration.rs b/payjoin/tests/integration.rs
index af2ec41f..ef0f3cba 100644
--- a/payjoin/tests/integration.rs
+++ b/payjoin/tests/integration.rs
@@ -180,7 +180,7 @@ mod integration {
         use bitcoin::Address;
         use http::StatusCode;
         use payjoin::receive::v2::{PayjoinProposal, Receiver, UncheckedProposal};
-        use payjoin::{OhttpKeys, PjUri, UriExt};
+        use payjoin::{HpkeKeyPair, OhttpKeys, PjUri, UriExt};
         use reqwest::{Client, ClientBuilder, Error, Response};
         use testcontainers_modules::redis::Redis;
         use testcontainers_modules::testcontainers::clients::Cli;
@@ -280,6 +280,7 @@ mod integration {
                 let expired_pj_uri = payjoin::PjUriBuilder::new(
                     address,
                     directory.clone(),
+                    Some(HpkeKeyPair::gen_keypair().public_key().clone()),
                     Some(ohttp_keys),
                     Some(std::time::SystemTime::now()),
                 )
@@ -601,9 +602,10 @@ mod integration {
             let (_bitcoind, sender, receiver) = init_bitcoind_sender_receiver(None, None)?;
             // Receiver creates the payjoin URI
             let pj_receiver_address = receiver.get_new_address(None, None)?.assume_checked();
-            let pj_uri = PjUriBuilder::new(pj_receiver_address, EXAMPLE_URL.to_owned(), None, None)
-                .amount(Amount::ONE_BTC)
-                .build();
+            let pj_uri =
+                PjUriBuilder::new(pj_receiver_address, EXAMPLE_URL.to_owned(), None, None, None)
+                    .amount(Amount::ONE_BTC)
+                    .build();
 
             // **********************
             // Inside the Sender: