1use axum::{
2 extract::{
3 ws::{Message, WebSocket, WebSocketUpgrade},
4 State,
5 },
6 response::IntoResponse,
7};
8use futures::{SinkExt, StreamExt};
9use nostr::{
10 ClientMessage as NostrClientMessage, Filter as NostrFilter, JsonUtil as NostrJsonUtil,
11 RelayMessage as NostrRelayMessage, SubscriptionId,
12};
13use std::collections::HashSet;
14use tokio::sync::{mpsc, watch};
15use tokio_tungstenite::{connect_async, tungstenite::Message as TungsteniteMessage};
16
17use super::auth::{AppState, UpstreamNostrSubscription};
18use crate::diagnostics::{
19 nostr_filters_summary, process_memory_snapshot, trim_process_allocations,
20};
21#[derive(Debug)]
22enum WsTextMessage<'a> {
23 Nostr(Box<NostrClientMessage<'a>>),
24}
25
26pub async fn ws_data(State(state): State<AppState>, ws: WebSocketUpgrade) -> impl IntoResponse {
27 ws_data_with_client_pubkey(state, ws, None)
28}
29
30pub fn ws_data_with_client_pubkey(
31 state: AppState,
32 ws: WebSocketUpgrade,
33 client_pubkey: Option<String>,
34) -> impl IntoResponse {
35 ws.on_upgrade(move |socket| handle_socket(socket, state, client_pubkey))
36}
37
38async fn handle_socket(socket: WebSocket, state: AppState, client_pubkey: Option<String>) {
39 let client_id = state
42 .nostr_relay
43 .as_ref()
44 .map(|relay| relay.next_client_id())
45 .unwrap_or_else(|| state.ws_relay.next_id());
46 let (tx, mut rx) = mpsc::unbounded_channel::<Message>();
47
48 {
49 let mut clients = state.ws_relay.clients.lock().await;
50 clients.insert(client_id, tx);
51 }
52 let mut nostr_rx = if let Some(relay) = state.nostr_relay.clone() {
53 let (nostr_tx, nostr_rx) = mpsc::unbounded_channel::<String>();
54 relay
55 .register_client(client_id, nostr_tx, client_pubkey.clone())
56 .await;
57 Some(nostr_rx)
58 } else {
59 None
60 };
61
62 let (mut sender, mut receiver) = socket.split();
63 loop {
64 tokio::select! {
65 maybe_msg = rx.recv() => {
66 let Some(msg) = maybe_msg else {
67 break;
68 };
69 if sender.send(msg).await.is_err() {
70 break;
71 }
72 }
73 maybe_text = async {
74 match &mut nostr_rx {
75 Some(rx) => rx.recv().await,
76 None => std::future::pending().await,
77 }
78 } => {
79 let Some(text) = maybe_text else {
80 nostr_rx = None;
81 continue;
82 };
83 if sender.send(Message::Text(text)).await.is_err() {
84 break;
85 }
86 }
87 maybe_incoming = receiver.next() => {
88 match maybe_incoming {
89 Some(Ok(msg)) => handle_message(client_id, msg, &state).await,
90 Some(Err(_)) | None => break,
91 }
92 }
93 }
94 }
95
96 close_all_upstream_nostr_subscriptions(&state, client_id).await;
97
98 {
99 let mut clients = state.ws_relay.clients.lock().await;
100 clients.remove(&client_id);
101 }
102 if let Some(relay) = &state.nostr_relay {
103 relay.unregister_client(client_id).await;
104 }
105}
106
107fn parse_ws_text_message(text: &str) -> Option<WsTextMessage<'static>> {
108 let trimmed = text.trim_start();
109 if trimmed.starts_with('[') {
110 if let Ok(msg) = NostrClientMessage::from_json(trimmed) {
111 return Some(WsTextMessage::Nostr(Box::new(msg)));
112 }
113 }
114
115 None
116}
117async fn close_upstream_nostr_subscription(
118 state: &AppState,
119 client_id: u64,
120 subscription_id: &SubscriptionId,
121) {
122 let key = (client_id, subscription_id.to_string());
123 let subscription = {
124 let mut subscriptions = state.ws_relay.upstream_nostr_subscriptions.lock().await;
125 subscriptions.remove(&key)
126 };
127 if let Some(subscription) = subscription {
128 let _ = subscription.close_tx.send(true);
129 for task in subscription.tasks {
130 task.abort();
131 }
132 }
133 state
134 .ws_relay
135 .upstream_pending_eose
136 .lock()
137 .await
138 .remove(&key);
139 state
140 .ws_relay
141 .upstream_seen_events
142 .lock()
143 .await
144 .remove(&key);
145}
146
147async fn close_all_upstream_nostr_subscriptions(state: &AppState, client_id: u64) {
148 let keys = {
149 let subscriptions = state.ws_relay.upstream_nostr_subscriptions.lock().await;
150 subscriptions
151 .keys()
152 .filter(|(id, _)| *id == client_id)
153 .cloned()
154 .collect::<Vec<_>>()
155 };
156 for (_, sub_id) in keys {
157 close_upstream_nostr_subscription(state, client_id, &SubscriptionId::new(sub_id)).await;
158 }
159}
160
161async fn forward_upstream_nostr_message(
162 state: &AppState,
163 client_id: u64,
164 subscription_id: &SubscriptionId,
165 text: &str,
166) {
167 let Ok(message) = NostrRelayMessage::from_json(text) else {
168 return;
169 };
170
171 match message {
172 NostrRelayMessage::Event {
173 subscription_id: sid,
174 event,
175 } if sid.as_ref() == subscription_id => {
176 let event = event.into_owned();
177 let key = (client_id, subscription_id.to_string());
178 let event_id = event.id.to_hex();
179 let inserted = {
180 let mut seen_events = state.ws_relay.upstream_seen_events.lock().await;
181 seen_events.entry(key).or_default().insert(event_id)
182 };
183 if !inserted {
184 return;
185 }
186 if let Some(relay) = &state.nostr_relay {
187 let _ = relay.ingest_trusted_event_silent(event.clone()).await;
188 }
189 send_nostr(
190 state,
191 client_id,
192 NostrRelayMessage::event(subscription_id.clone(), event),
193 )
194 .await;
195 }
196 NostrRelayMessage::Closed {
197 subscription_id: sid,
198 message,
199 } if sid.as_ref() == subscription_id => {
200 send_nostr(
201 state,
202 client_id,
203 NostrRelayMessage::closed(subscription_id.clone(), message.into_owned()),
204 )
205 .await;
206 }
207 _ => {}
208 }
209}
210
211async fn mark_upstream_nostr_relay_complete(
212 state: &AppState,
213 client_id: u64,
214 subscription_id: &SubscriptionId,
215) {
216 let key = (client_id, subscription_id.to_string());
217 let should_send_eose = {
218 let mut pending = state.ws_relay.upstream_pending_eose.lock().await;
219 let Some(remaining) = pending.get_mut(&key) else {
220 return;
221 };
222 if *remaining > 0 {
223 *remaining -= 1;
224 }
225 if *remaining == 0 {
226 pending.remove(&key);
227 true
228 } else {
229 false
230 }
231 };
232
233 if should_send_eose {
234 send_nostr(
235 state,
236 client_id,
237 NostrRelayMessage::eose(subscription_id.clone()),
238 )
239 .await;
240 }
241}
242
243async fn run_upstream_nostr_subscription(
244 state: AppState,
245 client_id: u64,
246 relay_url: String,
247 subscription_id: SubscriptionId,
248 filters: Vec<NostrFilter>,
249 mut close_rx: watch::Receiver<bool>,
250) {
251 let mut relay_complete = false;
252 let Ok((socket, _)) = connect_async(relay_url.as_str()).await else {
253 tracing::warn!(
254 "upstream nostr relay connect failed: client_id={} subscription_id={} relay={}",
255 client_id,
256 subscription_id,
257 relay_url,
258 );
259 mark_upstream_nostr_relay_complete(&state, client_id, &subscription_id).await;
260 return;
261 };
262 let (mut write, mut read) = socket.split();
263 let request = NostrClientMessage::req(subscription_id.clone(), filters).as_json();
264 state.ws_relay.note_upstream_relay_send(request.len());
265 if write.send(TungsteniteMessage::Text(request)).await.is_err() {
266 tracing::warn!(
267 "upstream nostr relay request send failed: client_id={} subscription_id={} relay={}",
268 client_id,
269 subscription_id,
270 relay_url,
271 );
272 mark_upstream_nostr_relay_complete(&state, client_id, &subscription_id).await;
273 return;
274 }
275
276 loop {
277 tokio::select! {
278 _ = close_rx.changed() => {
279 if *close_rx.borrow() {
280 let close = NostrClientMessage::close(subscription_id.clone()).as_json();
281 state
282 .ws_relay
283 .note_upstream_relay_send(close.len());
284 let _ = write.send(TungsteniteMessage::Text(close)).await;
285 let _ = write.close().await;
286 break;
287 }
288 }
289 message = read.next() => {
290 match message {
291 Some(Ok(TungsteniteMessage::Text(text))) => {
292 state
293 .ws_relay
294 .note_upstream_relay_receive(text.len());
295 if matches!(
296 NostrRelayMessage::from_json(text.as_str()),
297 Ok(NostrRelayMessage::EndOfStoredEvents(sid)) if sid.as_ref() == &subscription_id
298 ) {
299 if !relay_complete {
300 relay_complete = true;
301 mark_upstream_nostr_relay_complete(&state, client_id, &subscription_id).await;
302 }
303 continue;
304 }
305 forward_upstream_nostr_message(&state, client_id, &subscription_id, &text).await;
306 }
307 Some(Ok(TungsteniteMessage::Ping(payload))) => {
308 let _ = write.send(TungsteniteMessage::Pong(payload)).await;
309 }
310 Some(Ok(TungsteniteMessage::Close(_))) | None => {
311 if !relay_complete {
312 mark_upstream_nostr_relay_complete(&state, client_id, &subscription_id).await;
313 }
314 break;
315 }
316 Some(Err(_)) => {
317 if !relay_complete {
318 mark_upstream_nostr_relay_complete(&state, client_id, &subscription_id).await;
319 }
320 break;
321 }
322 _ => {}
323 }
324 }
325 }
326 }
327}
328
329async fn start_upstream_nostr_subscription(
330 state: &AppState,
331 client_id: u64,
332 subscription_id: SubscriptionId,
333 filters: Vec<NostrFilter>,
334) -> usize {
335 let memory_before = process_memory_snapshot();
336 close_upstream_nostr_subscription(state, client_id, &subscription_id).await;
337 if state.nostr_relay_urls.is_empty() || filters.is_empty() {
338 tracing::info!(
339 "upstream nostr relay skipped: client_id={} subscription_id={} relays={} filters={}",
340 client_id,
341 subscription_id,
342 state.nostr_relay_urls.len(),
343 filters.len(),
344 );
345 return 0;
346 }
347
348 let mut relay_urls = Vec::new();
349 let mut seen = HashSet::new();
350 for relay in &state.nostr_relay_urls {
351 let relay = relay.trim();
352 if relay.is_empty() || !seen.insert(relay.to_string()) {
353 continue;
354 }
355 relay_urls.push(relay.to_string());
356 }
357 if relay_urls.is_empty() {
358 tracing::info!(
359 "upstream nostr relay skipped after normalization: client_id={} subscription_id={}",
360 client_id,
361 subscription_id,
362 );
363 return 0;
364 }
365
366 tracing::info!(
367 "upstream nostr relay start: client_id={} subscription_id={} relays={}",
368 client_id,
369 subscription_id,
370 relay_urls.len(),
371 );
372
373 let filter_summary = nostr_filters_summary(&filters);
374 let key = (client_id, subscription_id.to_string());
375 state
376 .ws_relay
377 .upstream_seen_events
378 .lock()
379 .await
380 .insert(key.clone(), HashSet::new());
381 state
382 .ws_relay
383 .upstream_pending_eose
384 .lock()
385 .await
386 .insert(key.clone(), relay_urls.len());
387
388 let (close_tx, close_rx) = watch::channel(false);
389 let mut tasks = Vec::new();
390 for relay_url in &relay_urls {
391 tasks.push(tokio::spawn(run_upstream_nostr_subscription(
392 state.clone(),
393 client_id,
394 relay_url.clone(),
395 subscription_id.clone(),
396 filters.clone(),
397 close_rx.clone(),
398 )));
399 }
400
401 tracing::info!(
402 target: "hashtree_cli::server::ws_relay::upstream",
403 client_id,
404 subscription_id = %subscription_id,
405 relays = relay_urls.len(),
406 tasks = tasks.len(),
407 filters = filters.len(),
408 filter = %filter_summary,
409 memory_before = ?memory_before,
410 memory_after = ?process_memory_snapshot(),
411 "upstream nostr relay tasks spawned",
412 );
413
414 state
415 .ws_relay
416 .upstream_nostr_subscriptions
417 .lock()
418 .await
419 .insert(key, UpstreamNostrSubscription { close_tx, tasks });
420 relay_urls.len()
421}
422
423async fn handle_message(client_id: u64, msg: Message, state: &AppState) {
424 match msg {
425 Message::Text(text) => {
426 if let Some(msg) = parse_ws_text_message(&text) {
427 match msg {
428 WsTextMessage::Nostr(msg) => {
429 if let Some(relay) = &state.nostr_relay {
430 match *msg {
431 NostrClientMessage::Req {
432 subscription_id,
433 filters,
434 } => {
435 let subscription_id = subscription_id.into_owned();
436 let filters = filters
437 .into_iter()
438 .map(|filter| filter.into_owned())
439 .collect::<Vec<_>>();
440 let local_events = match relay
441 .register_subscription_query(
442 client_id,
443 subscription_id.clone(),
444 filters.clone(),
445 )
446 .await
447 {
448 Ok(events) => events,
449 Err(message) => {
450 send_nostr(
451 state,
452 client_id,
453 NostrRelayMessage::closed(subscription_id, message),
454 )
455 .await;
456 return;
457 }
458 };
459
460 let upstream_relays = start_upstream_nostr_subscription(
461 state,
462 client_id,
463 subscription_id.clone(),
464 filters,
465 )
466 .await;
467 if upstream_relays > 0 {
468 let key = (client_id, subscription_id.to_string());
469 let mut seen_events =
470 state.ws_relay.upstream_seen_events.lock().await;
471 seen_events.entry(key).or_default().extend(
472 local_events.iter().map(|event| event.id.to_hex()),
473 );
474 }
475 for event in local_events {
476 send_nostr(
477 state,
478 client_id,
479 NostrRelayMessage::event(
480 subscription_id.clone(),
481 event,
482 ),
483 )
484 .await;
485 }
486 trim_process_allocations();
487 if upstream_relays == 0 {
488 send_nostr(
489 state,
490 client_id,
491 NostrRelayMessage::eose(subscription_id),
492 )
493 .await;
494 }
495 }
496 NostrClientMessage::Close(subscription_id) => {
497 let subscription_id = subscription_id.into_owned();
498 close_upstream_nostr_subscription(
499 state,
500 client_id,
501 &subscription_id,
502 )
503 .await;
504 relay
505 .handle_client_message(
506 client_id,
507 NostrClientMessage::close(subscription_id.clone()),
508 )
509 .await;
510 }
511 other => {
512 relay.handle_client_message(client_id, other).await;
513 }
514 }
515 } else {
516 handle_nostr_message(client_id, *msg, state).await;
517 }
518 }
519 }
520 }
521 }
522 Message::Binary(_) => {}
523 Message::Close(_) => {}
524 _ => {}
525 }
526}
527
528async fn handle_nostr_message(client_id: u64, msg: NostrClientMessage<'_>, state: &AppState) {
529 let replies = nostr_responses_for(&msg);
530 for reply in replies {
531 send_nostr(state, client_id, reply).await;
532 }
533}
534
535fn nostr_responses_for(msg: &NostrClientMessage<'_>) -> Vec<NostrRelayMessage<'static>> {
536 match msg {
537 NostrClientMessage::Event(event) => {
538 let ok = event.verify().is_ok();
539 let message = if ok { "" } else { "invalid: signature" };
540 vec![NostrRelayMessage::ok(event.id, ok, message)]
541 }
542 NostrClientMessage::Req {
543 subscription_id, ..
544 } => {
545 vec![NostrRelayMessage::eose(
546 subscription_id.clone().into_owned(),
547 )]
548 }
549 NostrClientMessage::Count {
550 subscription_id, ..
551 } => {
552 vec![NostrRelayMessage::count(
553 subscription_id.clone().into_owned(),
554 0,
555 )]
556 }
557 NostrClientMessage::Close(_) => Vec::new(),
558 NostrClientMessage::Auth(event) => {
559 let ok = event.verify().is_ok();
560 let message = if ok { "" } else { "invalid auth" };
561 vec![NostrRelayMessage::ok(event.id, ok, message)]
562 }
563 NostrClientMessage::NegOpen { .. }
564 | NostrClientMessage::NegMsg { .. }
565 | NostrClientMessage::NegClose { .. } => {
566 vec![NostrRelayMessage::notice("negentropy not supported")]
567 }
568 }
569}
570
571async fn send_nostr(state: &AppState, client_id: u64, response: NostrRelayMessage<'_>) {
572 let text = response.as_json();
573 send_to_client(state, client_id, Message::Text(text)).await;
574}
575
576async fn send_to_client(state: &AppState, client_id: u64, msg: Message) {
577 let sender = {
578 let clients = state.ws_relay.clients.lock().await;
579 clients.get(&client_id).cloned()
580 };
581 if let Some(tx) = sender {
582 let _ = tx.send(msg);
583 }
584}
585
586#[cfg(test)]
587mod tests {
588 use super::*;
589 use crate::nostr_relay::{NostrRelay, NostrRelayConfig};
590 use anyhow::Result;
591 use futures::{SinkExt, StreamExt};
592 use nostr::secp256k1::schnorr::Signature;
593 use nostr::{EventBuilder, Filter, Keys, Kind, SubscriptionId};
594 use std::collections::HashSet;
595 use std::sync::Arc;
596 use tempfile::TempDir;
597 use tokio::net::TcpListener;
598 use tokio_tungstenite::{accept_async, tungstenite::Message as TungsteniteMessage};
599
600 #[test]
601 fn parse_ws_text_message_detects_nostr_req() {
602 let msg = r#"["REQ","sub-1",{"kinds":[1]}]"#;
603 match parse_ws_text_message(msg) {
604 Some(WsTextMessage::Nostr(_)) => {}
605 other => panic!("expected Nostr message, got {:?}", other),
606 }
607 }
608
609 #[test]
610 fn nostr_replies_for_req_is_eose() {
611 let sub = SubscriptionId::new("sub-1");
612 let msg = NostrClientMessage::req(sub.clone(), vec![]);
613 let replies = nostr_responses_for(&msg);
614 assert_eq!(replies.len(), 1);
615 match &replies[0] {
616 NostrRelayMessage::EndOfStoredEvents(id) => assert_eq!(id.as_ref(), &sub),
617 other => panic!("expected EOSE, got {:?}", other),
618 }
619 }
620
621 #[test]
622 fn nostr_replies_for_event_ok() {
623 let keys = Keys::generate();
624 let event = EventBuilder::new(Kind::TextNote, "hello")
625 .sign_with_keys(&keys)
626 .unwrap();
627 let msg = NostrClientMessage::event(event.clone());
628 let replies = nostr_responses_for(&msg);
629 assert_eq!(replies.len(), 1);
630 match &replies[0] {
631 NostrRelayMessage::Ok {
632 event_id, status, ..
633 } => {
634 assert_eq!(event_id, &event.id);
635 assert!(*status);
636 }
637 other => panic!("expected OK, got {:?}", other),
638 }
639 }
640
641 #[test]
642 fn nostr_replies_for_invalid_event_is_not_ok() {
643 let keys = Keys::generate();
644 let mut event = EventBuilder::new(Kind::TextNote, "hello")
645 .sign_with_keys(&keys)
646 .unwrap();
647 event.sig = Signature::from_slice(&[0u8; 64]).unwrap();
648 let msg = NostrClientMessage::event(event);
649 let replies = nostr_responses_for(&msg);
650 assert_eq!(replies.len(), 1);
651 match &replies[0] {
652 NostrRelayMessage::Ok { status, .. } => assert!(!*status),
653 other => panic!("expected OK=false, got {:?}", other),
654 }
655 }
656
657 async fn spawn_mock_upstream_relay(events: Vec<nostr::Event>) -> String {
658 let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind relay");
659 let addr = listener.local_addr().expect("relay addr");
660 tokio::spawn(async move {
661 let (stream, _) = listener.accept().await.expect("accept relay");
662 let ws = accept_async(stream).await.expect("accept websocket");
663 let (mut write, mut read) = ws.split();
664
665 while let Some(Ok(message)) = read.next().await {
666 let TungsteniteMessage::Text(text) = message else {
667 continue;
668 };
669 let Ok(parsed) = NostrClientMessage::from_json(text.as_bytes()) else {
670 continue;
671 };
672 if let NostrClientMessage::Req {
673 subscription_id,
674 filters,
675 } = parsed
676 {
677 let subscription_id = subscription_id.into_owned();
678 let filters = filters
679 .into_iter()
680 .map(|filter| filter.into_owned())
681 .collect::<Vec<_>>();
682 for event in events.iter().filter(|event| {
683 filters
684 .iter()
685 .any(|filter| filter.match_event(event, Default::default()))
686 }) {
687 let _ = write
688 .send(TungsteniteMessage::Text(
689 NostrRelayMessage::event(subscription_id.clone(), event.clone())
690 .as_json(),
691 ))
692 .await;
693 }
694 let _ = write
695 .send(TungsteniteMessage::Text(
696 NostrRelayMessage::eose(subscription_id).as_json(),
697 ))
698 .await;
699 }
700 }
701 });
702 format!("ws://{}", addr)
703 }
704
705 fn test_app_state(
706 tmp: &TempDir,
707 relay: Arc<NostrRelay>,
708 relay_url: String,
709 ) -> Result<AppState> {
710 let store = Arc::new(crate::storage::HashtreeStore::with_options(
711 tmp.path(),
712 None,
713 128 * 1024 * 1024,
714 )?);
715 Ok(AppState {
716 store,
717 auth: None,
718 daemon_started_at: 1_700_000_000,
719 peer_mode: crate::config::ServerMode::Normal,
720 hash_get_enabled: true,
721 fips_endpoint: None,
722 fips_blob_resolver: None,
723 fetch_from_fips_peers: true,
724 ws_relay: Arc::new(super::super::auth::WsRelayState::new()),
725 max_upload_bytes: 5 * 1024 * 1024,
726 public_writes: true,
727 public_plaintext_reads: true,
728 require_random_untrusted_ingest: false,
729 optimistic_blossom_uploads: false,
730 optimistic_upload_queue_bytes: 256 * 1024 * 1024,
731 optimistic_upload_queue: Arc::new(tokio::sync::Semaphore::new(256 * 1024 * 1024)),
732 allowed_pubkeys: HashSet::new(),
733 upstream_blossom: Vec::new(),
734 upstream_http_client: super::super::new_upstream_http_client(),
735 upstream_blossom_miss_cache: Arc::new(std::sync::Mutex::new(
736 super::super::auth::new_lookup_cache(),
737 )),
738 upstream_blossom_fetch_metrics: Arc::new(
739 super::super::auth::UpstreamBlossomFetchMetrics::default(),
740 ),
741 blossom_upload_replicas: Vec::new(),
742 blossom_upload_replica_queue_bytes: 256 * 1024 * 1024,
743 blossom_upload_replica_queue: Arc::new(tokio::sync::Semaphore::new(256 * 1024 * 1024)),
744 blossom_upload_replica_keys: None,
745 blossom_upload_replica_scheduler: Arc::new(
746 super::super::blossom::BlossomUploadReplicaScheduler::new(),
747 ),
748 social_graph: None,
749 social_graph_store: None,
750 social_graph_root: None,
751 socialgraph_snapshot_public: false,
752 nostr_relay: Some(relay),
753 nostr_provider: None,
754 nostr_event_transport: None,
755 nostr_relay_urls: vec![relay_url],
756 tree_root_cache: Arc::new(std::sync::Mutex::new(std::collections::HashMap::new())),
757 inflight_blob_fetches: Arc::new(tokio::sync::Mutex::new(
758 std::collections::HashMap::new(),
759 )),
760 inflight_blob_reads: Arc::new(
761 tokio::sync::Mutex::new(std::collections::HashMap::new()),
762 ),
763 blob_cache: Arc::new(crate::blob_cache::BlobCache::for_tests()),
764 directory_listing_cache: Arc::new(std::sync::Mutex::new(
765 super::super::auth::new_lookup_cache(),
766 )),
767 resolved_path_cache: Arc::new(std::sync::Mutex::new(
768 super::super::auth::new_lookup_cache(),
769 )),
770 thumbnail_path_cache: Arc::new(std::sync::Mutex::new(
771 super::super::auth::new_lookup_cache(),
772 )),
773 cid_size_cache: Arc::new(std::sync::Mutex::new(super::super::auth::new_lookup_cache())),
774 })
775 }
776
777 #[tokio::test]
778 async fn upstream_proxy_forwards_events_and_caches_them() -> Result<()> {
779 let tmp = TempDir::new()?;
780 let graph_store = {
781 let _guard = crate::socialgraph::test_lock().await;
782 crate::socialgraph::open_test_social_graph_store_with_mapsize(
783 tmp.path(),
784 Some(128 * 1024 * 1024),
785 )?
786 };
787 let backend: Arc<dyn crate::socialgraph::SocialGraphBackend> = graph_store.clone();
788 let access = Arc::new(crate::socialgraph::SocialGraphAccessControl::new(
789 Arc::clone(&backend),
790 0,
791 HashSet::new(),
792 ));
793
794 let keys = Keys::generate();
795 let relay = Arc::new(NostrRelay::new(
796 Arc::clone(&backend),
797 tmp.path().to_path_buf(),
798 HashSet::from([keys.public_key().to_hex()]),
799 Some(access),
800 NostrRelayConfig {
801 spambox_db_max_bytes: 0,
802 ..Default::default()
803 },
804 )?);
805
806 let event = EventBuilder::new(Kind::from(30078_u16), "")
807 .tags([
808 nostr::Tag::parse(["d", "videos/Test"]).expect("d tag"),
809 nostr::Tag::parse(["l", "hashtree"]).expect("label tag"),
810 ])
811 .sign_with_keys(&keys)?;
812
813 let relay_url = spawn_mock_upstream_relay(vec![event.clone()]).await;
814 let filter = Filter::new()
815 .authors(vec![event.pubkey])
816 .kinds(vec![event.kind]);
817 let state = test_app_state(&tmp, relay.clone(), relay_url)?;
818 let client_id = 7_u64;
819 let (tx, mut rx) = mpsc::unbounded_channel();
820 state.ws_relay.clients.lock().await.insert(client_id, tx);
821 let subscription_id = SubscriptionId::new("sub-1");
822
823 start_upstream_nostr_subscription(
824 &state,
825 client_id,
826 subscription_id.clone(),
827 vec![filter.clone()],
828 )
829 .await;
830
831 let forwarded = tokio::time::timeout(std::time::Duration::from_secs(2), rx.recv())
832 .await?
833 .expect("forwarded upstream event");
834 let Message::Text(text) = forwarded else {
835 panic!("expected text event");
836 };
837 match NostrRelayMessage::from_json(text.as_str())? {
838 NostrRelayMessage::Event {
839 subscription_id: sid,
840 event: forwarded_event,
841 } => {
842 assert_eq!(sid.as_ref(), &subscription_id);
843 assert_eq!(forwarded_event.id, event.id);
844 }
845 other => panic!("expected forwarded EVENT, got {:?}", other),
846 }
847
848 let eose = tokio::time::timeout(std::time::Duration::from_secs(2), rx.recv())
849 .await?
850 .expect("forwarded upstream eose");
851 let Message::Text(eose_text) = eose else {
852 panic!("expected text eose");
853 };
854 match NostrRelayMessage::from_json(eose_text.as_str())? {
855 NostrRelayMessage::EndOfStoredEvents(sid) => {
856 assert_eq!(sid.as_ref(), &subscription_id);
857 }
858 other => panic!("expected forwarded EOSE, got {:?}", other),
859 }
860
861 let events = relay.query_events(&filter, 10).await;
862 assert_eq!(events.len(), 1);
863 assert_eq!(events[0].id, event.id);
864
865 close_upstream_nostr_subscription(&state, client_id, &subscription_id).await;
866 assert!(state
867 .ws_relay
868 .upstream_nostr_subscriptions
869 .lock()
870 .await
871 .is_empty());
872 Ok(())
873 }
874
875 #[tokio::test]
876 async fn req_waits_for_upstream_event_before_eose() -> Result<()> {
877 let tmp = TempDir::new()?;
878 let graph_store = {
879 let _guard = crate::socialgraph::test_lock().await;
880 crate::socialgraph::open_test_social_graph_store_with_mapsize(
881 tmp.path(),
882 Some(128 * 1024 * 1024),
883 )?
884 };
885 let backend: Arc<dyn crate::socialgraph::SocialGraphBackend> = graph_store.clone();
886 let access = Arc::new(crate::socialgraph::SocialGraphAccessControl::new(
887 Arc::clone(&backend),
888 0,
889 HashSet::new(),
890 ));
891
892 let keys = Keys::generate();
893 let relay = Arc::new(NostrRelay::new(
894 Arc::clone(&backend),
895 tmp.path().to_path_buf(),
896 HashSet::from([keys.public_key().to_hex()]),
897 Some(access),
898 NostrRelayConfig {
899 spambox_db_max_bytes: 0,
900 ..Default::default()
901 },
902 )?);
903
904 let event = EventBuilder::new(Kind::from(30078_u16), "")
905 .tags([
906 nostr::Tag::parse(["d", "videos/Test"]).expect("d tag"),
907 nostr::Tag::parse(["l", "hashtree"]).expect("label tag"),
908 ])
909 .sign_with_keys(&keys)?;
910
911 let relay_url = spawn_mock_upstream_relay(vec![event.clone()]).await;
912 let state = test_app_state(&tmp, relay.clone(), relay_url)?;
913 let client_id = 11_u64;
914 let (ws_tx, mut ws_rx) = mpsc::unbounded_channel();
915 let (relay_tx, _relay_rx) = mpsc::unbounded_channel();
916 state.ws_relay.clients.lock().await.insert(client_id, ws_tx);
917 relay.register_client(client_id, relay_tx, None).await;
918
919 let request = NostrClientMessage::req(
920 SubscriptionId::new("feed"),
921 vec![Filter::new()
922 .authors(vec![event.pubkey])
923 .kinds(vec![event.kind])],
924 )
925 .as_json();
926
927 handle_message(client_id, Message::Text(request), &state).await;
928
929 let first = tokio::time::timeout(std::time::Duration::from_secs(2), ws_rx.recv())
930 .await?
931 .expect("first forwarded message");
932 let Message::Text(first_text) = first else {
933 panic!("expected text event");
934 };
935 match NostrRelayMessage::from_json(first_text.as_str())? {
936 NostrRelayMessage::Event {
937 event: forwarded_event,
938 ..
939 } => {
940 assert_eq!(forwarded_event.id, event.id);
941 }
942 other => panic!("expected upstream EVENT before EOSE, got {:?}", other),
943 }
944
945 let second = tokio::time::timeout(std::time::Duration::from_secs(2), ws_rx.recv())
946 .await?
947 .expect("second forwarded message");
948 let Message::Text(second_text) = second else {
949 panic!("expected text eose");
950 };
951 match NostrRelayMessage::from_json(second_text.as_str())? {
952 NostrRelayMessage::EndOfStoredEvents(sid) => {
953 assert_eq!(sid.as_ref(), &SubscriptionId::new("feed"));
954 }
955 other => panic!("expected aggregated EOSE, got {:?}", other),
956 }
957
958 Ok(())
959 }
960
961 #[tokio::test]
962 async fn websocket_publish_returns_ok_for_trusted_event() -> Result<()> {
963 let tmp = TempDir::new()?;
964 let graph_store = {
965 let _guard = crate::socialgraph::test_lock().await;
966 crate::socialgraph::open_test_social_graph_store_with_mapsize(
967 tmp.path(),
968 Some(128 * 1024 * 1024),
969 )?
970 };
971 let author_keys = Keys::generate();
972 let backend: Arc<dyn crate::socialgraph::SocialGraphBackend> = graph_store.clone();
973 let access = Arc::new(crate::socialgraph::SocialGraphAccessControl::new(
974 Arc::clone(&backend),
975 0,
976 HashSet::from([author_keys.public_key().to_hex()]),
977 ));
978 let relay = Arc::new(NostrRelay::new(
979 Arc::clone(&backend),
980 tmp.path().to_path_buf(),
981 HashSet::from([author_keys.public_key().to_hex()]),
982 Some(access),
983 NostrRelayConfig {
984 spambox_db_max_bytes: 0,
985 ..Default::default()
986 },
987 )?);
988
989 let state = test_app_state(&tmp, relay.clone(), String::new())?;
990 let listener = TcpListener::bind("127.0.0.1:0").await?;
991 let addr = listener.local_addr()?;
992 let client_pubkey = author_keys.public_key().to_hex();
993 let app = axum::Router::new().route(
994 "/ws",
995 axum::routing::get({
996 let state = state.clone();
997 let client_pubkey = client_pubkey.clone();
998 move |ws: WebSocketUpgrade| {
999 let state = state.clone();
1000 let client_pubkey = client_pubkey.clone();
1001 async move { ws_data_with_client_pubkey(state, ws, Some(client_pubkey)) }
1002 }
1003 }),
1004 );
1005 tokio::spawn(async move {
1006 let _ = axum::serve(listener, app).await;
1007 });
1008
1009 let (mut socket, _) = connect_async(format!("ws://{addr}/ws")).await?;
1010 let event = EventBuilder::new(Kind::TextNote, "websocket publish ack")
1011 .sign_with_keys(&author_keys)?;
1012 socket
1013 .send(TungsteniteMessage::Text(
1014 NostrClientMessage::event(event.clone()).as_json(),
1015 ))
1016 .await?;
1017
1018 let reply = tokio::time::timeout(std::time::Duration::from_secs(2), socket.next())
1019 .await?
1020 .ok_or_else(|| anyhow::anyhow!("websocket closed before publish ack"))??;
1021 let TungsteniteMessage::Text(text) = reply else {
1022 anyhow::bail!("expected text publish ack");
1023 };
1024
1025 match NostrRelayMessage::from_json(text.as_str())? {
1026 NostrRelayMessage::Ok {
1027 event_id, status, ..
1028 } => {
1029 assert_eq!(event_id, event.id);
1030 assert!(status);
1031 }
1032 other => anyhow::bail!("expected OK publish ack, got {:?}", other),
1033 }
1034
1035 let stored = relay
1036 .query_events(
1037 &Filter::new()
1038 .authors(vec![event.pubkey])
1039 .kinds(vec![event.kind]),
1040 10,
1041 )
1042 .await;
1043 assert!(stored.iter().any(|candidate| candidate.id == event.id));
1044 Ok(())
1045 }
1046
1047 #[tokio::test]
1048 async fn websocket_req_is_rate_limited_after_configured_quota() -> Result<()> {
1049 let tmp = TempDir::new()?;
1050 let graph_store = {
1051 let _guard = crate::socialgraph::test_lock().await;
1052 crate::socialgraph::open_test_social_graph_store_with_mapsize(
1053 tmp.path(),
1054 Some(128 * 1024 * 1024),
1055 )?
1056 };
1057 let backend: Arc<dyn crate::socialgraph::SocialGraphBackend> = graph_store.clone();
1058 let access = Arc::new(crate::socialgraph::SocialGraphAccessControl::new(
1059 Arc::clone(&backend),
1060 0,
1061 HashSet::new(),
1062 ));
1063 let relay = Arc::new(NostrRelay::new(
1064 Arc::clone(&backend),
1065 tmp.path().to_path_buf(),
1066 HashSet::new(),
1067 Some(access),
1068 NostrRelayConfig {
1069 spambox_db_max_bytes: 0,
1070 spambox_max_reqs_per_min: 1,
1071 ..Default::default()
1072 },
1073 )?);
1074
1075 let state = test_app_state(&tmp, relay, String::new())?;
1076 let listener = TcpListener::bind("127.0.0.1:0").await?;
1077 let addr = listener.local_addr()?;
1078 let app = axum::Router::new().route(
1079 "/ws",
1080 axum::routing::get({
1081 let state = state.clone();
1082 move |ws: WebSocketUpgrade| {
1083 let state = state.clone();
1084 async move { ws_data_with_client_pubkey(state, ws, None) }
1085 }
1086 }),
1087 );
1088 tokio::spawn(async move {
1089 let _ = axum::serve(listener, app).await;
1090 });
1091
1092 let (mut socket, _) = connect_async(format!("ws://{addr}/ws")).await?;
1093 socket
1094 .send(TungsteniteMessage::Text(
1095 NostrClientMessage::req(SubscriptionId::new("sub-1"), vec![Filter::new()])
1096 .as_json(),
1097 ))
1098 .await?;
1099
1100 let first = tokio::time::timeout(std::time::Duration::from_secs(2), socket.next())
1101 .await?
1102 .ok_or_else(|| anyhow::anyhow!("websocket closed before first relay reply"))??;
1103 let TungsteniteMessage::Text(first_text) = first else {
1104 anyhow::bail!("expected text EOSE reply");
1105 };
1106 match NostrRelayMessage::from_json(first_text.as_str())? {
1107 NostrRelayMessage::EndOfStoredEvents(subscription_id) => {
1108 assert_eq!(subscription_id.as_ref(), &SubscriptionId::new("sub-1"));
1109 }
1110 other => anyhow::bail!("expected EOSE for first request, got {:?}", other),
1111 }
1112
1113 socket
1114 .send(TungsteniteMessage::Text(
1115 NostrClientMessage::req(SubscriptionId::new("sub-2"), vec![Filter::new()])
1116 .as_json(),
1117 ))
1118 .await?;
1119
1120 let second = tokio::time::timeout(std::time::Duration::from_secs(2), socket.next())
1121 .await?
1122 .ok_or_else(|| anyhow::anyhow!("websocket closed before rate-limit reply"))??;
1123 let TungsteniteMessage::Text(second_text) = second else {
1124 anyhow::bail!("expected text CLOSED reply");
1125 };
1126 let second_value: serde_json::Value = serde_json::from_str(second_text.as_str())?;
1127 assert_eq!(
1128 second_value,
1129 serde_json::json!(["CLOSED", "sub-2", "rate limited"])
1130 );
1131
1132 Ok(())
1133 }
1134
1135 #[tokio::test]
1136 async fn websocket_publish_is_rate_limited_for_untrusted_spambox_events() -> Result<()> {
1137 let tmp = TempDir::new()?;
1138 let graph_store = {
1139 let _guard = crate::socialgraph::test_lock().await;
1140 crate::socialgraph::open_test_social_graph_store_with_mapsize(
1141 tmp.path(),
1142 Some(128 * 1024 * 1024),
1143 )?
1144 };
1145 let backend: Arc<dyn crate::socialgraph::SocialGraphBackend> = graph_store.clone();
1146 let access = Arc::new(crate::socialgraph::SocialGraphAccessControl::new(
1147 Arc::clone(&backend),
1148 0,
1149 HashSet::new(),
1150 ));
1151 let relay = Arc::new(NostrRelay::new(
1152 Arc::clone(&backend),
1153 tmp.path().to_path_buf(),
1154 HashSet::new(),
1155 Some(access),
1156 NostrRelayConfig {
1157 spambox_db_max_bytes: 0,
1158 spambox_max_events_per_min: 1,
1159 ..Default::default()
1160 },
1161 )?);
1162
1163 let state = test_app_state(&tmp, relay, String::new())?;
1164 let listener = TcpListener::bind("127.0.0.1:0").await?;
1165 let addr = listener.local_addr()?;
1166 let app = axum::Router::new().route(
1167 "/ws",
1168 axum::routing::get({
1169 let state = state.clone();
1170 move |ws: WebSocketUpgrade| {
1171 let state = state.clone();
1172 async move { ws_data_with_client_pubkey(state, ws, None) }
1173 }
1174 }),
1175 );
1176 tokio::spawn(async move {
1177 let _ = axum::serve(listener, app).await;
1178 });
1179
1180 let (mut socket, _) = connect_async(format!("ws://{addr}/ws")).await?;
1181 let author_keys = Keys::generate();
1182 let event_a =
1183 EventBuilder::new(Kind::TextNote, "spambox-a").sign_with_keys(&author_keys)?;
1184 let event_b =
1185 EventBuilder::new(Kind::TextNote, "spambox-b").sign_with_keys(&author_keys)?;
1186
1187 socket
1188 .send(TungsteniteMessage::Text(
1189 NostrClientMessage::event(event_a.clone()).as_json(),
1190 ))
1191 .await?;
1192
1193 let first = tokio::time::timeout(std::time::Duration::from_secs(2), socket.next())
1194 .await?
1195 .ok_or_else(|| anyhow::anyhow!("websocket closed before first publish ack"))??;
1196 let TungsteniteMessage::Text(first_text) = first else {
1197 anyhow::bail!("expected text publish ack");
1198 };
1199 match NostrRelayMessage::from_json(first_text.as_str())? {
1200 NostrRelayMessage::Ok {
1201 event_id,
1202 status,
1203 message,
1204 } => {
1205 assert_eq!(event_id, event_a.id);
1206 assert!(status);
1207 assert_eq!(message, "spambox");
1208 }
1209 other => anyhow::bail!("expected OK publish ack, got {:?}", other),
1210 }
1211
1212 socket
1213 .send(TungsteniteMessage::Text(
1214 NostrClientMessage::event(event_b.clone()).as_json(),
1215 ))
1216 .await?;
1217
1218 let second = tokio::time::timeout(std::time::Duration::from_secs(2), socket.next())
1219 .await?
1220 .ok_or_else(|| anyhow::anyhow!("websocket closed before rate-limit publish ack"))??;
1221 let TungsteniteMessage::Text(second_text) = second else {
1222 anyhow::bail!("expected text publish ack");
1223 };
1224 match NostrRelayMessage::from_json(second_text.as_str())? {
1225 NostrRelayMessage::Ok {
1226 event_id,
1227 status,
1228 message,
1229 } => {
1230 assert_eq!(event_id, event_b.id);
1231 assert!(!status);
1232 assert_eq!(message, "rate limited");
1233 }
1234 other => anyhow::bail!("expected OK=false publish ack, got {:?}", other),
1235 }
1236
1237 Ok(())
1238 }
1239}