From 2490ef203191ae08328cc4c418aa3dbf6ea05347 Mon Sep 17 00:00:00 2001
From: Armin Sabouri <armins88@gmail.com>
Date: Thu, 28 Aug 2025 14:44:54 -0400
Subject: [PATCH] Look up RK instead of looping

Follow up from https://github.com/payjoin/rust-payjoin/pull/995.
We can replace a loop over all active session ids with a WHERE clause.
---
 payjoin-cli/src/app/v2/mod.rs | 27 ++++++++++-----------------
 payjoin-cli/src/db/v2.rs      | 24 +++++++++++++++---------
 2 files changed, 25 insertions(+), 26 deletions(-)

diff --git a/payjoin-cli/src/app/v2/mod.rs b/payjoin-cli/src/app/v2/mod.rs
index 0eb4a044d..6c0ceabfd 100644
--- a/payjoin-cli/src/app/v2/mod.rs
+++ b/payjoin-cli/src/app/v2/mod.rs
@@ -206,23 +206,16 @@ impl AppTrait for App {
             }
             PjParam::V2(pj_param) => {
                 let receiver_pubkey = pj_param.receiver_pubkey();
-                let sender_state =
-                    self.db.get_send_session_ids()?.into_iter().find_map(|session_id| {
-                        let session_receiver_pubkey = self
-                            .db
-                            .get_send_session_receiver_pk(&session_id)
-                            .expect("Receiver pubkey should exist if session id exists");
-                        if session_receiver_pubkey == *receiver_pubkey {
-                            let sender_persister =
-                                SenderPersister::from_id(self.db.clone(), session_id);
-                            let (send_session, _) = replay_sender_event_log(&sender_persister)
-                                .map_err(|e| anyhow!("Failed to replay sender event log: {:?}", e))
-                                .ok()?;
-
-                            Some((send_session, sender_persister))
-                        } else {
-                            None
-                        }
+                let sender_state = self
+                    .db
+                    .get_send_session_id_by_receiver_pk(receiver_pubkey)?
+                    .and_then(|session_id| {
+                        let sender_persister =
+                            SenderPersister::from_id(self.db.clone(), session_id);
+                        let (send_session, _) = replay_sender_event_log(&sender_persister)
+                            .map_err(|e| anyhow!("Failed to replay sender event log: {:?}", e))
+                            .ok()?;
+                        Some((send_session, sender_persister))
                     });
 
                 let (sender_state, persister) = match sender_state {
diff --git a/payjoin-cli/src/db/v2.rs b/payjoin-cli/src/db/v2.rs
index 7ec7b8e8c..8ac33bd71 100644
--- a/payjoin-cli/src/db/v2.rs
+++ b/payjoin-cli/src/db/v2.rs
@@ -4,7 +4,7 @@ use payjoin::persist::SessionPersister;
 use payjoin::receive::v2::SessionEvent as ReceiverSessionEvent;
 use payjoin::send::v2::SessionEvent as SenderSessionEvent;
 use payjoin::HpkePublicKey;
-use rusqlite::params;
+use rusqlite::{params, OptionalExtension};
 
 use super::*;
 
@@ -239,16 +239,22 @@ impl Database {
         Ok(session_ids)
     }
 
-    pub(crate) fn get_send_session_receiver_pk(
+    pub(crate) fn get_send_session_id_by_receiver_pk(
         &self,
-        session_id: &SessionId,
-    ) -> Result<HpkePublicKey> {
+        receiver_pubkey: &HpkePublicKey,
+    ) -> Result<Option<SessionId>> {
         let conn = self.get_connection()?;
-        let mut stmt =
-            conn.prepare("SELECT receiver_pubkey FROM send_sessions WHERE session_id = ?1")?;
-        let receiver_pubkey: Vec<u8> =
-            stmt.query_row(params![session_id.0.to_string()], |row| row.get(0))?;
-        Ok(HpkePublicKey::from_compressed_bytes(&receiver_pubkey).expect("Valid receiver pubkey"))
+        let receiver_pubkey_bytes = receiver_pubkey.to_compressed_bytes();
+        let mut stmt = conn.prepare(
+            "SELECT session_id FROM send_sessions WHERE receiver_pubkey = ?1 AND completed_at IS NULL",
+        )?;
+        let result = stmt.query_row(params![&receiver_pubkey_bytes], |row| {
+            let session_id: String = row.get(0)?;
+            let session_id = uuid::Uuid::parse_str(&session_id)
+                .expect("Database corruption: invalid session_id UUID");
+            Ok(SessionId(session_id))
+        });
+        Ok(result.optional()?)
     }
 
     pub(crate) fn get_inactive_send_session_ids(&self) -> Result<Vec<(SessionId, u64)>> {