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_relay_urls: vec![relay_url],
755 tree_root_cache: Arc::new(std::sync::Mutex::new(std::collections::HashMap::new())),
756 inflight_blob_fetches: Arc::new(tokio::sync::Mutex::new(
757 std::collections::HashMap::new(),
758 )),
759 inflight_blob_reads: Arc::new(
760 tokio::sync::Mutex::new(std::collections::HashMap::new()),
761 ),
762 blob_cache: Arc::new(crate::blob_cache::BlobCache::for_tests()),
763 directory_listing_cache: Arc::new(std::sync::Mutex::new(
764 super::super::auth::new_lookup_cache(),
765 )),
766 resolved_path_cache: Arc::new(std::sync::Mutex::new(
767 super::super::auth::new_lookup_cache(),
768 )),
769 thumbnail_path_cache: Arc::new(std::sync::Mutex::new(
770 super::super::auth::new_lookup_cache(),
771 )),
772 cid_size_cache: Arc::new(std::sync::Mutex::new(super::super::auth::new_lookup_cache())),
773 })
774 }
775
776 #[tokio::test]
777 async fn upstream_proxy_forwards_events_and_caches_them() -> Result<()> {
778 let tmp = TempDir::new()?;
779 let graph_store = {
780 let _guard = crate::socialgraph::test_lock().await;
781 crate::socialgraph::open_test_social_graph_store_with_mapsize(
782 tmp.path(),
783 Some(128 * 1024 * 1024),
784 )?
785 };
786 let backend: Arc<dyn crate::socialgraph::SocialGraphBackend> = graph_store.clone();
787 let access = Arc::new(crate::socialgraph::SocialGraphAccessControl::new(
788 Arc::clone(&backend),
789 0,
790 HashSet::new(),
791 ));
792
793 let keys = Keys::generate();
794 let relay = Arc::new(NostrRelay::new(
795 Arc::clone(&backend),
796 tmp.path().to_path_buf(),
797 HashSet::from([keys.public_key().to_hex()]),
798 Some(access),
799 NostrRelayConfig {
800 spambox_db_max_bytes: 0,
801 ..Default::default()
802 },
803 )?);
804
805 let event = EventBuilder::new(Kind::from(30078_u16), "")
806 .tags([
807 nostr::Tag::parse(["d", "videos/Test"]).expect("d tag"),
808 nostr::Tag::parse(["l", "hashtree"]).expect("label tag"),
809 ])
810 .sign_with_keys(&keys)?;
811
812 let relay_url = spawn_mock_upstream_relay(vec![event.clone()]).await;
813 let filter = Filter::new()
814 .authors(vec![event.pubkey])
815 .kinds(vec![event.kind]);
816 let state = test_app_state(&tmp, relay.clone(), relay_url)?;
817 let client_id = 7_u64;
818 let (tx, mut rx) = mpsc::unbounded_channel();
819 state.ws_relay.clients.lock().await.insert(client_id, tx);
820 let subscription_id = SubscriptionId::new("sub-1");
821
822 start_upstream_nostr_subscription(
823 &state,
824 client_id,
825 subscription_id.clone(),
826 vec![filter.clone()],
827 )
828 .await;
829
830 let forwarded = tokio::time::timeout(std::time::Duration::from_secs(2), rx.recv())
831 .await?
832 .expect("forwarded upstream event");
833 let Message::Text(text) = forwarded else {
834 panic!("expected text event");
835 };
836 match NostrRelayMessage::from_json(text.as_str())? {
837 NostrRelayMessage::Event {
838 subscription_id: sid,
839 event: forwarded_event,
840 } => {
841 assert_eq!(sid.as_ref(), &subscription_id);
842 assert_eq!(forwarded_event.id, event.id);
843 }
844 other => panic!("expected forwarded EVENT, got {:?}", other),
845 }
846
847 let eose = tokio::time::timeout(std::time::Duration::from_secs(2), rx.recv())
848 .await?
849 .expect("forwarded upstream eose");
850 let Message::Text(eose_text) = eose else {
851 panic!("expected text eose");
852 };
853 match NostrRelayMessage::from_json(eose_text.as_str())? {
854 NostrRelayMessage::EndOfStoredEvents(sid) => {
855 assert_eq!(sid.as_ref(), &subscription_id);
856 }
857 other => panic!("expected forwarded EOSE, got {:?}", other),
858 }
859
860 let events = relay.query_events(&filter, 10).await;
861 assert_eq!(events.len(), 1);
862 assert_eq!(events[0].id, event.id);
863
864 close_upstream_nostr_subscription(&state, client_id, &subscription_id).await;
865 assert!(state
866 .ws_relay
867 .upstream_nostr_subscriptions
868 .lock()
869 .await
870 .is_empty());
871 Ok(())
872 }
873
874 #[tokio::test]
875 async fn req_waits_for_upstream_event_before_eose() -> Result<()> {
876 let tmp = TempDir::new()?;
877 let graph_store = {
878 let _guard = crate::socialgraph::test_lock().await;
879 crate::socialgraph::open_test_social_graph_store_with_mapsize(
880 tmp.path(),
881 Some(128 * 1024 * 1024),
882 )?
883 };
884 let backend: Arc<dyn crate::socialgraph::SocialGraphBackend> = graph_store.clone();
885 let access = Arc::new(crate::socialgraph::SocialGraphAccessControl::new(
886 Arc::clone(&backend),
887 0,
888 HashSet::new(),
889 ));
890
891 let keys = Keys::generate();
892 let relay = Arc::new(NostrRelay::new(
893 Arc::clone(&backend),
894 tmp.path().to_path_buf(),
895 HashSet::from([keys.public_key().to_hex()]),
896 Some(access),
897 NostrRelayConfig {
898 spambox_db_max_bytes: 0,
899 ..Default::default()
900 },
901 )?);
902
903 let event = EventBuilder::new(Kind::from(30078_u16), "")
904 .tags([
905 nostr::Tag::parse(["d", "videos/Test"]).expect("d tag"),
906 nostr::Tag::parse(["l", "hashtree"]).expect("label tag"),
907 ])
908 .sign_with_keys(&keys)?;
909
910 let relay_url = spawn_mock_upstream_relay(vec![event.clone()]).await;
911 let state = test_app_state(&tmp, relay.clone(), relay_url)?;
912 let client_id = 11_u64;
913 let (ws_tx, mut ws_rx) = mpsc::unbounded_channel();
914 let (relay_tx, _relay_rx) = mpsc::unbounded_channel();
915 state.ws_relay.clients.lock().await.insert(client_id, ws_tx);
916 relay.register_client(client_id, relay_tx, None).await;
917
918 let request = NostrClientMessage::req(
919 SubscriptionId::new("feed"),
920 vec![Filter::new()
921 .authors(vec![event.pubkey])
922 .kinds(vec![event.kind])],
923 )
924 .as_json();
925
926 handle_message(client_id, Message::Text(request), &state).await;
927
928 let first = tokio::time::timeout(std::time::Duration::from_secs(2), ws_rx.recv())
929 .await?
930 .expect("first forwarded message");
931 let Message::Text(first_text) = first else {
932 panic!("expected text event");
933 };
934 match NostrRelayMessage::from_json(first_text.as_str())? {
935 NostrRelayMessage::Event {
936 event: forwarded_event,
937 ..
938 } => {
939 assert_eq!(forwarded_event.id, event.id);
940 }
941 other => panic!("expected upstream EVENT before EOSE, got {:?}", other),
942 }
943
944 let second = tokio::time::timeout(std::time::Duration::from_secs(2), ws_rx.recv())
945 .await?
946 .expect("second forwarded message");
947 let Message::Text(second_text) = second else {
948 panic!("expected text eose");
949 };
950 match NostrRelayMessage::from_json(second_text.as_str())? {
951 NostrRelayMessage::EndOfStoredEvents(sid) => {
952 assert_eq!(sid.as_ref(), &SubscriptionId::new("feed"));
953 }
954 other => panic!("expected aggregated EOSE, got {:?}", other),
955 }
956
957 Ok(())
958 }
959
960 #[tokio::test]
961 async fn websocket_publish_returns_ok_for_trusted_event() -> Result<()> {
962 let tmp = TempDir::new()?;
963 let graph_store = {
964 let _guard = crate::socialgraph::test_lock().await;
965 crate::socialgraph::open_test_social_graph_store_with_mapsize(
966 tmp.path(),
967 Some(128 * 1024 * 1024),
968 )?
969 };
970 let author_keys = Keys::generate();
971 let backend: Arc<dyn crate::socialgraph::SocialGraphBackend> = graph_store.clone();
972 let access = Arc::new(crate::socialgraph::SocialGraphAccessControl::new(
973 Arc::clone(&backend),
974 0,
975 HashSet::from([author_keys.public_key().to_hex()]),
976 ));
977 let relay = Arc::new(NostrRelay::new(
978 Arc::clone(&backend),
979 tmp.path().to_path_buf(),
980 HashSet::from([author_keys.public_key().to_hex()]),
981 Some(access),
982 NostrRelayConfig {
983 spambox_db_max_bytes: 0,
984 ..Default::default()
985 },
986 )?);
987
988 let state = test_app_state(&tmp, relay.clone(), String::new())?;
989 let listener = TcpListener::bind("127.0.0.1:0").await?;
990 let addr = listener.local_addr()?;
991 let client_pubkey = author_keys.public_key().to_hex();
992 let app = axum::Router::new().route(
993 "/ws",
994 axum::routing::get({
995 let state = state.clone();
996 let client_pubkey = client_pubkey.clone();
997 move |ws: WebSocketUpgrade| {
998 let state = state.clone();
999 let client_pubkey = client_pubkey.clone();
1000 async move { ws_data_with_client_pubkey(state, ws, Some(client_pubkey)) }
1001 }
1002 }),
1003 );
1004 tokio::spawn(async move {
1005 let _ = axum::serve(listener, app).await;
1006 });
1007
1008 let (mut socket, _) = connect_async(format!("ws://{addr}/ws")).await?;
1009 let event = EventBuilder::new(Kind::TextNote, "websocket publish ack")
1010 .sign_with_keys(&author_keys)?;
1011 socket
1012 .send(TungsteniteMessage::Text(
1013 NostrClientMessage::event(event.clone()).as_json(),
1014 ))
1015 .await?;
1016
1017 let reply = tokio::time::timeout(std::time::Duration::from_secs(2), socket.next())
1018 .await?
1019 .ok_or_else(|| anyhow::anyhow!("websocket closed before publish ack"))??;
1020 let TungsteniteMessage::Text(text) = reply else {
1021 anyhow::bail!("expected text publish ack");
1022 };
1023
1024 match NostrRelayMessage::from_json(text.as_str())? {
1025 NostrRelayMessage::Ok {
1026 event_id, status, ..
1027 } => {
1028 assert_eq!(event_id, event.id);
1029 assert!(status);
1030 }
1031 other => anyhow::bail!("expected OK publish ack, got {:?}", other),
1032 }
1033
1034 let stored = relay
1035 .query_events(
1036 &Filter::new()
1037 .authors(vec![event.pubkey])
1038 .kinds(vec![event.kind]),
1039 10,
1040 )
1041 .await;
1042 assert!(stored.iter().any(|candidate| candidate.id == event.id));
1043 Ok(())
1044 }
1045
1046 #[tokio::test]
1047 async fn websocket_req_is_rate_limited_after_configured_quota() -> Result<()> {
1048 let tmp = TempDir::new()?;
1049 let graph_store = {
1050 let _guard = crate::socialgraph::test_lock().await;
1051 crate::socialgraph::open_test_social_graph_store_with_mapsize(
1052 tmp.path(),
1053 Some(128 * 1024 * 1024),
1054 )?
1055 };
1056 let backend: Arc<dyn crate::socialgraph::SocialGraphBackend> = graph_store.clone();
1057 let access = Arc::new(crate::socialgraph::SocialGraphAccessControl::new(
1058 Arc::clone(&backend),
1059 0,
1060 HashSet::new(),
1061 ));
1062 let relay = Arc::new(NostrRelay::new(
1063 Arc::clone(&backend),
1064 tmp.path().to_path_buf(),
1065 HashSet::new(),
1066 Some(access),
1067 NostrRelayConfig {
1068 spambox_db_max_bytes: 0,
1069 spambox_max_reqs_per_min: 1,
1070 ..Default::default()
1071 },
1072 )?);
1073
1074 let state = test_app_state(&tmp, relay, String::new())?;
1075 let listener = TcpListener::bind("127.0.0.1:0").await?;
1076 let addr = listener.local_addr()?;
1077 let app = axum::Router::new().route(
1078 "/ws",
1079 axum::routing::get({
1080 let state = state.clone();
1081 move |ws: WebSocketUpgrade| {
1082 let state = state.clone();
1083 async move { ws_data_with_client_pubkey(state, ws, None) }
1084 }
1085 }),
1086 );
1087 tokio::spawn(async move {
1088 let _ = axum::serve(listener, app).await;
1089 });
1090
1091 let (mut socket, _) = connect_async(format!("ws://{addr}/ws")).await?;
1092 socket
1093 .send(TungsteniteMessage::Text(
1094 NostrClientMessage::req(SubscriptionId::new("sub-1"), vec![Filter::new()])
1095 .as_json(),
1096 ))
1097 .await?;
1098
1099 let first = tokio::time::timeout(std::time::Duration::from_secs(2), socket.next())
1100 .await?
1101 .ok_or_else(|| anyhow::anyhow!("websocket closed before first relay reply"))??;
1102 let TungsteniteMessage::Text(first_text) = first else {
1103 anyhow::bail!("expected text EOSE reply");
1104 };
1105 match NostrRelayMessage::from_json(first_text.as_str())? {
1106 NostrRelayMessage::EndOfStoredEvents(subscription_id) => {
1107 assert_eq!(subscription_id.as_ref(), &SubscriptionId::new("sub-1"));
1108 }
1109 other => anyhow::bail!("expected EOSE for first request, got {:?}", other),
1110 }
1111
1112 socket
1113 .send(TungsteniteMessage::Text(
1114 NostrClientMessage::req(SubscriptionId::new("sub-2"), vec![Filter::new()])
1115 .as_json(),
1116 ))
1117 .await?;
1118
1119 let second = tokio::time::timeout(std::time::Duration::from_secs(2), socket.next())
1120 .await?
1121 .ok_or_else(|| anyhow::anyhow!("websocket closed before rate-limit reply"))??;
1122 let TungsteniteMessage::Text(second_text) = second else {
1123 anyhow::bail!("expected text CLOSED reply");
1124 };
1125 let second_value: serde_json::Value = serde_json::from_str(second_text.as_str())?;
1126 assert_eq!(
1127 second_value,
1128 serde_json::json!(["CLOSED", "sub-2", "rate limited"])
1129 );
1130
1131 Ok(())
1132 }
1133
1134 #[tokio::test]
1135 async fn websocket_publish_is_rate_limited_for_untrusted_spambox_events() -> Result<()> {
1136 let tmp = TempDir::new()?;
1137 let graph_store = {
1138 let _guard = crate::socialgraph::test_lock().await;
1139 crate::socialgraph::open_test_social_graph_store_with_mapsize(
1140 tmp.path(),
1141 Some(128 * 1024 * 1024),
1142 )?
1143 };
1144 let backend: Arc<dyn crate::socialgraph::SocialGraphBackend> = graph_store.clone();
1145 let access = Arc::new(crate::socialgraph::SocialGraphAccessControl::new(
1146 Arc::clone(&backend),
1147 0,
1148 HashSet::new(),
1149 ));
1150 let relay = Arc::new(NostrRelay::new(
1151 Arc::clone(&backend),
1152 tmp.path().to_path_buf(),
1153 HashSet::new(),
1154 Some(access),
1155 NostrRelayConfig {
1156 spambox_db_max_bytes: 0,
1157 spambox_max_events_per_min: 1,
1158 ..Default::default()
1159 },
1160 )?);
1161
1162 let state = test_app_state(&tmp, relay, String::new())?;
1163 let listener = TcpListener::bind("127.0.0.1:0").await?;
1164 let addr = listener.local_addr()?;
1165 let app = axum::Router::new().route(
1166 "/ws",
1167 axum::routing::get({
1168 let state = state.clone();
1169 move |ws: WebSocketUpgrade| {
1170 let state = state.clone();
1171 async move { ws_data_with_client_pubkey(state, ws, None) }
1172 }
1173 }),
1174 );
1175 tokio::spawn(async move {
1176 let _ = axum::serve(listener, app).await;
1177 });
1178
1179 let (mut socket, _) = connect_async(format!("ws://{addr}/ws")).await?;
1180 let author_keys = Keys::generate();
1181 let event_a =
1182 EventBuilder::new(Kind::TextNote, "spambox-a").sign_with_keys(&author_keys)?;
1183 let event_b =
1184 EventBuilder::new(Kind::TextNote, "spambox-b").sign_with_keys(&author_keys)?;
1185
1186 socket
1187 .send(TungsteniteMessage::Text(
1188 NostrClientMessage::event(event_a.clone()).as_json(),
1189 ))
1190 .await?;
1191
1192 let first = tokio::time::timeout(std::time::Duration::from_secs(2), socket.next())
1193 .await?
1194 .ok_or_else(|| anyhow::anyhow!("websocket closed before first publish ack"))??;
1195 let TungsteniteMessage::Text(first_text) = first else {
1196 anyhow::bail!("expected text publish ack");
1197 };
1198 match NostrRelayMessage::from_json(first_text.as_str())? {
1199 NostrRelayMessage::Ok {
1200 event_id,
1201 status,
1202 message,
1203 } => {
1204 assert_eq!(event_id, event_a.id);
1205 assert!(status);
1206 assert_eq!(message, "spambox");
1207 }
1208 other => anyhow::bail!("expected OK publish ack, got {:?}", other),
1209 }
1210
1211 socket
1212 .send(TungsteniteMessage::Text(
1213 NostrClientMessage::event(event_b.clone()).as_json(),
1214 ))
1215 .await?;
1216
1217 let second = tokio::time::timeout(std::time::Duration::from_secs(2), socket.next())
1218 .await?
1219 .ok_or_else(|| anyhow::anyhow!("websocket closed before rate-limit publish ack"))??;
1220 let TungsteniteMessage::Text(second_text) = second else {
1221 anyhow::bail!("expected text publish ack");
1222 };
1223 match NostrRelayMessage::from_json(second_text.as_str())? {
1224 NostrRelayMessage::Ok {
1225 event_id,
1226 status,
1227 message,
1228 } => {
1229 assert_eq!(event_id, event_b.id);
1230 assert!(!status);
1231 assert_eq!(message, "rate limited");
1232 }
1233 other => anyhow::bail!("expected OK=false publish ack, got {:?}", other),
1234 }
1235
1236 Ok(())
1237 }
1238}