Enable payjoin expiration

DanGould · DanGould · commit a8a24f9e988c · 2024-04-04T20:24:55.000-04:00
diff --git a/mutiny-core/src/lib.rs b/mutiny-core/src/lib.rs
@@ -1550,17 +1550,19 @@ impl<S: MutinyStorage> MutinyWallet<S> {
             let enrolled = enroller
                 .process_res(ohttp_response.as_ref(), context)
                 .map_err(|_| MutinyError::PayjoinCreateRequest)?;
-            self.node_manager
+            let session = self
+                .node_manager
                 .storage
                 .persist_payjoin(enrolled.clone())?;
             let pj_uri = enrolled.fallback_target();
             log_debug!(self.logger, "{pj_uri}");
             let wallet = self.node_manager.wallet.clone();
             let stop = self.node_manager.stop.clone();
+            let storage = Arc::new(self.node_manager.storage.clone());
             // run await payjoin task in the background as it'll keep polling the relay
             let logger = self.logger.clone();
             utils::spawn(async move {
-                match NodeManager::receive_payjoin(wallet, stop, enrolled).await {
+                match NodeManager::receive_payjoin(wallet, stop, storage, session).await {
                     Ok(pj_txid) => log_info!(logger, "Received payjoin txid: {}", pj_txid),
                     Err(e) => log_error!(logger, "Payjoin error: {e}"),
                 }
diff --git a/mutiny-core/src/nodemanager.rs b/mutiny-core/src/nodemanager.rs
@@ -635,8 +635,11 @@ impl<S: MutinyStorage> NodeManager<S> {
         for payjoin in all {
             let wallet = nm.wallet.clone();
             let stop = nm.stop.clone();
+            let storage = Arc::new(nm.storage.clone());
             utils::spawn(async move {
-                let pj_txid = Self::receive_payjoin(wallet, stop, payjoin).await.unwrap();
+                let pj_txid = Self::receive_payjoin(wallet, stop, storage, payjoin)
+                    .await
+                    .unwrap();
                 log::info!("Received payjoin txid: {}", pj_txid);
             });
         }
@@ -810,15 +813,16 @@ impl<S: MutinyStorage> NodeManager<S> {
     pub async fn receive_payjoin(
         wallet: Arc<OnChainWallet<S>>,
         stop: Arc<AtomicBool>,
-        mut enrolled: payjoin::receive::v2::Enrolled,
+        storage: Arc<S>,
+        mut session: crate::payjoin::Session,
     ) -> Result<Txid, MutinyError> {
         use crate::payjoin::Error as PayjoinError;
 
         let http_client = reqwest::Client::builder()
             .build()
             .map_err(PayjoinError::Reqwest)?;
         let proposal: payjoin::receive::v2::UncheckedProposal =
-            Self::poll_for_fallback_psbt(stop, &http_client, &mut enrolled)
+            Self::poll_for_fallback_psbt(stop, storage, &http_client, &mut session)
                 .await
                 .map_err(|e| PayjoinError::ReceiverStateMachine(e.to_string()))?;
         let mut payjoin_proposal = wallet
@@ -845,22 +849,30 @@ impl<S: MutinyStorage> NodeManager<S> {
 
     async fn poll_for_fallback_psbt(
         stop: Arc<AtomicBool>,
+        storage: Arc<S>,
         client: &reqwest::Client,
-        enroller: &mut payjoin::receive::v2::Enrolled,
+        session: &mut crate::payjoin::Session,
     ) -> Result<payjoin::receive::v2::UncheckedProposal, crate::payjoin::Error> {
         loop {
             if stop.load(Ordering::Relaxed) {
                 return Err(crate::payjoin::Error::Shutdown);
             }
-            let (req, context) = enroller.extract_req()?;
+
+            if session.expiry < utils::now() {
+                let _ = storage.delete_payjoin(&session.enrolled.pubkey());
+                return Err(crate::payjoin::Error::SessionExpired);
+            }
+            let (req, context) = session.enrolled.extract_req()?;
             let ohttp_response = client
                 .post(req.url)
                 .header("Content-Type", "message/ohttp-req")
                 .body(req.body)
                 .send()
                 .await?;
             let ohttp_response = ohttp_response.bytes().await?;
-            let proposal = enroller.process_res(ohttp_response.as_ref(), context)?;
+            let proposal = session
+                .enrolled
+                .process_res(ohttp_response.as_ref(), context)?;
             match proposal {
                 Some(proposal) => return Ok(proposal),
                 None => utils::sleep(5000).await,
diff --git a/mutiny-core/src/payjoin.rs b/mutiny-core/src/payjoin.rs
@@ -2,10 +2,12 @@ use std::collections::HashMap;
 
 use crate::error::MutinyError;
 use crate::storage::MutinyStorage;
+use core::time::Duration;
 use hex_conservative::DisplayHex;
 use once_cell::sync::Lazy;
 use payjoin::receive::v2::Enrolled;
 use payjoin::OhttpKeys;
+use serde::{Deserialize, Serialize};
 use url::Url;
 
 pub(crate) static OHTTP_RELAYS: [Lazy<Url>; 2] = [
@@ -16,10 +18,22 @@ pub(crate) static OHTTP_RELAYS: [Lazy<Url>; 2] = [
 pub(crate) static PAYJOIN_DIR: Lazy<Url> =
     Lazy::new(|| Url::parse("https://payjo.in").expect("Invalid URL"));
 
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+pub struct Session {
+    pub enrolled: Enrolled,
+    pub expiry: Duration,
+}
+
+impl Session {
+    pub fn pubkey(&self) -> [u8; 33] {
+        self.enrolled.pubkey()
+    }
+}
 pub trait PayjoinStorage {
-    fn get_payjoin(&self, id: &[u8; 33]) -> Result<Option<Enrolled>, MutinyError>;
-    fn get_payjoins(&self) -> Result<Vec<Enrolled>, MutinyError>;
-    fn persist_payjoin(&self, session: Enrolled) -> Result<(), MutinyError>;
+    fn get_payjoin(&self, id: &[u8; 33]) -> Result<Option<Session>, MutinyError>;
+    fn get_payjoins(&self) -> Result<Vec<Session>, MutinyError>;
+    fn persist_payjoin(&self, session: Enrolled) -> Result<Session, MutinyError>;
+    fn delete_payjoin(&self, id: &[u8; 33]) -> Result<(), MutinyError>;
 }
 
 const PAYJOIN_KEY_PREFIX: &str = "payjoin/";
@@ -29,18 +43,28 @@ fn get_payjoin_key(id: &[u8; 33]) -> String {
 }
 
 impl<S: MutinyStorage> PayjoinStorage for S {
-    fn get_payjoin(&self, id: &[u8; 33]) -> Result<Option<Enrolled>, MutinyError> {
+    fn get_payjoin(&self, id: &[u8; 33]) -> Result<Option<Session>, MutinyError> {
         let sessions = self.get_data(get_payjoin_key(id))?;
         Ok(sessions)
     }
 
-    fn get_payjoins(&self) -> Result<Vec<Enrolled>, MutinyError> {
-        let map: HashMap<String, Enrolled> = self.scan(PAYJOIN_KEY_PREFIX, None)?;
+    fn get_payjoins(&self) -> Result<Vec<Session>, MutinyError> {
+        let map: HashMap<String, Session> = self.scan(PAYJOIN_KEY_PREFIX, None)?;
         Ok(map.values().map(|v| v.to_owned()).collect())
     }
 
-    fn persist_payjoin(&self, session: Enrolled) -> Result<(), MutinyError> {
-        self.set_data(get_payjoin_key(&session.pubkey()), session, None)
+    fn persist_payjoin(&self, enrolled: Enrolled) -> Result<Session, MutinyError> {
+        let in_24_hours = crate::utils::now() + Duration::from_secs(60 * 60 * 24);
+        let session = Session {
+            enrolled,
+            expiry: in_24_hours,
+        };
+        self.set_data(get_payjoin_key(&session.pubkey()), session.clone(), None)
+            .map(|_| session)
+    }
+
+    fn delete_payjoin(&self, id: &[u8; 33]) -> Result<(), MutinyError> {
+        self.delete(&[get_payjoin_key(id)])
     }
 }
 
@@ -65,6 +89,7 @@ pub enum Error {
     ReceiverStateMachine(String),
     Txid(bitcoin::hashes::hex::Error),
     Shutdown,
+    SessionExpired,
 }
 
 impl std::error::Error for Error {}
@@ -76,6 +101,7 @@ impl std::fmt::Display for Error {
             Error::ReceiverStateMachine(e) => write!(f, "Payjoin state machine error: {}", e),
             Error::Txid(e) => write!(f, "Payjoin txid error: {}", e),
             Error::Shutdown => write!(f, "Payjoin stopped by application shutdown"),
+            Error::SessionExpired => write!(f, "Payjoin session expired. Create a new payment request and have the sender try again."),
         }
     }
 }
