1use std::collections::{BTreeMap, HashMap, HashSet};
2use std::path::{Path, PathBuf};
3use std::sync::mpsc as std_mpsc;
4use std::sync::{Arc, Mutex, MutexGuard};
5use std::thread;
6use std::time::Duration as StdDuration;
7
8use anyhow::{anyhow, Context, Result};
9use futures_util::{SinkExt, StreamExt};
10use serde_json::{json, Value};
11use tokio::net::TcpListener;
12use tokio::sync::mpsc;
13use tokio_tungstenite::accept_async;
14use tokio_tungstenite::tungstenite::Message;
15
16#[derive(Default)]
17struct RelayState {
18 #[cfg(test)]
19 wire: Option<Vec<(usize, Value)>>,
20 events_by_id: BTreeMap<String, Value>,
21 subscriptions: HashMap<usize, HashMap<String, Vec<Value>>>,
22 clients: HashMap<usize, mpsc::UnboundedSender<Message>>,
23 faults: RelayFaults,
24 dropped_event_ids: HashSet<String>,
25}
26
27#[derive(Clone, Default)]
28struct RelayFaults {
29 drop_event_ids_file: Option<PathBuf>,
30 drop_matching_events_once: bool,
31 unacknowledged_kinds: HashSet<u64>,
32}
33
34impl RelayState {
35 fn from_env() -> Self {
36 Self {
37 faults: RelayFaults::from_env(),
38 ..Self::default()
39 }
40 }
41
42 fn should_drop_event(&mut self, event_id: &str) -> bool {
43 let Some(path) = self.faults.drop_event_ids_file.as_ref() else {
44 return false;
45 };
46 if self.faults.drop_matching_events_once && self.dropped_event_ids.contains(event_id) {
47 return false;
48 }
49 if !drop_event_ids(path).contains(event_id) {
50 return false;
51 }
52 self.dropped_event_ids.insert(event_id.to_string());
53 true
54 }
55}
56
57impl RelayFaults {
58 fn from_env() -> Self {
59 let drop_event_ids_file = std::env::var_os("IRIS_LOCAL_RELAY_DROP_EVENT_IDS_FILE")
60 .filter(|value| !value.is_empty())
61 .map(PathBuf::from);
62 let drop_matching_events_once = !env_flag("IRIS_LOCAL_RELAY_DROP_EVENT_IDS_ALWAYS");
63 Self {
64 drop_event_ids_file,
65 drop_matching_events_once,
66 unacknowledged_kinds: HashSet::new(),
67 }
68 }
69}
70
71fn env_flag(name: &str) -> bool {
72 matches!(
73 std::env::var(name)
74 .unwrap_or_default()
75 .trim()
76 .to_ascii_lowercase()
77 .as_str(),
78 "1" | "true" | "yes" | "on"
79 )
80}
81
82fn drop_event_ids(path: &Path) -> HashSet<String> {
83 let Ok(raw) = std::fs::read_to_string(path) else {
84 return HashSet::new();
85 };
86 raw.lines()
87 .filter_map(|line| line.split('#').next())
88 .map(str::trim)
89 .filter(|line| !line.is_empty())
90 .map(str::to_string)
91 .collect()
92}
93
94fn lock_relay_state(state: &Arc<Mutex<RelayState>>) -> MutexGuard<'_, RelayState> {
95 state.lock().unwrap_or_else(|poison| poison.into_inner())
96}
97
98enum RelayControl {
99 #[cfg(test)]
100 Wire(std_mpsc::Sender<Vec<(usize, Value)>>),
101 #[cfg(test)]
102 CloseSubscription(String, String, std_mpsc::Sender<usize>),
103 #[cfg(test)]
104 DisconnectClients(std_mpsc::Sender<usize>),
105 IgnoreAcknowledgements(u64, std_mpsc::Sender<()>),
106 ReplayStored,
107 Snapshot(std_mpsc::Sender<Vec<Value>>),
108 Shutdown,
109}
110
111pub struct TestRelay {
112 control_tx: mpsc::UnboundedSender<RelayControl>,
113 join: Option<thread::JoinHandle<()>>,
114 url: String,
115}
116
117impl TestRelay {
118 pub fn start() -> Self {
119 match Self::start_with_bind("127.0.0.1:0") {
120 Ok(relay) => relay,
121 Err(error) => {
122 eprintln!("failed to start local relay: {error}");
123 let (control_tx, _) = mpsc::unbounded_channel();
124 Self {
125 control_tx,
126 join: None,
127 url: String::new(),
128 }
129 }
130 }
131 }
132
133 pub fn start_with_bind(bind_addr: &str) -> Result<Self> {
134 let (control_tx, mut control_rx) = mpsc::unbounded_channel();
135 let (ready_tx, ready_rx) = std_mpsc::channel();
136 let bind_addr = bind_addr.to_string();
137
138 let join = thread::spawn(move || {
139 let runtime = match tokio::runtime::Builder::new_multi_thread()
140 .enable_all()
141 .build()
142 {
143 Ok(runtime) => runtime,
144 Err(error) => {
145 let _ = ready_tx.send(Err(anyhow!("relay runtime: {error}")));
146 return;
147 }
148 };
149
150 runtime.block_on(async move {
151 let listener = match TcpListener::bind(&bind_addr)
152 .await
153 .with_context(|| format!("bind relay listener {bind_addr}"))
154 {
155 Ok(listener) => listener,
156 Err(error) => {
157 let _ = ready_tx.send(Err(error));
158 return;
159 }
160 };
161 let local_addr = match listener.local_addr() {
162 Ok(addr) => addr,
163 Err(error) => {
164 let _ = ready_tx.send(Err(anyhow!("relay local addr: {error}")));
165 return;
166 }
167 };
168 let state = Arc::new(Mutex::new(RelayState::default()));
169 let next_client_id = Arc::new(std::sync::atomic::AtomicUsize::new(1));
170 let _ = ready_tx.send(Ok(format!("ws://{local_addr}")));
171
172 loop {
173 tokio::select! {
174 Some(control) = control_rx.recv() => {
175 match control {
176 #[cfg(test)]
177 RelayControl::Wire(reply) => {
178 let _ = reply.send(lock_relay_state(&state).wire.get_or_insert_default().clone());
179 }
180 #[cfg(test)]
181 RelayControl::DisconnectClients(reply) => {
182 let state = lock_relay_state(&state);
183 for sender in state.clients.values() {
184 let _ = sender.send(Message::Close(None));
185 }
186 let _ = reply.send(state.clients.len());
187 }
188 #[cfg(test)]
189 RelayControl::CloseSubscription(id, reason, reply) => {
190 let mut state = lock_relay_state(&state);
191 let clients = state.subscriptions.iter_mut()
192 .filter_map(|(client, subscriptions)| subscriptions.remove(&id).map(|_| *client))
193 .collect::<Vec<_>>();
194 for client in &clients {
195 if let Some(sender) = state.clients.get(client) {
196 let _ = sender.send(Message::Text(json!(["CLOSED", id, reason]).to_string()));
197 }
198 }
199 let _ = reply.send(clients.len());
200 }
201 RelayControl::IgnoreAcknowledgements(kind, reply) => {
202 lock_relay_state(&state).faults.unacknowledged_kinds.insert(kind);
203 let _ = reply.send(());
204 }
205 RelayControl::ReplayStored => replay_stored_events(&state),
206 RelayControl::Snapshot(reply_tx) => {
207 let events = lock_relay_state(&state)
208 .events_by_id
209 .values()
210 .cloned()
211 .collect::<Vec<_>>();
212 let _ = reply_tx.send(events);
213 }
214 RelayControl::Shutdown => break,
215 }
216 }
217 accept_result = listener.accept() => {
218 let Ok((stream, _)) = accept_result else {
219 break;
220 };
221 let websocket = match accept_async(stream).await {
222 Ok(websocket) => websocket,
223 Err(error) => {
224 eprintln!("Ignoring failed test relay websocket handshake: {error}");
225 continue;
226 }
227 };
228 let state = state.clone();
229 let client_id = next_client_id.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
230 tokio::spawn(async move {
231 handle_connection(client_id, websocket, state).await;
232 });
233 }
234 }
235 }
236 });
237 });
238
239 let url = ready_rx
240 .recv_timeout(StdDuration::from_secs(5))
241 .context("relay ready")??;
242
243 Ok(Self {
244 control_tx,
245 join: Some(join),
246 url,
247 })
248 }
249
250 pub fn url(&self) -> &str {
251 &self.url
252 }
253
254 pub fn ignore_acknowledgements(&self, kind: u64) -> anyhow::Result<()> {
255 let (tx, rx) = std_mpsc::channel();
256 self.control_tx
257 .send(RelayControl::IgnoreAcknowledgements(kind, tx))
258 .context("configure test relay acknowledgements")?;
259 rx.recv_timeout(StdDuration::from_secs(5))
260 .context("test relay acknowledgement configuration")?;
261 Ok(())
262 }
263
264 pub fn replay_stored(&self) {
265 let _ = self.control_tx.send(RelayControl::ReplayStored);
266 }
267
268 pub fn events(&self) -> Vec<Value> {
269 let (reply_tx, reply_rx) = std_mpsc::channel();
270 let _ = self.control_tx.send(RelayControl::Snapshot(reply_tx));
271 reply_rx
272 .recv_timeout(StdDuration::from_secs(5))
273 .unwrap_or_default()
274 }
275
276 #[cfg(test)]
277 pub(crate) fn wire(&self) -> Vec<(usize, Value)> {
279 let (tx, rx) = std_mpsc::channel();
280 self.control_tx
281 .send(RelayControl::Wire(tx))
282 .expect("live test relay");
283 rx.recv_timeout(StdDuration::from_secs(5))
284 .expect("relay wire snapshot")
285 }
286
287 #[cfg(test)]
288 pub(crate) fn close_subscription(&self, id: &str, reason: &str) -> usize {
289 let (tx, rx) = std_mpsc::channel();
290 self.control_tx
291 .send(RelayControl::CloseSubscription(
292 id.into(),
293 reason.into(),
294 tx,
295 ))
296 .expect("live test relay");
297 rx.recv_timeout(StdDuration::from_secs(5))
298 .expect("relay closed subscription")
299 }
300
301 #[cfg(test)]
302 pub(crate) fn disconnect_clients(&self) -> usize {
303 let (tx, rx) = std_mpsc::channel();
304 self.control_tx
305 .send(RelayControl::DisconnectClients(tx))
306 .expect("live test relay");
307 rx.recv_timeout(StdDuration::from_secs(5))
308 .expect("relay disconnect")
309 }
310}
311
312impl Drop for TestRelay {
313 fn drop(&mut self) {
314 let _ = self.control_tx.send(RelayControl::Shutdown);
315 if let Some(join) = self.join.take() {
316 let _ = join.join();
317 }
318 }
319}
320
321pub fn run_forever(bind_addr: &str) -> Result<()> {
322 let runtime = tokio::runtime::Builder::new_multi_thread()
323 .enable_all()
324 .build()
325 .context("relay runtime")?;
326 let bind_addr = bind_addr.to_string();
327
328 runtime.block_on(async move {
329 let listener = TcpListener::bind(&bind_addr)
330 .await
331 .with_context(|| format!("bind relay listener {bind_addr}"))?;
332 let state = Arc::new(Mutex::new(RelayState::from_env()));
333 let next_client_id = Arc::new(std::sync::atomic::AtomicUsize::new(1));
334
335 println!("Local Nostr relay listening on ws://{bind_addr}");
336
337 loop {
338 let (stream, _) = listener
339 .accept()
340 .await
341 .with_context(|| format!("accept relay client on {bind_addr}"))?;
342 let websocket = match accept_async(stream).await {
343 Ok(websocket) => websocket,
344 Err(error) => {
345 eprintln!("Ignoring failed websocket handshake on {bind_addr}: {error}");
346 continue;
347 }
348 };
349 let state = state.clone();
350 let client_id = next_client_id.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
351 tokio::spawn(async move {
352 handle_connection(client_id, websocket, state).await;
353 });
354 }
355 })
356}
357
358async fn handle_connection(
359 client_id: usize,
360 websocket: tokio_tungstenite::WebSocketStream<tokio::net::TcpStream>,
361 state: Arc<Mutex<RelayState>>,
362) {
363 let (mut sink, mut stream) = websocket.split();
364 let (client_tx, mut client_rx) = mpsc::unbounded_channel::<Message>();
365
366 {
367 let mut relay = lock_relay_state(&state);
368 relay.clients.insert(client_id, client_tx);
369 }
370
371 let writer = tokio::spawn(async move {
372 while let Some(message) = client_rx.recv().await {
373 if sink.send(message).await.is_err() {
374 break;
375 }
376 }
377 });
378
379 while let Some(message) = stream.next().await {
380 let Ok(message) = message else {
381 break;
382 };
383 match message {
384 Message::Text(text) => handle_client_message(client_id, &text, &state),
385 Message::Ping(payload) => {
386 let sender = {
387 let relay = lock_relay_state(&state);
388 relay.clients.get(&client_id).cloned()
389 };
390 if let Some(sender) = sender {
391 let _ = sender.send(Message::Pong(payload));
392 }
393 }
394 Message::Close(_) => break,
395 _ => {}
396 }
397 }
398
399 {
400 let mut relay = lock_relay_state(&state);
401 relay.clients.remove(&client_id);
402 relay.subscriptions.remove(&client_id);
403 }
404
405 writer.abort();
406}
407
408fn handle_client_message(client_id: usize, raw_message: &str, state: &Arc<Mutex<RelayState>>) {
409 let Ok(message) = serde_json::from_str::<Value>(raw_message) else {
410 return;
411 };
412 let Some(parts) = message.as_array() else {
413 return;
414 };
415 let Some(kind) = parts.first().and_then(Value::as_str) else {
416 return;
417 };
418 #[cfg(test)]
419 if let Some(wire) = lock_relay_state(state).wire.as_mut() {
420 wire.push((client_id, message.clone()));
421 }
422
423 match kind {
424 "REQ" if parts.len() >= 2 => {
425 let Some(subscription_id) = parts.get(1).and_then(Value::as_str) else {
426 return;
427 };
428 let filters: Vec<Value> = parts
429 .iter()
430 .skip(2)
431 .filter(|value| value.is_object())
432 .cloned()
433 .collect();
434 let (sender, events) = {
435 let mut relay = lock_relay_state(state);
436 relay
437 .subscriptions
438 .entry(client_id)
439 .or_default()
440 .insert(subscription_id.to_string(), filters.clone());
441 (
442 relay.clients.get(&client_id).cloned(),
443 relay.events_by_id.values().cloned().collect::<Vec<_>>(),
444 )
445 };
446
447 if let Some(sender) = sender {
448 for event in events {
449 if matches_any_filter(&event, &filters) {
450 let payload =
451 Message::Text(json!(["EVENT", subscription_id, event]).to_string());
452 let _ = sender.send(payload);
453 }
454 }
455 let _ = sender.send(Message::Text(json!(["EOSE", subscription_id]).to_string()));
456 }
457 }
458 "CLOSE" if parts.len() >= 2 => {
459 let Some(subscription_id) = parts.get(1).and_then(Value::as_str) else {
460 return;
461 };
462 let mut relay = lock_relay_state(state);
463 if let Some(subscriptions) = relay.subscriptions.get_mut(&client_id) {
464 subscriptions.remove(subscription_id);
465 }
466 }
467 "EVENT" if parts.get(1).is_some_and(Value::is_object) => {
468 let Some(event) = parts.get(1).cloned() else {
469 return;
470 };
471 let Some(event_id) = event.get("id").and_then(Value::as_str) else {
472 return;
473 };
474 let event_id = event_id.to_string();
475 let (sender, deliveries, dropped) = {
476 let mut relay = lock_relay_state(state);
477 let sender = if event
478 .get("kind")
479 .and_then(Value::as_u64)
480 .is_some_and(|kind| relay.faults.unacknowledged_kinds.contains(&kind))
481 {
482 None
483 } else {
484 relay.clients.get(&client_id).cloned()
485 };
486 if relay.should_drop_event(&event_id) {
487 (sender, Vec::new(), true)
488 } else {
489 relay.events_by_id.insert(event_id.clone(), event.clone());
490 let deliveries = matching_deliveries(&relay, &event);
491 (sender, deliveries, false)
492 }
493 };
494 if dropped {
495 eprintln!("Local relay fault dropped event_id={event_id}");
496 }
497 if let Some(sender) = sender {
498 let message = if dropped {
499 "fault: dropped by local relay"
500 } else {
501 ""
502 };
503 let _ = sender.send(Message::Text(
504 json!(["OK", event_id, true, message]).to_string(),
505 ));
506 }
507 if dropped {
508 return;
509 }
510
511 for (target, payload) in deliveries {
512 let _ = target.send(payload);
513 }
514 }
515 _ => {}
516 }
517}
518
519fn replay_stored_events(state: &Arc<Mutex<RelayState>>) {
520 let deliveries = {
521 let relay = lock_relay_state(state);
522 relay
523 .events_by_id
524 .values()
525 .flat_map(|event| matching_deliveries(&relay, event))
526 .collect::<Vec<_>>()
527 };
528
529 for (target, payload) in deliveries {
530 let _ = target.send(payload);
531 }
532}
533
534fn matching_deliveries(
535 relay: &RelayState,
536 event: &Value,
537) -> Vec<(mpsc::UnboundedSender<Message>, Message)> {
538 let mut deliveries = Vec::new();
539 for (client_id, subscriptions) in &relay.subscriptions {
540 let Some(target) = relay.clients.get(client_id).cloned() else {
541 continue;
542 };
543 for (subscription_id, filters) in subscriptions {
544 if matches_any_filter(event, filters) {
545 deliveries.push((
546 target.clone(),
547 Message::Text(json!(["EVENT", subscription_id, event]).to_string()),
548 ));
549 }
550 }
551 }
552 deliveries
553}
554
555pub fn matches_any_filter(event: &Value, filters: &[Value]) -> bool {
556 if filters.is_empty() {
557 return true;
558 }
559
560 filters.iter().any(|filter| matches_filter(event, filter))
561}
562
563pub fn matches_filter(event: &Value, filter: &Value) -> bool {
564 let Some(filter_object) = filter.as_object() else {
565 return false;
566 };
567
568 if let Some(ids) = filter_object.get("ids").and_then(Value::as_array) {
569 let Some(event_id) = event.get("id").and_then(Value::as_str) else {
570 return false;
571 };
572 if !ids
573 .iter()
574 .filter_map(Value::as_str)
575 .any(|id| id == event_id)
576 {
577 return false;
578 }
579 }
580
581 if let Some(authors) = filter_object.get("authors").and_then(Value::as_array) {
582 let Some(pubkey) = event.get("pubkey").and_then(Value::as_str) else {
583 return false;
584 };
585 if !authors
586 .iter()
587 .filter_map(Value::as_str)
588 .any(|author| author == pubkey)
589 {
590 return false;
591 }
592 }
593
594 if let Some(kinds) = filter_object.get("kinds").and_then(Value::as_array) {
595 let Some(kind) = event.get("kind").and_then(Value::as_u64) else {
596 return false;
597 };
598 if !kinds
599 .iter()
600 .filter_map(Value::as_u64)
601 .any(|value| value == kind)
602 {
603 return false;
604 }
605 }
606
607 if let Some(since) = filter_object.get("since").and_then(Value::as_u64) {
608 let Some(created_at) = event.get("created_at").and_then(Value::as_u64) else {
609 return false;
610 };
611 if created_at < since {
612 return false;
613 }
614 }
615
616 if let Some(until) = filter_object.get("until").and_then(Value::as_u64) {
617 let Some(created_at) = event.get("created_at").and_then(Value::as_u64) else {
618 return false;
619 };
620 if created_at > until {
621 return false;
622 }
623 }
624
625 for (key, value) in filter_object {
626 let Some(tag_name) = key.strip_prefix('#') else {
627 continue;
628 };
629
630 let Some(expected_values) = value.as_array() else {
631 return false;
632 };
633 if expected_values.is_empty() {
634 continue;
635 }
636
637 let Some(tags) = event.get("tags").and_then(Value::as_array) else {
638 return false;
639 };
640 let matched = tags.iter().any(|tag| {
641 let Some(tag_values) = tag.as_array() else {
642 return false;
643 };
644 if tag_values.first().and_then(Value::as_str) != Some(tag_name) {
645 return false;
646 }
647 tag_values
648 .iter()
649 .skip(1)
650 .filter_map(Value::as_str)
651 .any(|tag_value| {
652 expected_values
653 .iter()
654 .filter_map(Value::as_str)
655 .any(|expected| expected == tag_value)
656 })
657 });
658 if !matched {
659 return false;
660 }
661 }
662
663 true
664}
665
666#[cfg(test)]
667mod tests {
668 use super::*;
669 use std::io::Write;
670
671 #[test]
672 fn drop_event_ids_file_ignores_comments_and_blank_lines() {
673 let mut file = tempfile::NamedTempFile::new().expect("temp drop file");
674 writeln!(file, "\n# comment\nabc\n def # inline comment\n").expect("write drop file");
675
676 let ids = drop_event_ids(file.path());
677
678 assert!(ids.contains("abc"));
679 assert!(ids.contains("def"));
680 assert!(!ids.contains("# comment"));
681 }
682
683 #[test]
684 fn relay_fault_drops_matching_event_once_by_default() {
685 let mut file = tempfile::NamedTempFile::new().expect("temp drop file");
686 writeln!(file, "event-to-drop").expect("write drop file");
687 let mut state = RelayState {
688 faults: RelayFaults {
689 drop_event_ids_file: Some(file.path().to_path_buf()),
690 drop_matching_events_once: true,
691 ..RelayFaults::default()
692 },
693 ..RelayState::default()
694 };
695
696 assert!(state.should_drop_event("event-to-drop"));
697 assert!(!state.should_drop_event("event-to-drop"));
698 assert!(!state.should_drop_event("different-event"));
699 }
700
701 #[test]
702 fn relay_fault_can_drop_matching_event_every_time() {
703 let mut file = tempfile::NamedTempFile::new().expect("temp drop file");
704 writeln!(file, "event-to-drop").expect("write drop file");
705 let mut state = RelayState {
706 faults: RelayFaults {
707 drop_event_ids_file: Some(file.path().to_path_buf()),
708 drop_matching_events_once: false,
709 ..RelayFaults::default()
710 },
711 ..RelayState::default()
712 };
713
714 assert!(state.should_drop_event("event-to-drop"));
715 assert!(state.should_drop_event("event-to-drop"));
716 }
717}