Thread SelectContext through relay selection

Add a SelectContext argument to select() and pass an empty, random
context at each call site.

Author: bc1cindy

.cargo/mutants.toml

@@ -12,6 +12,8 @@ exclude_re = [
   # Timeout loops
   # src/receive/v1/mod.rs
   "interleave_shuffle", # Replacing index += 1 with index *= 1 in a loop causes a timeout due to an infinite loop
+  # src/core/relay.rs
+  "RelaySelector::mark_failed", # A no-op makes fetch_ohttp_keys_inner never exhaust relays, looping until timeout
 
   # Trivial mutations
   # These exclusions are allowing code blocks to run with arithmetic involving zero and as a result are no-ops

payjoin-cli/src/app/v2/mod.rs

@@ -1084,9 +1084,11 @@ impl App {
         E: Into<anyhow::Error>,
     {
         let mut selector = RelaySelector::new(self.config.v2()?.ohttp_relays.clone());
+        // A later policy fills Post/Poll, receiver key and time window in here.
+        let select_ctx = payjoin::relay::SelectContext::random();
         loop {
             let relay = selector
-                .select(&mut payjoin::bitcoin::key::rand::thread_rng())
+                .select(&select_ctx, &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 {

payjoin/src/core/io.rs

@@ -6,7 +6,7 @@ use http::header::ACCEPT;
 use reqwest::{Client, Proxy};
 
 use crate::into_url::IntoUrl;
-use crate::relay::RelaySelector;
+use crate::relay::{RelaySelector, SelectContext};
 use crate::{OhttpKeys, Url};
 
 /// Fetch the ohttp keys from the specified payjoin directory via proxy.
@@ -54,9 +54,10 @@ async fn fetch_ohttp_keys_inner(
 ) -> Result<(OhttpKeys, Url), Error> {
     let ohttp_keys_url = payjoin_directory.join("/.well-known/ohttp-gateway")?;
     let mut selector = RelaySelector::new(relays.to_vec());
+    let ctx = SelectContext::random();
     let mut last_err: Option<Error> = None;
     loop {
-        let relay = match selector.select(&mut rand::thread_rng()) {
+        let relay = match selector.select(&ctx, &mut rand::thread_rng()) {
             Some(relay) => relay,
             None => return Err(last_err.unwrap_or(Error::NoRelaysAvailable)),
         };

payjoin/src/core/relay.rs

@@ -3,8 +3,19 @@ use bitcoin::secp256k1::rand::{self};
 
 use crate::Url;
 
-/// Picks an OHTTP relay at random, excluding relays marked failed, so clients
-/// share one selection policy instead of each diverging.
+/// Per-request inputs for [`RelaySelector::select`]. Empty, so selection is
+/// unconstrained and uniform-random.
+#[derive(Clone, Debug)]
+#[non_exhaustive]
+pub struct SelectContext {}
+
+impl SelectContext {
+    /// An unconstrained context: selection is uniform-random.
+    pub fn random() -> Self { Self {} }
+}
+
+/// Picks an OHTTP relay, excluding relays marked failed, so clients share one
+/// selection policy instead of each diverging.
 #[derive(Clone, Debug)]
 pub struct RelaySelector {
     relays: Vec<Url>,
@@ -14,11 +25,12 @@ pub struct RelaySelector {
 impl RelaySelector {
     pub fn new(relays: Vec<Url>) -> Self { Self { relays, failed: Vec::new() } }
 
-    /// Pick a random relay not marked failed, or `None` when none remain.
-    pub fn select<R: rand::Rng>(&self, rng: &mut R) -> Option<Url> {
+    /// Pick a relay for `ctx`, never one marked failed, or `None` when none remain.
+    pub fn select<R: rand::Rng>(&self, _ctx: &SelectContext, rng: &mut R) -> Option<Url> {
         self.relays.iter().filter(|r| !self.failed.contains(r)).choose(rng).cloned()
     }
 
+    /// Record a relay transport failure so `select` avoids it.
     pub fn mark_failed(&mut self, relay: &Url) {
         if !self.failed.contains(relay) {
             self.failed.push(relay.clone());
@@ -44,7 +56,7 @@ mod tests {
     fn select_returns_a_configured_relay() {
         let selector = RelaySelector::new(relays());
         let mut rng = StdRng::seed_from_u64(1);
-        let picked = selector.select(&mut rng).expect("a relay");
+        let picked = selector.select(&SelectContext::random(), &mut rng).expect("a relay");
         assert!(relays().contains(&picked));
     }
 
@@ -55,7 +67,7 @@ mod tests {
         let failed = Url::parse("https://a.example").unwrap();
         selector.mark_failed(&failed);
         for _ in 0..50 {
-            assert_ne!(selector.select(&mut rng), Some(failed.clone()));
+            assert_ne!(selector.select(&SelectContext::random(), &mut rng), Some(failed.clone()));
         }
     }
 
@@ -66,13 +78,27 @@ mod tests {
             selector.mark_failed(&r);
         }
         let mut rng = StdRng::seed_from_u64(3);
-        assert_eq!(selector.select(&mut rng), None);
+        assert_eq!(selector.select(&SelectContext::random(), &mut rng), None);
     }
 
     #[test]
     fn select_is_none_when_empty() {
         let selector = RelaySelector::new(Vec::new());
         let mut rng = StdRng::seed_from_u64(4);
-        assert_eq!(selector.select(&mut rng), None);
+        assert_eq!(selector.select(&SelectContext::random(), &mut rng), None);
+    }
+
+    // An empty context selects uniformly across all relays.
+    #[test]
+    fn random_context_is_unconstrained() {
+        let selector = RelaySelector::new(relays());
+        let mut seen = std::collections::BTreeSet::new();
+        let mut rng = StdRng::seed_from_u64(5);
+        for _ in 0..200 {
+            if let Some(r) = selector.select(&SelectContext::random(), &mut rng) {
+                seen.insert(r.to_string());
+            }
+        }
+        assert_eq!(seen.len(), relays().len(), "random must reach every relay");
     }
 }