Persistable payjoin types (payjoin#552)

Introduce persistence as a first class abstraction in payjoin crate payjoin#477

Persistence is implemented for both the sender and reciever.
The trait design is meant to be generic over the Value persisted.
 For example if a application is expecting to save the state at
init and finalize and non of the intermediate steps that will be an option.
And for the reciever I would advocate in a seperate PR that we persist

This takes a step toward completing payjoin#336

Author: DanGould

payjoin-cli/src/app/v2.rs

@@ -4,7 +4,7 @@ use anyhow::{anyhow, Context, Result};
 use payjoin::bitcoin::consensus::encode::serialize_hex;
 use payjoin::bitcoin::psbt::Psbt;
 use payjoin::bitcoin::{Amount, FeeRate};
-use payjoin::receive::v2::{Receiver, UncheckedProposal};
+use payjoin::receive::v2::{NewReceiver, Receiver, UncheckedProposal};
 use payjoin::receive::{Error, ImplementationError, ReplyableError};
 use payjoin::send::v2::{Sender, SenderBuilder};
 use payjoin::Uri;
@@ -14,6 +14,7 @@ use super::config::Config;
 use super::wallet::BitcoindWallet;
 use super::App as AppTrait;
 use crate::app::{handle_interrupt, http_agent};
+use crate::db::v2::{ReceiverPersister, SenderPersister};
 use crate::db::Database;
 
 #[derive(Clone)]
@@ -52,11 +53,15 @@ impl AppTrait for App {
             Some(send_session) => send_session,
             None => {
                 let psbt = self.create_original_psbt(&uri, fee_rate)?;
-                let mut req_ctx = SenderBuilder::new(psbt, uri.clone())
+                let mut persister = SenderPersister::new(self.db.clone());
+                let new_sender = SenderBuilder::new(psbt, uri.clone())
                     .build_recommended(fee_rate)
                     .with_context(|| "Failed to build payjoin request")?;
-                self.db.insert_send_session(&mut req_ctx, url)?;
-                req_ctx
+                let storage_token = new_sender
+                    .persist(&mut persister)
+                    .map_err(|e| anyhow!("Failed to persist sender: {}", e))?;
+                Sender::load(storage_token, &persister)
+                    .map_err(|e| anyhow!("Failed to load sender: {}", e))?
             }
         };
         self.spawn_payjoin_sender(req_ctx).await
@@ -65,13 +70,18 @@ impl AppTrait for App {
     async fn receive_payjoin(&self, amount: Amount) -> Result<()> {
         let address = self.wallet().get_new_address()?;
         let ohttp_keys = unwrap_ohttp_keys_or_else_fetch(&self.config).await?;
-        let session = Receiver::new(
+        let mut persister = ReceiverPersister::new(self.db.clone());
+        let new_receiver = NewReceiver::new(
             address,
             self.config.v2()?.pj_directory.clone(),
             ohttp_keys.clone(),
             None,
         )?;
-        self.db.insert_recv_session(session.clone())?;
+        let storage_token = new_receiver
+            .persist(&mut persister)
+            .map_err(|e| anyhow!("Failed to persist receiver: {}", e))?;
+        let session = Receiver::load(storage_token, &persister)
+            .map_err(|e| anyhow!("Failed to load receiver: {}", e))?;
         self.spawn_payjoin_receiver(session, Some(amount)).await
     }
 

payjoin-cli/src/db/error.rs

@@ -13,6 +13,8 @@ pub(crate) enum Error {
     Serialize(serde_json::Error),
     #[cfg(feature = "v2")]
     Deserialize(serde_json::Error),
+    #[cfg(feature = "v2")]
+    NotFound(String),
 }
 
 impl fmt::Display for Error {
@@ -23,6 +25,8 @@ impl fmt::Display for Error {
             Error::Serialize(e) => write!(f, "Serialization failed: {}", e),
             #[cfg(feature = "v2")]
             Error::Deserialize(e) => write!(f, "Deserialization failed: {}", e),
+            #[cfg(feature = "v2")]
+            Error::NotFound(key) => write!(f, "Key not found: {}", key),
         }
     }
 }

payjoin-cli/src/db/mod.rs

@@ -27,4 +27,4 @@ impl Database {
 }
 
 #[cfg(feature = "v2")]
-mod v2;
+pub(crate) mod v2;

payjoin-cli/src/db/v2.rs

@@ -1,21 +1,62 @@
+use std::sync::Arc;
+
 use bitcoincore_rpc::jsonrpc::serde_json;
-use payjoin::receive::v2::Receiver;
-use payjoin::send::v2::Sender;
-use sled::{IVec, Tree};
+use payjoin::persist::{Persister, Value};
+use payjoin::receive::v2::{Receiver, ReceiverToken};
+use payjoin::send::v2::{Sender, SenderToken};
+use sled::Tree;
 use url::Url;
 
 use super::*;
 
-impl Database {
-    pub(crate) fn insert_recv_session(&self, session: Receiver) -> Result<()> {
-        let recv_tree = self.0.open_tree("recv_sessions")?;
-        let key = &session.id();
-        let value = serde_json::to_string(&session).map_err(Error::Serialize)?;
-        recv_tree.insert(key.as_slice(), IVec::from(value.as_str()))?;
+pub(crate) struct SenderPersister(Arc<Database>);
+impl SenderPersister {
+    pub fn new(db: Arc<Database>) -> Self { Self(db) }
+}
+
+impl Persister<Sender> for SenderPersister {
+    type Token = SenderToken;
+    type Error = crate::db::error::Error;
+    fn save(&mut self, value: Sender) -> std::result::Result<SenderToken, Self::Error> {
+        let send_tree = self.0 .0.open_tree("send_sessions")?;
+        let key = value.key();
+        let value = serde_json::to_vec(&value).map_err(Error::Serialize)?;
+        send_tree.insert(key.clone(), value.as_slice())?;
+        send_tree.flush()?;
+        Ok(key)
+    }
+
+    fn load(&self, key: SenderToken) -> std::result::Result<Sender, Self::Error> {
+        let send_tree = self.0 .0.open_tree("send_sessions")?;
+        let value = send_tree.get(key.as_ref())?.ok_or(Error::NotFound(key.to_string()))?;
+        serde_json::from_slice(&value).map_err(Error::Deserialize)
+    }
+}
+
+pub(crate) struct ReceiverPersister(Arc<Database>);
+impl ReceiverPersister {
+    pub fn new(db: Arc<Database>) -> Self { Self(db) }
+}
+
+impl Persister<Receiver> for ReceiverPersister {
+    type Token = ReceiverToken;
+    type Error = crate::db::error::Error;
+    fn save(&mut self, value: Receiver) -> std::result::Result<ReceiverToken, Self::Error> {
+        let recv_tree = self.0 .0.open_tree("recv_sessions")?;
+        let key = value.key();
+        let value = serde_json::to_vec(&value).map_err(Error::Serialize)?;
+        recv_tree.insert(key.clone(), value.as_slice())?;
         recv_tree.flush()?;
-        Ok(())
+        Ok(key)
     }
+    fn load(&self, key: ReceiverToken) -> std::result::Result<Receiver, Self::Error> {
+        let recv_tree = self.0 .0.open_tree("recv_sessions")?;
+        let value = recv_tree.get(key.as_ref())?.ok_or(Error::NotFound(key.to_string()))?;
+        serde_json::from_slice(&value).map_err(Error::Deserialize)
+    }
+}
 
+impl Database {
     pub(crate) fn get_recv_sessions(&self) -> Result<Vec<Receiver>> {
         let recv_tree = self.0.open_tree("recv_sessions")?;
         let mut sessions = Vec::new();
@@ -34,14 +75,6 @@ impl Database {
         Ok(())
     }
 
-    pub(crate) fn insert_send_session(&self, session: &mut Sender, pj_url: &Url) -> Result<()> {
-        let send_tree: Tree = self.0.open_tree("send_sessions")?;
-        let value = serde_json::to_string(session).map_err(Error::Serialize)?;
-        send_tree.insert(pj_url.to_string(), IVec::from(value.as_str()))?;
-        send_tree.flush()?;
-        Ok(())
-    }
-
     pub(crate) fn get_send_sessions(&self) -> Result<Vec<Sender>> {
         let send_tree: Tree = self.0.open_tree("send_sessions")?;
         let mut sessions = Vec::new();
@@ -55,7 +88,7 @@ impl Database {
 
     pub(crate) fn get_send_session(&self, pj_url: &Url) -> Result<Option<Sender>> {
         let send_tree = self.0.open_tree("send_sessions")?;
-        if let Some(val) = send_tree.get(pj_url.to_string())? {
+        if let Some(val) = send_tree.get(pj_url.as_str())? {
             let session: Sender = serde_json::from_slice(&val).map_err(Error::Deserialize)?;
             Ok(Some(session))
         } else {
@@ -65,7 +98,7 @@ impl Database {
 
     pub(crate) fn clear_send_session(&self, pj_url: &Url) -> Result<()> {
         let send_tree: Tree = self.0.open_tree("send_sessions")?;
-        send_tree.remove(pj_url.to_string())?;
+        send_tree.remove(pj_url.as_str())?;
         send_tree.flush()?;
         Ok(())
     }

payjoin/src/lib.rs

@@ -28,6 +28,9 @@ pub mod receive;
 #[cfg(feature = "_core")]
 pub mod send;
 
+#[cfg(feature = "v2")]
+pub mod persist;
+
 #[cfg(feature = "v2")]
 pub(crate) mod hpke;
 #[cfg(feature = "v2")]

payjoin/src/persist/mod.rs

@@ -0,0 +1,61 @@
+use std::fmt::Display;
+
+/// Types that can generate their own keys for persistent storage
+pub trait Value: serde::Serialize + serde::de::DeserializeOwned + Sized + Clone {
+    type Key: AsRef<[u8]> + Clone + Display;
+
+    /// Unique identifier for this persisted value
+    fn key(&self) -> Self::Key;
+}
+
+/// Implemented types that should be persisted by the application.
+pub trait Persister<V: Value> {
+    type Token: From<V>;
+    type Error: std::error::Error + Send + Sync + 'static;
+
+    fn save(&mut self, value: V) -> Result<Self::Token, Self::Error>;
+    fn load(&self, token: Self::Token) -> Result<V, Self::Error>;
+}
+
+/// A key type that stores the value itself for no-op persistence
+#[derive(Debug, Clone, serde::Serialize)]
+pub struct NoopToken<V: Value>(V);
+
+impl<V: Value> AsRef<[u8]> for NoopToken<V> {
+    fn as_ref(&self) -> &[u8] {
+        // Since this is a no-op implementation, we can return an empty slice
+        // as we never actually need to use the bytes
+        &[]
+    }
+}
+
+impl<'de, V: Value> serde::Deserialize<'de> for NoopToken<V> {
+    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
+    where
+        D: serde::Deserializer<'de>,
+    {
+        Ok(NoopToken(V::deserialize(deserializer)?))
+    }
+}
+
+impl<V: Value> Value for NoopToken<V> {
+    type Key = V::Key;
+
+    fn key(&self) -> Self::Key { self.0.key() }
+}
+
+/// A persister that does nothing but store values in memory
+#[derive(Debug, Clone)]
+pub struct NoopPersister;
+
+impl<V: Value> From<V> for NoopToken<V> {
+    fn from(value: V) -> Self { NoopToken(value) }
+}
+impl<V: Value> Persister<V> for NoopPersister {
+    type Token = NoopToken<V>;
+    type Error = std::convert::Infallible;
+
+    fn save(&mut self, value: V) -> Result<Self::Token, Self::Error> { Ok(NoopToken(value)) }
+
+    fn load(&self, token: Self::Token) -> Result<V, Self::Error> { Ok(token.0) }
+}

payjoin/src/receive/v2/mod.rs

@@ -1,4 +1,5 @@
 //! Receive BIP 77 Payjoin v2
+use std::fmt::{self, Display};
 use std::str::FromStr;
 use std::time::{Duration, SystemTime};
 
@@ -19,6 +20,7 @@ use super::{
 use crate::hpke::{decrypt_message_a, encrypt_message_b, HpkeKeyPair, HpkePublicKey};
 use crate::ohttp::{ohttp_decapsulate, ohttp_encapsulate, OhttpEncapsulationError, OhttpKeys};
 use crate::output_substitution::OutputSubstitution;
+use crate::persist::{self, Persister};
 use crate::receive::{parse_payload, InputPair};
 use crate::uri::ShortId;
 use crate::{IntoUrl, IntoUrlError, Request};
@@ -71,12 +73,12 @@ fn subdir_path_from_pubkey(pubkey: &HpkePublicKey) -> ShortId {
 
 /// A payjoin V2 receiver, allowing for polled requests to the
 /// payjoin directory and response processing.
-#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
-pub struct Receiver {
+#[derive(Debug)]
+pub struct NewReceiver {
     context: SessionContext,
 }
 
-impl Receiver {
+impl NewReceiver {
     /// Creates a new `Receiver` with the provided parameters.
     ///
     /// # Parameters
@@ -96,7 +98,7 @@ impl Receiver {
         ohttp_keys: OhttpKeys,
         expire_after: Option<Duration>,
     ) -> Result<Self, IntoUrlError> {
-        Ok(Self {
+        let receiver = Self {
             context: SessionContext {
                 address,
                 directory: directory.into_url()?,
@@ -107,9 +109,53 @@ impl Receiver {
                 s: HpkeKeyPair::gen_keypair(),
                 e: None,
             },
-        })
+        };
+        Ok(receiver)
+    }
+
+    pub fn persist<P: Persister<Receiver>>(
+        &self,
+        persister: &mut P,
+    ) -> Result<P::Token, ImplementationError> {
+        let receiver = Receiver { context: self.context.clone() };
+        Ok(persister.save(receiver)?)
     }
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+pub struct Receiver {
+    context: SessionContext,
+}
 
+/// Opaque key type for the receiver
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct ReceiverToken(ShortId);
+
+impl Display for ReceiverToken {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { write!(f, "{}", self.0) }
+}
+
+impl From<Receiver> for ReceiverToken {
+    fn from(receiver: Receiver) -> Self { ReceiverToken(id(&receiver.context.s)) }
+}
+
+impl AsRef<[u8]> for ReceiverToken {
+    fn as_ref(&self) -> &[u8] { self.0.as_bytes() }
+}
+
+impl persist::Value for Receiver {
+    type Key = ReceiverToken;
+
+    fn key(&self) -> Self::Key { ReceiverToken(id(&self.context.s)) }
+}
+
+impl Receiver {
+    pub fn load<P: Persister<Receiver>>(
+        token: P::Token,
+        persister: &P,
+    ) -> Result<Self, ImplementationError> {
+        persister.load(token).map_err(ImplementationError::from)
+    }
     /// Extract an OHTTP Encapsulated HTTP GET request for the Original PSBT
     pub fn extract_req(
         &mut self,