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