| 1 | use std::collections::HashMap; |
| 2 | use std::convert::Infallible; |
| 3 | use std::net::{IpAddr, SocketAddr}; |
| 4 | use std::sync::{Arc, Mutex, PoisonError}; |
| 5 | use std::time::Duration; |
| 6 | |
| 7 | use anyhow::Context as _; |
| 8 | use http_body_util::Full; |
| 9 | use hyper::body::{Bytes, Incoming}; |
| 10 | use hyper::header::{ACCEPT, ACCESS_CONTROL_ALLOW_ORIGIN, CONTENT_TYPE, HeaderValue, UPGRADE}; |
| 11 | use hyper::server::conn::http1::Builder as Http1Builder; |
| 12 | use hyper::service::service_fn; |
| 13 | use hyper::upgrade::on as upgrade; |
| 14 | use hyper::{Method, Request, Response, StatusCode}; |
| 15 | use hyper_util::rt::{TokioIo, TokioTimer}; |
| 16 | use nostr_database::NostrDatabase; |
| 17 | use serde_json::json; |
| 18 | use tokio::net::TcpListener; |
| 19 | use tokio::sync::{OwnedSemaphorePermit, Semaphore}; |
| 20 | use tokio::task::{JoinHandle, spawn as spawn_task}; |
| 21 | use tokio::time::timeout; |
| 22 | use tokio_tungstenite::tungstenite::handshake::server::create_response; |
| 23 | use tokio_util::sync::CancellationToken; |
| 24 | use tokio_util::task::TaskTracker; |
| 25 | |
| 26 | use crate::config::Config; |
| 27 | use crate::federation::spawn as spawn_federation; |
| 28 | use crate::policy::Policy; |
| 29 | use crate::relay::Relay; |
| 30 | use crate::store::Store; |
| 31 | |
| 32 | type Body = Full<Bytes>; |
| 33 | |
| 34 | struct ConnectionLease { |
| 35 | ip: IpAddr, |
| 36 | peers: Arc<Mutex<HashMap<IpAddr, usize>>>, |
| 37 | _permit: OwnedSemaphorePermit, |
| 38 | } |
| 39 | |
| 40 | impl Drop for ConnectionLease { |
| 41 | fn drop(&mut self) { |
| 42 | let mut peers = self.peers.lock().unwrap_or_else(PoisonError::into_inner); |
| 43 | if let Some(count) = peers.get_mut(&self.ip) { |
| 44 | *count -= 1; |
| 45 | if *count == 0 { |
| 46 | peers.remove(&self.ip); |
| 47 | } |
| 48 | } |
| 49 | } |
| 50 | } |
| 51 | |
| 52 | #[derive(Clone)] |
| 53 | struct HttpState { |
| 54 | relay: Relay, |
| 55 | store: Arc<Store>, |
| 56 | config: Arc<Config>, |
| 57 | websocket_path: String, |
| 58 | tasks: TaskTracker, |
| 59 | stop: CancellationToken, |
| 60 | } |
| 61 | |
| 62 | pub struct RunningRelay { |
| 63 | address: SocketAddr, |
| 64 | path: String, |
| 65 | state: HttpState, |
| 66 | listener: JoinHandle<anyhow::Result<()>>, |
| 67 | } |
| 68 | |
| 69 | impl RunningRelay { |
| 70 | #[must_use] |
| 71 | pub fn url(&self) -> String { |
| 72 | format!("ws://{}{}", self.address, self.path) |
| 73 | } |
| 74 | |
| 75 | pub async fn shutdown(self) -> anyhow::Result<()> { |
| 76 | self.state.stop.cancel(); |
| 77 | self.listener.await.context("joining listener")??; |
| 78 | self.state.tasks.close(); |
| 79 | timeout(Duration::from_secs(30), self.state.tasks.wait()) |
| 80 | .await |
| 81 | .context("relay connections did not drain")?; |
| 82 | Ok(()) |
| 83 | } |
| 84 | } |
| 85 | |
| 86 | pub async fn spawn(config: Config) -> anyhow::Result<RunningRelay> { |
| 87 | config.validate()?; |
| 88 | let store = Store::open(&config).await?; |
| 89 | let policy_database: Arc<dyn NostrDatabase> = Arc::<Store>::clone(&store); |
| 90 | let policy = Policy::new(policy_database, &config)?; |
| 91 | let listener = TcpListener::bind(config.listen) |
| 92 | .await |
| 93 | .context("binding relay listener")?; |
| 94 | let address = listener.local_addr()?; |
| 95 | let path = config.websocket_path()?; |
| 96 | let stop = CancellationToken::new(); |
| 97 | let relay = Relay::new( |
| 98 | Arc::clone(&store), |
| 99 | policy.clone(), |
| 100 | config.limits.clone(), |
| 101 | stop.clone(), |
| 102 | ); |
| 103 | let state = HttpState { |
| 104 | relay, |
| 105 | store, |
| 106 | config: Arc::new(config), |
| 107 | websocket_path: path.clone(), |
| 108 | tasks: TaskTracker::new(), |
| 109 | stop, |
| 110 | }; |
| 111 | let task_state = state.clone(); |
| 112 | spawn_federation( |
| 113 | Arc::clone(&state.config), |
| 114 | Arc::clone(&state.store), |
| 115 | state.relay.clone(), |
| 116 | policy, |
| 117 | &state.tasks, |
| 118 | state.stop.clone(), |
| 119 | ); |
| 120 | let listener_task = spawn_task(async move { accept_connections(listener, task_state).await }); |
| 121 | tracing::info!(%address, %path, "relay is listening"); |
| 122 | Ok(RunningRelay { |
| 123 | address, |
| 124 | path, |
| 125 | state, |
| 126 | listener: listener_task, |
| 127 | }) |
| 128 | } |
| 129 | |
| 130 | async fn accept_connections(listener: TcpListener, state: HttpState) -> anyhow::Result<()> { |
| 131 | let slots = Arc::new(Semaphore::new(state.config.limits.max_connections)); |
| 132 | let peers = Arc::new(Mutex::new(HashMap::<IpAddr, usize>::new())); |
| 133 | loop { |
| 134 | let (stream, address) = tokio::select! { |
| 135 | biased; |
| 136 | () = state.stop.cancelled() => return Ok(()), |
| 137 | accepted = listener.accept() => accepted?, |
| 138 | }; |
| 139 | let Ok(permit) = Arc::clone(&slots).try_acquire_owned() else { |
| 140 | continue; |
| 141 | }; |
| 142 | let mut peer_counts = peers.lock().unwrap_or_else(PoisonError::into_inner); |
| 143 | let count = peer_counts.entry(address.ip()).or_default(); |
| 144 | if *count >= state.config.limits.max_connections_per_ip { |
| 145 | continue; |
| 146 | } |
| 147 | *count += 1; |
| 148 | drop(peer_counts); |
| 149 | let lease = Arc::new(ConnectionLease { |
| 150 | ip: address.ip(), |
| 151 | peers: Arc::clone(&peers), |
| 152 | _permit: permit, |
| 153 | }); |
| 154 | let handler_state = state.clone(); |
| 155 | let stop = state.stop.clone(); |
| 156 | state.tasks.spawn(async move { |
| 157 | let handler = service_fn(move |request| { |
| 158 | handle_request( |
| 159 | request, |
| 160 | address, |
| 161 | handler_state.clone(), |
| 162 | Arc::clone(&lease), |
| 163 | ) |
| 164 | }); |
| 165 | let mut builder = Http1Builder::new(); |
| 166 | builder.timer(TokioTimer::new()).header_read_timeout(Duration::from_secs(10)); |
| 167 | let connection = builder.serve_connection(TokioIo::new(stream), handler).with_upgrades(); |
| 168 | tokio::pin!(connection); |
| 169 | tokio::select! { |
| 170 | result = &mut connection => { |
| 171 | if let Err(error) = result { tracing::debug!(%address, %error, "http connection ended"); } |
| 172 | } |
| 173 | () = stop.cancelled() => {} |
| 174 | } |
| 175 | }); |
| 176 | } |
| 177 | } |
| 178 | |
| 179 | async fn handle_request( |
| 180 | mut request: Request<Incoming>, |
| 181 | address: SocketAddr, |
| 182 | state: HttpState, |
| 183 | lease: Arc<ConnectionLease>, |
| 184 | ) -> Result<Response<Body>, Infallible> { |
| 185 | if state.stop.is_cancelled() { |
| 186 | return Ok(response( |
| 187 | StatusCode::SERVICE_UNAVAILABLE, |
| 188 | "text/plain", |
| 189 | "relay is stopping", |
| 190 | )); |
| 191 | } |
| 192 | if request.method() != Method::GET { |
| 193 | return Ok(response( |
| 194 | StatusCode::METHOD_NOT_ALLOWED, |
| 195 | "text/plain", |
| 196 | "GET required", |
| 197 | )); |
| 198 | } |
| 199 | if request.uri().path() == "/healthz" { |
| 200 | let healthy = timeout(Duration::from_secs(2), state.store.healthy()) |
| 201 | .await |
| 202 | .unwrap_or(false); |
| 203 | return Ok(response( |
| 204 | if healthy { |
| 205 | StatusCode::OK |
| 206 | } else { |
| 207 | StatusCode::SERVICE_UNAVAILABLE |
| 208 | }, |
| 209 | "text/plain", |
| 210 | if healthy { |
| 211 | "ok" |
| 212 | } else { |
| 213 | "database unavailable" |
| 214 | }, |
| 215 | )); |
| 216 | } |
| 217 | if request.uri().path() != state.websocket_path { |
| 218 | return Ok(response(StatusCode::NOT_FOUND, "text/plain", "not found")); |
| 219 | } |
| 220 | if !request.headers().contains_key(UPGRADE) { |
| 221 | let wants_information = request |
| 222 | .headers() |
| 223 | .get(ACCEPT) |
| 224 | .and_then(|value| value.to_str().ok()) |
| 225 | .is_some_and(|value| { |
| 226 | value.split(',').any(|media_type| { |
| 227 | media_type.trim().split(';').next() == Some("application/nostr+json") |
| 228 | }) |
| 229 | }); |
| 230 | return Ok(if wants_information { |
| 231 | response( |
| 232 | StatusCode::OK, |
| 233 | "application/nostr+json", |
| 234 | &information(&state.config), |
| 235 | ) |
| 236 | } else { |
| 237 | response( |
| 238 | StatusCode::OK, |
| 239 | "text/plain", |
| 240 | "vineyard\na persistent nostr relay for git forges\n", |
| 241 | ) |
| 242 | }); |
| 243 | } |
| 244 | let upgrade = upgrade(&mut request); |
| 245 | let Ok(handshake) = create_response(&request.map(|_body| ())) else { |
| 246 | return Ok(response( |
| 247 | StatusCode::BAD_REQUEST, |
| 248 | "text/plain", |
| 249 | "invalid WebSocket handshake", |
| 250 | )); |
| 251 | }; |
| 252 | let relay = state.relay.clone(); |
| 253 | state.tasks.spawn(async move { |
| 254 | let _lease = lease; |
| 255 | if let Ok(Ok(stream)) = timeout(Duration::from_secs(10), upgrade).await { |
| 256 | if let Err(error) = relay.take_connection(TokioIo::new(stream), address).await { |
| 257 | tracing::debug!(%address, %error, "websocket connection ended"); |
| 258 | } |
| 259 | } else { |
| 260 | tracing::debug!(%address, "websocket upgrade did not complete"); |
| 261 | } |
| 262 | }); |
| 263 | Ok(handshake.map(|()| Full::new(Bytes::new()))) |
| 264 | } |
| 265 | |
| 266 | fn response(status: StatusCode, content_type: &str, body: &str) -> Response<Body> { |
| 267 | let mut response = Response::new(Full::new(Bytes::copy_from_slice(body.as_bytes()))); |
| 268 | *response.status_mut() = status; |
| 269 | if let Ok(value) = HeaderValue::from_str(content_type) { |
| 270 | response.headers_mut().insert(CONTENT_TYPE, value); |
| 271 | } |
| 272 | response |
| 273 | .headers_mut() |
| 274 | .insert(ACCESS_CONTROL_ALLOW_ORIGIN, HeaderValue::from_static("*")); |
| 275 | response |
| 276 | } |
| 277 | |
| 278 | fn information(config: &Config) -> String { |
| 279 | json!({ |
| 280 | "name": "vineyard", |
| 281 | "description": "a persistent nostr relay for git forges", |
| 282 | "software": "vineyard", |
| 283 | "version": env!("CARGO_PKG_VERSION"), |
| 284 | "supported_nips": [1, 9, 11, 34, 65], |
| 285 | "forge_pubkeys": config.forge_pubkeys, |
| 286 | "limitation": { |
| 287 | "max_message_length": config.limits.max_message_bytes, |
| 288 | "max_event_tags": config.limits.max_tags, |
| 289 | "max_subscriptions": config.limits.max_subscriptions, |
| 290 | "max_filters": config.limits.max_filters, |
| 291 | "max_limit": config.limits.max_query_results, |
| 292 | "max_subid_length": 64, |
| 293 | "auth_required": false, |
| 294 | "payment_required": false, |
| 295 | "restricted_writes": true |
| 296 | } |
| 297 | }) |
| 298 | .to_string() |
| 299 | } |