Select the OHTTP relay inside fetch_ohttp_keys

bc1cindy · bc1cindy · commit 2f4453456f2e · 2026-06-30T15:25:06.000-03:00
Consumers inherit one selection and failover policy instead of each
diverging.
diff --git a/payjoin-cli/src/app/v2/mod.rs b/payjoin-cli/src/app/v2/mod.rs
@@ -12,6 +12,7 @@ use payjoin::receive::v2::{
     ReceiverBuilder, SessionOutcome as ReceiverSessionOutcome, UncheckedOriginalPayload,
     WantsFeeRange, WantsInputs, WantsOutputs,
 };
+use payjoin::relay::RelaySelector;
 use payjoin::send::v2::{
     replay_event_log as replay_sender_event_log, PendingFallback as SenderPendingFallback,
     PollingForProposal, SendSession, Sender, SenderBuilder, SessionOutcome as SenderSessionOutcome,
@@ -278,7 +279,6 @@ impl AppTrait for App {
                 Err(e) => {
                     tracing::debug!("Directory {directory} failed: {e:#}");
                     self.mailroom_manager.add_failed_directory(directory);
-                    self.mailroom_manager.clear_failed_relays();
                     continue;
                 }
             }
@@ -1055,14 +1055,17 @@ impl App {
         F: FnMut(&str) -> std::result::Result<(payjoin::Request, T), E>,
         E: Into<anyhow::Error>,
     {
+        let mut selector = RelaySelector::new(self.config.v2()?.ohttp_relays.clone());
         loop {
-            let relay = self.mailroom_manager.choose_relay()?;
+            let relay = selector
+                .select(&mut payjoin::bitcoin::key::rand::thread_rng())
+                .ok_or_else(|| anyhow!("No valid relays available"))?;
             let (req, ctx) = build(relay.as_str()).map_err(Into::into)?;
             match self.post_request(req).await {
                 Ok(resp) => return Ok((resp, ctx)),
                 Err(e) => {
                     tracing::debug!("Request to relay {relay} failed: {e:?}");
-                    self.mailroom_manager.add_failed_relay(relay);
+                    selector.mark_failed(&relay);
                 }
             }
         }
diff --git a/payjoin-cli/src/app/v2/ohttp.rs b/payjoin-cli/src/app/v2/ohttp.rs
@@ -1,16 +1,14 @@
-//! OHTTP relay and payjoin directory selection / key bootstrapping for the payjoin-cli.
+//! Payjoin directory selection / OHTTP key bootstrapping for the payjoin-cli.
 //!
-//! [`MailroomManager`] tracks relays and directories that have failed,
-//! excluding them from future selections for the lifetime of the [`MailroomManager`].
+//! [`MailroomManager`] tracks directories that have failed, excluding them from
+//! future selections for the lifetime of the [`MailroomManager`]. Relay
+//! selection and failover are delegated to [`payjoin::io::fetch_ohttp_keys`].
 //!
 //! `unwrap_ohttp_keys_or_else_fetch_from_directory` returns user-supplied keys
-//! when present, otherwise selects a relay at random from the configured list
-//! (excluding failed relays) to fetch OHTTP keys from the given directory.
-//!
-//! `fetch_ohttp_keys_from_directory` retries on relay failures (e.g. connection
-//! errors) by selecting another relay. Once a directory is chosen for a session
-//! it must not change — the directory is embedded in the BIP21 URI at session
-//! creation and recovered from the session event log on resume.
+//! when present, otherwise fetches OHTTP keys from the given directory via a
+//! relay chosen from the configured list. Once a directory is chosen for a
+//! session it must not change — the directory is embedded in the BIP21 URI at
+//! session creation and recovered from the session event log on resume.
 use std::sync::{Arc, Mutex};
 
 use anyhow::{anyhow, Result};
@@ -21,48 +19,18 @@ use super::Config;
 #[derive(Debug, Clone)]
 pub struct MailroomManager {
     config: Config,
-    failed_relays: Arc<Mutex<Vec<Url>>>,
     failed_directories: Arc<Mutex<Vec<Url>>>,
 }
 
 impl MailroomManager {
     pub fn new(config: Config) -> Self {
-        MailroomManager {
-            config,
-            failed_relays: Arc::new(Mutex::new(Vec::new())),
-            failed_directories: Arc::new(Mutex::new(Vec::new())),
-        }
-    }
-
-    pub fn add_failed_relay(&self, relay: Url) {
-        self.failed_relays.lock().expect("Lock should not be poisoned").push(relay);
-    }
-
-    pub fn clear_failed_relays(&self) {
-        self.failed_relays.lock().expect("Lock should not be poisoned").clear();
+        MailroomManager { config, failed_directories: Arc::new(Mutex::new(Vec::new())) }
     }
 
     pub fn add_failed_directory(&self, directory: Url) {
         self.failed_directories.lock().expect("Lock should not be poisoned").push(directory);
     }
 
-    pub fn choose_relay(&self) -> Result<Url> {
-        use payjoin::bitcoin::secp256k1::rand::prelude::SliceRandom;
-        let relays = &self.config.v2()?.ohttp_relays;
-        let failed_relays = self.failed_relays.lock().expect("Lock should not be poisoned");
-        let remaining_relays: Vec<_> =
-            relays.iter().filter(|r| !failed_relays.contains(r)).cloned().collect();
-
-        if remaining_relays.is_empty() {
-            return Err(anyhow!("No valid relays available"));
-        }
-
-        remaining_relays
-            .choose(&mut payjoin::bitcoin::key::rand::thread_rng())
-            .cloned()
-            .ok_or_else(|| anyhow!("Failed to select from remaining relays"))
-    }
-
     pub fn choose_directory(&self) -> Result<Url> {
         use payjoin::bitcoin::secp256k1::rand::prelude::SliceRandom;
         let directories = &self.config.v2()?.pj_directories;
@@ -92,43 +60,33 @@ impl MailroomManager {
     }
 
     async fn fetch_ohttp_keys_from_directory(&self, directory: &Url) -> Result<ValidatedOhttpKeys> {
-        loop {
-            let selected_relay = self.choose_relay()?;
-
-            let ohttp_keys = {
-                #[cfg(feature = "_manual-tls")]
-                {
-                    if let Some(cert_path) = self.config.root_certificate.as_ref() {
-                        let cert_der = std::fs::read(cert_path)?;
-                        payjoin::io::fetch_ohttp_keys_with_cert(
-                            selected_relay.as_str(),
-                            directory.as_str(),
-                            &cert_der,
-                        )
+        let relays = &self.config.v2()?.ohttp_relays;
+        let result = {
+            #[cfg(feature = "_manual-tls")]
+            {
+                if let Some(cert_path) = self.config.root_certificate.as_ref() {
+                    let cert_der = std::fs::read(cert_path)?;
+                    payjoin::io::fetch_ohttp_keys_with_cert(relays, directory.as_str(), &cert_der)
                         .await
-                    } else {
-                        payjoin::io::fetch_ohttp_keys(selected_relay.as_str(), directory.as_str())
-                            .await
-                    }
-                }
-                #[cfg(not(feature = "_manual-tls"))]
-                payjoin::io::fetch_ohttp_keys(selected_relay.as_str(), directory.as_str()).await
-            };
-
-            match ohttp_keys {
-                Ok(keys) => return Ok(ValidatedOhttpKeys { ohttp_keys: keys }),
-                Err(payjoin::io::Error::UnexpectedStatusCode(e)) => {
-                    tracing::debug!(
-                        "Directory {directory} returned unexpected status via relay {selected_relay}: {e:?}"
-                    );
-                    self.add_failed_directory(directory.clone());
-                    return Err(anyhow!("Directory {directory} returned unexpected status: {e}"));
-                }
-                Err(e) => {
-                    tracing::debug!("Failed to connect to relay: {selected_relay}, {e:?}");
-                    self.add_failed_relay(selected_relay);
+                } else {
+                    payjoin::io::fetch_ohttp_keys(relays, directory.as_str()).await
                 }
             }
+            #[cfg(not(feature = "_manual-tls"))]
+            payjoin::io::fetch_ohttp_keys(relays, directory.as_str()).await
+        };
+
+        match result {
+            Ok((ohttp_keys, _relay)) => Ok(ValidatedOhttpKeys { ohttp_keys }),
+            Err(payjoin::io::Error::UnexpectedStatusCode(e)) => {
+                tracing::debug!("Directory {directory} returned unexpected status: {e:?}");
+                self.add_failed_directory(directory.clone());
+                Err(anyhow!("Directory {directory} returned unexpected status: {e}"))
+            }
+            Err(e) => {
+                tracing::debug!("Failed to fetch ohttp keys from directory {directory}: {e:?}");
+                Err(anyhow!("Failed to fetch ohttp keys: {e}"))
+            }
         }
     }
 }
diff --git a/payjoin-test-utils/src/v2.rs b/payjoin-test-utils/src/v2.rs
@@ -10,7 +10,7 @@ use ohttp::hpke::{Aead, Kdf, Kem};
 use ohttp::{KeyId, SymmetricSuite};
 use payjoin::io::{fetch_ohttp_keys_with_cert, Error as IOError};
 pub use payjoin::persist::{InMemoryPersister, SessionPersister};
-use payjoin::OhttpKeys;
+use payjoin::{OhttpKeys, Url};
 use rcgen::Certificate;
 use reqwest::{Client, ClientBuilder};
 use rustls::pki_types::CertificateDer;
@@ -110,12 +110,16 @@ impl TestServices {
     }
 
     pub async fn fetch_ohttp_keys(&self) -> Result<OhttpKeys, IOError> {
-        fetch_ohttp_keys_with_cert(
-            self.ohttp_relay_url().as_str(),
-            self.directory_url().as_str(),
-            &self.cert(),
-        )
-        .await
+        let relays: Vec<Url> = self
+            .ohttp_relays
+            .iter()
+            .map(|r| format!("http://localhost:{}", r.0))
+            .map(|s| Url::parse(&s).expect("valid relay URL"))
+            .collect();
+        let (ohttp_keys, _relay) =
+            fetch_ohttp_keys_with_cert(&relays, self.directory_url().as_str(), &self.cert())
+                .await?;
+        Ok(ohttp_keys)
     }
 }
 
diff --git a/payjoin/CHANGELOG.md b/payjoin/CHANGELOG.md
@@ -44,6 +44,10 @@ Selected Improvements:
 - Document BIP77 v1 fallback behavior in `create_post_request` (#1593)
 - Remove redundant language from `finalize_proposal` rustdocs (#1567)
 
+### Relay Selection
+
+- Centralize OHTTP relay selection in `payjoin::relay::RelaySelector`; `io::fetch_ohttp_keys` now takes a relay slice and returns the relay that served the keys (breaking: `fetch_ohttp_keys(relay, dir) -> OhttpKeys` becomes `fetch_ohttp_keys(relays, dir) -> (OhttpKeys, Url)`)
+
 ## 0.25.0
 
 Introduce monitoring typestates, replyable error handling, async
diff --git a/payjoin/src/core/io.rs b/payjoin/src/core/io.rs
@@ -1,60 +1,93 @@
 //! IO-related types and functions. Specifically, fetching OHTTP keys from a payjoin directory.
 use std::time::Duration;
 
+use bitcoin::secp256k1::rand;
 use http::header::ACCEPT;
 use reqwest::{Client, Proxy};
 
 use crate::into_url::IntoUrl;
-use crate::OhttpKeys;
+use crate::relay::RelaySelector;
+use crate::{OhttpKeys, Url};
 
 /// Fetch the ohttp keys from the specified payjoin directory via proxy.
 ///
-/// * `ohttp_relay`: The http CONNECT method proxy to request the ohttp keys from a payjoin
-///   directory.  Proxying requests for ohttp keys ensures a client IP address is never revealed to
-///   the payjoin directory.
+/// * `relays`: The http CONNECT method proxies to request the ohttp keys from a payjoin
+///   directory, tried in random order until one succeeds.  Proxying requests for ohttp keys
+///   ensures a client IP address is never revealed to the payjoin directory.
 ///
 /// * `payjoin_directory`: The payjoin directory from which to fetch the ohttp keys.  This
 ///   directory stores and forwards payjoin client payloads.
+///
+/// Returns the ohttp keys and the relay that served them.
 pub async fn fetch_ohttp_keys(
-    ohttp_relay: impl IntoUrl,
+    relays: &[Url],
     payjoin_directory: impl IntoUrl,
-) -> Result<OhttpKeys, Error> {
-    let ohttp_keys_url = payjoin_directory.into_url()?.join("/.well-known/ohttp-gateway")?;
-    let proxy = Proxy::all(ohttp_relay.into_url()?.as_str())?;
-    let client = Client::builder().proxy(proxy).http1_only().build()?;
-    let res = client
-        .get(ohttp_keys_url.as_str())
-        .timeout(Duration::from_secs(10))
-        .header(ACCEPT, "application/ohttp-keys")
-        .send()
-        .await?;
-    parse_ohttp_keys_response(res).await
+) -> Result<(OhttpKeys, Url), Error> {
+    fetch_ohttp_keys_inner(relays, payjoin_directory.into_url()?, None).await
 }
 
 /// Fetch the ohttp keys from the specified payjoin directory via proxy.
 ///
-/// * `ohttp_relay`: The http CONNECT method proxy to request the ohttp keys from a payjoin
-///   directory.  Proxying requests for ohttp keys ensures a client IP address is never revealed to
-///   the payjoin directory.
+/// * `relays`: The http CONNECT method proxies to request the ohttp keys from a payjoin
+///   directory, tried in random order until one succeeds.  Proxying requests for ohttp keys
+///   ensures a client IP address is never revealed to the payjoin directory.
 ///
 /// * `payjoin_directory`: The payjoin directory from which to fetch the ohttp keys.  This
 ///   directory stores and forwards payjoin client payloads.
 ///
 /// * `cert_der`: The DER-encoded certificate to use for local HTTPS connections.
+///
+/// Returns the ohttp keys and the relay that served them.
 #[cfg(feature = "_manual-tls")]
 pub async fn fetch_ohttp_keys_with_cert(
-    ohttp_relay: impl IntoUrl,
+    relays: &[Url],
     payjoin_directory: impl IntoUrl,
     cert_der: &[u8],
+) -> Result<(OhttpKeys, Url), Error> {
+    fetch_ohttp_keys_inner(relays, payjoin_directory.into_url()?, Some(cert_der)).await
+}
+
+async fn fetch_ohttp_keys_inner(
+    relays: &[Url],
+    payjoin_directory: Url,
+    cert_der: Option<&[u8]>,
+) -> Result<(OhttpKeys, Url), Error> {
+    let ohttp_keys_url = payjoin_directory.join("/.well-known/ohttp-gateway")?;
+    let mut selector = RelaySelector::new(relays.to_vec());
+    let mut last_err: Option<Error> = None;
+    loop {
+        let relay = match selector.select(&mut rand::thread_rng()) {
+            Some(relay) => relay,
+            None => return Err(last_err.unwrap_or(Error::NoRelaysAvailable)),
+        };
+        match fetch_keys_once(&relay, &ohttp_keys_url, cert_der).await {
+            Ok(keys) => return Ok((keys, relay)),
+            Err(e @ Error::UnexpectedStatusCode(_)) => return Err(e),
+            Err(e) => {
+                selector.mark_failed(&relay);
+                last_err = Some(e);
+            }
+        }
+    }
+}
+
+async fn fetch_keys_once(
+    relay: &Url,
+    ohttp_keys_url: &Url,
+    cert_der: Option<&[u8]>,
 ) -> Result<OhttpKeys, Error> {
-    let ohttp_keys_url = payjoin_directory.into_url()?.join("/.well-known/ohttp-gateway")?;
-    let proxy = Proxy::all(ohttp_relay.into_url()?.as_str())?;
-    let client = Client::builder()
-        .use_rustls_tls()
-        .add_root_certificate(reqwest::tls::Certificate::from_der(cert_der)?)
-        .proxy(proxy)
-        .http1_only()
-        .build()?;
+    #[cfg(not(feature = "_manual-tls"))]
+    let _ = cert_der;
+    let proxy = Proxy::all(relay.as_str())?;
+    let builder = Client::builder().proxy(proxy).http1_only();
+    #[cfg(feature = "_manual-tls")]
+    let builder = match cert_der {
+        Some(cert_der) => builder
+            .use_rustls_tls()
+            .add_root_certificate(reqwest::tls::Certificate::from_der(cert_der)?),
+        None => builder,
+    };
+    let client = builder.build()?;
     let res = client
         .get(ohttp_keys_url.as_str())
         .timeout(Duration::from_secs(10))
@@ -80,6 +113,8 @@ async fn parse_ohttp_keys_response(res: reqwest::Response) -> Result<OhttpKeys,
 pub enum Error {
     /// When the payjoin directory returns an unexpected status code
     UnexpectedStatusCode(http::StatusCode),
+    /// No relay was available to reach the payjoin directory.
+    NoRelaysAvailable,
     /// Internal errors that should not be pattern matched by users
     #[doc(hidden)]
     Internal(InternalError),
@@ -126,6 +161,7 @@ impl std::fmt::Display for Error {
             Self::UnexpectedStatusCode(code) => {
                 write!(f, "Unexpected status code from payjoin directory: {code}")
             }
+            Self::NoRelaysAvailable => write!(f, "No relay was available to fetch ohttp keys"),
             Self::Internal(InternalError(e)) => e.fmt(f),
         }
     }
@@ -153,6 +189,7 @@ impl std::error::Error for Error {
         match self {
             Self::Internal(InternalError(e)) => e.source(),
             Self::UnexpectedStatusCode(_) => None,
+            Self::NoRelaysAvailable => None,
         }
     }
 }
@@ -240,4 +277,27 @@ mod tests {
             "expected InvalidOhttpKeys error"
         );
     }
+
+    #[tokio::test]
+    async fn fetch_with_no_relays_returns_no_relays_available() {
+        assert!(matches!(
+            fetch_ohttp_keys(&[], "https://directory.example").await,
+            Err(Error::NoRelaysAvailable)
+        ));
+    }
+
+    #[tokio::test]
+    async fn fetch_exhausts_unreachable_relays_and_surfaces_last_error() {
+        // Every relay fails to connect: the loop marks each failed and terminates
+        // once none remain, surfacing the last transport error. NoRelaysAvailable
+        // is reserved for an empty relay list, where no attempt was ever made.
+        let relays = [
+            Url::parse("http://127.0.0.1:1").expect("valid url"),
+            Url::parse("http://127.0.0.1:2").expect("valid url"),
+        ];
+        assert!(matches!(
+            fetch_ohttp_keys(&relays, "https://directory.example").await,
+            Err(Error::Internal(_))
+        ));
+    }
 }

Original file line number	Diff line number	Diff line change
`@@ -12,6 +12,7 @@ use payjoin::receive::v2::{`
`12`	`12`	`ReceiverBuilder, SessionOutcome as ReceiverSessionOutcome, UncheckedOriginalPayload,`
`13`	`13`	`WantsFeeRange, WantsInputs, WantsOutputs,`
`14`	`14`	`};`
	`15`	`+use payjoin::relay::RelaySelector;`
`15`	`16`	`use payjoin::send::v2::{`
`16`	`17`	`replay_event_log as replay_sender_event_log, PendingFallback as SenderPendingFallback,`
`17`	`18`	`PollingForProposal, SendSession, Sender, SenderBuilder, SessionOutcome as SenderSessionOutcome,`
`@@ -278,7 +279,6 @@ impl AppTrait for App {`
`278`	`279`	`Err(e) => {`
`279`	`280`	`tracing::debug!("Directory {directory} failed: {e:#}");`
`280`	`281`	`self.mailroom_manager.add_failed_directory(directory);`
`281`		`- self.mailroom_manager.clear_failed_relays();`
`282`	`282`	`continue;`
`283`	`283`	`}`
`284`	`284`	`}`
`@@ -1055,14 +1055,17 @@ impl App {`
`1055`	`1055`	`F: FnMut(&str) -> std::result::Result<(payjoin::Request, T), E>,`
`1056`	`1056`	`E: Into<anyhow::Error>,`
`1057`	`1057`	`{`
	`1058`	`+ let mut selector = RelaySelector::new(self.config.v2()?.ohttp_relays.clone());`
`1058`	`1059`	`loop {`
`1059`		`- let relay = self.mailroom_manager.choose_relay()?;`
	`1060`	`+ let relay = selector`
	`1061`	`+ .select(&mut payjoin::bitcoin::key::rand::thread_rng())`
	`1062`	`+ .ok_or_else(\|\| anyhow!("No valid relays available"))?;`
`1060`	`1063`	`let (req, ctx) = build(relay.as_str()).map_err(Into::into)?;`
`1061`	`1064`	`match self.post_request(req).await {`
`1062`	`1065`	`Ok(resp) => return Ok((resp, ctx)),`
`1063`	`1066`	`Err(e) => {`
`1064`	`1067`	`tracing::debug!("Request to relay {relay} failed: {e:?}");`
`1065`		`- self.mailroom_manager.add_failed_relay(relay);`
	`1068`	`+ selector.mark_failed(&relay);`
`1066`	`1069`	`}`
`1067`	`1070`	`}`
`1068`	`1071`	`}`