| 1 | use std::collections::{BTreeSet, HashMap}; |
| 2 | use std::net::SocketAddr; |
| 3 | use std::sync::Arc; |
| 4 | use std::time::Duration; |
| 5 | |
| 6 | use anyhow::{Context as _, Result}; |
| 7 | use futures_util::{Sink, SinkExt as _, StreamExt as _}; |
| 8 | use nostr::event::Event; |
| 9 | use nostr::filter::{Filter, MatchEventOptions}; |
| 10 | use nostr_database::{NostrDatabase as _, RejectedReason, SaveEventStatus}; |
| 11 | use serde_json::{Value, json}; |
| 12 | use tokio::io::{AsyncRead, AsyncWrite}; |
| 13 | use tokio::sync::broadcast; |
| 14 | use tokio::time::timeout; |
| 15 | use tokio_tungstenite::WebSocketStream; |
| 16 | use tokio_tungstenite::tungstenite::{ |
| 17 | Error as WebSocketError, Message, |
| 18 | protocol::{Role, WebSocketConfig}, |
| 19 | }; |
| 20 | use tokio_util::sync::CancellationToken; |
| 21 | |
| 22 | use crate::config::Limits; |
| 23 | use crate::policy::{Policy, QueryPolicyResult, WritePolicyResult}; |
| 24 | use crate::store::Store; |
| 25 | |
| 26 | #[derive(Clone)] |
| 27 | pub struct Relay { |
| 28 | store: Arc<Store>, |
| 29 | policy: Policy, |
| 30 | limits: Limits, |
| 31 | events: broadcast::Sender<Event>, |
| 32 | stop: CancellationToken, |
| 33 | } |
| 34 | |
| 35 | struct Subscription { |
| 36 | filters: Vec<Filter>, |
| 37 | bytes: usize, |
| 38 | } |
| 39 | |
| 40 | #[derive(Default)] |
| 41 | struct Session { |
| 42 | subscriptions: HashMap<String, Subscription>, |
| 43 | } |
| 44 | |
| 45 | impl Relay { |
| 46 | pub(crate) fn new( |
| 47 | store: Arc<Store>, |
| 48 | policy: Policy, |
| 49 | limits: Limits, |
| 50 | stop: CancellationToken, |
| 51 | ) -> Self { |
| 52 | let (events, _) = broadcast::channel(limits.max_connections.max(32)); |
| 53 | Self { |
| 54 | store, |
| 55 | policy, |
| 56 | limits, |
| 57 | events, |
| 58 | stop, |
| 59 | } |
| 60 | } |
| 61 | |
| 62 | pub(crate) async fn add_event(&self, event: Event) -> Result<SaveEventStatus> { |
| 63 | event.verify().context("invalid event")?; |
| 64 | let status = self.store.save_event(&event).await?; |
| 65 | if status.is_success() { |
| 66 | match self.events.send(event) { |
| 67 | Ok(_) | Err(_) => {} |
| 68 | } |
| 69 | } |
| 70 | Ok(status) |
| 71 | } |
| 72 | |
| 73 | pub(crate) async fn take_connection<S>(&self, stream: S, address: SocketAddr) -> Result<()> |
| 74 | where |
| 75 | S: AsyncRead + AsyncWrite + Unpin, |
| 76 | { |
| 77 | let websocket = WebSocketStream::from_raw_socket( |
| 78 | stream, |
| 79 | Role::Server, |
| 80 | Some( |
| 81 | WebSocketConfig::default() |
| 82 | .max_message_size(Some(self.limits.max_message_bytes)) |
| 83 | .max_frame_size(Some(self.limits.max_message_bytes)), |
| 84 | ), |
| 85 | ) |
| 86 | .await; |
| 87 | let (mut sink, mut source) = websocket.split(); |
| 88 | let mut events = self.events.subscribe(); |
| 89 | let mut session = Session::default(); |
| 90 | |
| 91 | loop { |
| 92 | tokio::select! { |
| 93 | biased; |
| 94 | () = self.stop.cancelled() => return Ok(()), |
| 95 | received = source.next() => { |
| 96 | let Some(message) = received else { return Ok(()); }; |
| 97 | match message? { |
| 98 | Message::Text(text) => { |
| 99 | for response in self.handle_text(&mut session, address, &text).await { |
| 100 | self.send(&mut sink, response).await?; |
| 101 | } |
| 102 | } |
| 103 | Message::Ping(payload) => { |
| 104 | self.send_message(&mut sink, Message::Pong(payload)).await?; |
| 105 | } |
| 106 | Message::Close(_) => return Ok(()), |
| 107 | Message::Binary(_) | Message::Pong(_) | Message::Frame(_) => {} |
| 108 | } |
| 109 | } |
| 110 | received = events.recv() => { |
| 111 | let event = received.context("live event stream lagged")?; |
| 112 | let ids = session.subscriptions.iter() |
| 113 | .filter(|entry| entry.1.filters.iter().any(|filter| filter.match_event(&event, MatchEventOptions::new()))) |
| 114 | .map(|entry| entry.0.clone()) |
| 115 | .collect::<Vec<_>>(); |
| 116 | for id in ids { |
| 117 | self.send(&mut sink, json!(["EVENT", id, event])).await?; |
| 118 | } |
| 119 | } |
| 120 | } |
| 121 | } |
| 122 | } |
| 123 | |
| 124 | async fn send<S>(&self, sink: &mut S, value: Value) -> Result<()> |
| 125 | where |
| 126 | S: Sink<Message, Error = WebSocketError> + Unpin, |
| 127 | { |
| 128 | self.send_message(sink, Message::Text(value.to_string().into())) |
| 129 | .await |
| 130 | } |
| 131 | |
| 132 | async fn send_message<S>(&self, sink: &mut S, message: Message) -> Result<()> |
| 133 | where |
| 134 | S: Sink<Message, Error = WebSocketError> + Unpin, |
| 135 | { |
| 136 | timeout( |
| 137 | Duration::from_secs(self.limits.write_timeout_seconds), |
| 138 | sink.send(message), |
| 139 | ) |
| 140 | .await |
| 141 | .context("websocket write timed out")??; |
| 142 | Ok(()) |
| 143 | } |
| 144 | |
| 145 | async fn handle_text( |
| 146 | &self, |
| 147 | session: &mut Session, |
| 148 | address: SocketAddr, |
| 149 | text: &str, |
| 150 | ) -> Vec<Value> { |
| 151 | let Ok(value) = serde_json::from_str::<Value>(text) else { |
| 152 | return vec![json!(["NOTICE", "invalid json"])]; |
| 153 | }; |
| 154 | let Some(frame) = value.as_array() else { |
| 155 | return vec![json!(["NOTICE", "message must be an array"])]; |
| 156 | }; |
| 157 | match frame.first().and_then(Value::as_str) { |
| 158 | Some("EVENT") => vec![self.handle_event(frame, address).await], |
| 159 | Some("REQ") => self.handle_req(session, frame, address, text.len()).await, |
| 160 | Some("CLOSE") => { |
| 161 | if let Some(id) = frame.get(1).and_then(Value::as_str) { |
| 162 | session.subscriptions.remove(id); |
| 163 | } |
| 164 | Vec::new() |
| 165 | } |
| 166 | _ => vec![json!(["NOTICE", "unsupported message"])], |
| 167 | } |
| 168 | } |
| 169 | |
| 170 | async fn handle_event(&self, frame: &[Value], address: SocketAddr) -> Value { |
| 171 | let Some(value) = frame.get(1) else { |
| 172 | return json!(["NOTICE", "EVENT requires an event"]); |
| 173 | }; |
| 174 | let id = value |
| 175 | .get("id") |
| 176 | .and_then(Value::as_str) |
| 177 | .unwrap_or_default() |
| 178 | .to_owned(); |
| 179 | let Ok(event) = serde_json::from_value::<Event>(value.clone()) else { |
| 180 | return json!(["OK", id, false, "invalid: malformed event"]); |
| 181 | }; |
| 182 | if event.verify().is_err() { |
| 183 | return json!(["OK", event.id, false, "invalid: bad id or signature"]); |
| 184 | } |
| 185 | match self.policy.admit_event(&event, &address).await { |
| 186 | WritePolicyResult::Accept => {} |
| 187 | WritePolicyResult::Respond { |
| 188 | prefix, |
| 189 | message, |
| 190 | status, |
| 191 | } => return json!(["OK", event.id, status, format!("{prefix}: {message}")]), |
| 192 | } |
| 193 | let event_id = event.id; |
| 194 | match self.add_event(event).await { |
| 195 | Ok(SaveEventStatus::Success) => json!(["OK", event_id, true, ""]), |
| 196 | Ok(SaveEventStatus::Rejected(RejectedReason::Duplicate)) => { |
| 197 | json!(["OK", event_id, true, "duplicate: already stored"]) |
| 198 | } |
| 199 | Ok(SaveEventStatus::Rejected(reason)) => { |
| 200 | json!(["OK", event_id, false, format!("blocked: {reason:?}")]) |
| 201 | } |
| 202 | Err(error) => { |
| 203 | tracing::error!(event = %event_id, %error, "event storage failed"); |
| 204 | json!(["OK", event_id, false, "error: database unavailable"]) |
| 205 | } |
| 206 | } |
| 207 | } |
| 208 | |
| 209 | async fn handle_req( |
| 210 | &self, |
| 211 | session: &mut Session, |
| 212 | frame: &[Value], |
| 213 | address: SocketAddr, |
| 214 | request_bytes: usize, |
| 215 | ) -> Vec<Value> { |
| 216 | let Some(id) = frame.get(1).and_then(Value::as_str) else { |
| 217 | return vec![json!(["NOTICE", "REQ requires a subscription id"])]; |
| 218 | }; |
| 219 | if id.len() > 64 { |
| 220 | return vec![closed(id, "blocked: subscription id is too long")]; |
| 221 | } |
| 222 | if frame.len() < 3 { |
| 223 | return vec![closed(id, "blocked: REQ requires a filter")]; |
| 224 | } |
| 225 | if frame.len() - 2 > self.limits.max_filters { |
| 226 | return vec![closed(id, "rate-limited: too many filters")]; |
| 227 | } |
| 228 | if session.subscriptions.len() >= self.limits.max_subscriptions |
| 229 | && !session.subscriptions.contains_key(id) |
| 230 | { |
| 231 | return vec![closed(id, "rate-limited: too many subscriptions")]; |
| 232 | } |
| 233 | let old_bytes = session.subscriptions.get(id).map_or(0, |entry| entry.bytes); |
| 234 | let active_bytes = session |
| 235 | .subscriptions |
| 236 | .values() |
| 237 | .map(|entry| entry.bytes) |
| 238 | .sum::<usize>() |
| 239 | .saturating_sub(old_bytes); |
| 240 | let max_subscription_bytes = self |
| 241 | .limits |
| 242 | .max_filter_bytes |
| 243 | .saturating_mul(self.limits.max_subscriptions); |
| 244 | if active_bytes.saturating_add(request_bytes) > max_subscription_bytes { |
| 245 | return vec![closed(id, "rate-limited: subscription byte limit exceeded")]; |
| 246 | } |
| 247 | |
| 248 | let mut filters = Vec::with_capacity(frame.len() - 2); |
| 249 | for value in &frame[2..] { |
| 250 | let Ok(filter) = serde_json::from_value::<Filter>(value.clone()) else { |
| 251 | return vec![closed(id, "blocked: invalid filter")]; |
| 252 | }; |
| 253 | match self.policy.admit_query(&filter, &address).await { |
| 254 | QueryPolicyResult::Accept => filters.push(filter), |
| 255 | QueryPolicyResult::Reject { prefix, message } => { |
| 256 | return vec![closed(id, &format!("{prefix}: {message}"))]; |
| 257 | } |
| 258 | } |
| 259 | } |
| 260 | |
| 261 | let mut events = BTreeSet::new(); |
| 262 | for filter in &filters { |
| 263 | let implicit = filter.limit.is_none(); |
| 264 | let limit = filter.limit.unwrap_or(self.limits.max_query_results); |
| 265 | let mut query = filter.clone(); |
| 266 | query.limit = Some(if implicit { |
| 267 | limit.saturating_add(1) |
| 268 | } else { |
| 269 | limit |
| 270 | }); |
| 271 | let found = match self.store.query(query).await { |
| 272 | Ok(found) => found, |
| 273 | Err(error) => { |
| 274 | let message = if error.to_string().contains("history byte limit exceeded") { |
| 275 | "rate-limited: history byte limit exceeded" |
| 276 | } else { |
| 277 | tracing::error!(%error, "history query failed"); |
| 278 | "error: database unavailable" |
| 279 | }; |
| 280 | return vec![closed(id, message)]; |
| 281 | } |
| 282 | }; |
| 283 | if implicit && found.len() > limit { |
| 284 | return vec![closed(id, "rate-limited: history result limit exceeded")]; |
| 285 | } |
| 286 | events.extend(found); |
| 287 | if events.len() > self.limits.max_query_results { |
| 288 | return vec![closed(id, "rate-limited: too many results")]; |
| 289 | } |
| 290 | } |
| 291 | let bytes = events |
| 292 | .iter() |
| 293 | .map(|event| event.as_json().len()) |
| 294 | .sum::<usize>(); |
| 295 | if bytes > self.limits.max_query_bytes { |
| 296 | return vec![closed(id, "rate-limited: history byte limit exceeded")]; |
| 297 | } |
| 298 | |
| 299 | let mut response = Vec::with_capacity(events.len() + 1); |
| 300 | for event in events { |
| 301 | if !event.is_expired() { |
| 302 | response.push(json!(["EVENT", id, event])); |
| 303 | } |
| 304 | } |
| 305 | response.push(json!(["EOSE", id])); |
| 306 | session.subscriptions.insert( |
| 307 | id.to_owned(), |
| 308 | Subscription { |
| 309 | filters, |
| 310 | bytes: request_bytes, |
| 311 | }, |
| 312 | ); |
| 313 | response |
| 314 | } |
| 315 | } |
| 316 | |
| 317 | fn closed(id: &str, message: &str) -> Value { |
| 318 | json!(["CLOSED", id, message]) |
| 319 | } |