1use std::{
4 borrow::Cow,
5 convert::Infallible,
6 io,
7 net::{Ipv4Addr, Ipv6Addr, SocketAddr},
8 sync::Arc,
9 time::Duration,
10};
11
12use bytes::Bytes;
13use http_body_util::Full;
14use hyper::{
15 Request, Response, StatusCode,
16 body::Incoming,
17 header::{CONTENT_TYPE, HeaderValue},
18 server::conn::http1,
19 service::service_fn,
20};
21use hyper_util::rt::TokioIo;
22use socket2::{Domain, Protocol, Socket, Type};
23use tokio::{
24 net::TcpListener,
25 sync::watch,
26 task::{JoinHandle, JoinSet},
27};
28
29const MAX_CONNECTIONS: usize = 1024;
30const TIMEOUT: Duration = Duration::from_secs(5);
31
32pub struct Listener {
35 address: SocketAddr,
36 stop: watch::Sender<bool>,
37 task: Option<JoinHandle<io::Result<()>>>,
38}
39
40pub struct HealthProbe(JoinHandle<()>);
42
43impl HealthProbe {
44 pub fn start(address: &str, registry: crate::metrics::Registry) -> Self {
45 let address = if address.starts_with(':') {
46 format!("localhost{address}")
47 } else {
48 address.to_owned()
49 };
50 Self(tokio::spawn(async move {
51 let mut interval = tokio::time::interval(Duration::from_secs(1));
52 interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
53 interval.tick().await;
54 loop {
55 interval.tick().await;
56 let start = std::time::Instant::now();
57 let result =
58 tokio::time::timeout(Duration::from_secs(3), health_request(&address)).await;
59 let failed = !matches!(result, Ok(Ok(())));
60 let elapsed = start.elapsed();
61 registry.observe_health(elapsed.as_secs_f64(), failed);
62 if failed {
63 crate::logging::log(
64 "WARN",
65 "health",
66 &format!("Local health request to http://{address}/health failed"),
67 );
68 } else if elapsed > Duration::from_secs(1) {
69 crate::logging::log(
70 "WARN",
71 "health",
72 &format!(
73 "Local health request to http://{address}/health took more than 1s: {elapsed:?}"
74 ),
75 );
76 }
77 }
78 }))
79 }
80}
81
82impl Drop for HealthProbe {
83 fn drop(&mut self) {
84 self.0.abort();
85 }
86}
87
88async fn health_request(address: &str) -> io::Result<()> {
89 let stream = tokio::net::TcpStream::connect(address).await?;
90 let (mut sender, connection) = hyper::client::conn::http1::handshake(TokioIo::new(stream))
91 .await
92 .map_err(io::Error::other)?;
93 let request = Request::builder()
94 .uri("/health")
95 .header("host", address)
96 .body(Full::<Bytes>::default())
97 .map_err(io::Error::other)?;
98 let response = sender.send_request(request);
99 tokio::pin!(connection, response);
100 tokio::select! {
101 biased;
102 result = &mut response => result.map(|_| ()).map_err(io::Error::other),
103 result = &mut connection => result.map_err(io::Error::other).and_then(|()| Err(io::Error::other("health connection closed"))),
104 }
105}
106
107impl Listener {
108 pub async fn bind<F>(address: &str, route: F) -> io::Result<Self>
110 where
111 F: Fn(&str) -> Option<(&'static str, String)> + Send + Sync + 'static,
112 {
113 Self::bind_with_status(address, move |path| {
114 route(path).map(|(content_type, body)| (200, content_type, body))
115 })
116 .await
117 }
118
119 pub async fn bind_with_status<F>(address: &str, route: F) -> io::Result<Self>
121 where
122 F: Fn(&str) -> Option<(u16, &'static str, String)> + Send + Sync + 'static,
123 {
124 Self::bind_with_methods(address, move |method, path| {
125 (method == "GET").then(|| route(path)).flatten()
126 })
127 .await
128 }
129
130 pub async fn bind_with_methods<F>(address: &str, route: F) -> io::Result<Self>
132 where
133 F: Fn(&str, &str) -> Option<(u16, &'static str, String)> + Send + Sync + 'static,
134 {
135 let address = if address.ends_with(':') {
136 Cow::Owned(format!("{address}0"))
137 } else {
138 Cow::Borrowed(address)
139 };
140 let wildcard_port = address
141 .strip_prefix(':')
142 .and_then(|port| port.parse().ok())
143 .or_else(|| match address.parse::<SocketAddr>() {
144 Ok(SocketAddr::V6(addr)) if addr.ip().is_unspecified() => Some(addr.port()),
145 _ => None,
146 });
147 let listener = if let Some(port) = wildcard_port {
148 match dual_stack_listener(port)? {
149 Some(listener) => TcpListener::from_std(listener)?,
150 None => TcpListener::bind((Ipv4Addr::UNSPECIFIED, port)).await?,
151 }
152 } else {
153 TcpListener::bind(address.as_ref()).await?
154 };
155 let address = listener.local_addr()?;
156 let (stop, mut receiver) = watch::channel(false);
157 let route = Arc::new(route);
158 let task = tokio::spawn(async move {
159 let mut connections = JoinSet::new();
160 let result = loop {
161 tokio::select! {
162 _ = receiver.changed() => break Ok(()),
163 Some(result) = connections.join_next(), if !connections.is_empty() => {
164 if let Err(error) = result {
165 crate::logging::log("ERROR", "http", &format!("HTTP connection failed: {error}"));
166 }
167 },
168 accepted = listener.accept(), if connections.len() < MAX_CONNECTIONS => {
169 let (stream, _) = match accepted {
170 Ok(value) => value,
171 Err(error) => {
172 crate::logging::log("ERROR", "http", &format!("HTTP accept failed: {error}"));
173 tokio::select! {
174 _ = receiver.changed() => break Ok(()),
175 _ = tokio::time::sleep(Duration::from_millis(100)) => {},
176 }
177 continue;
178 }
179 };
180 let route = Arc::clone(&route);
181 let mut stop = receiver.clone();
182 connections.spawn(async move {
183 let service = service_fn(move |request: Request<Incoming>| {
184 let response = route(request.method().as_str(), request.uri().path());
185 async move {
186 let mut result = Response::new(Full::default());
187 match response {
188 Some((status, content_type, body)) => match (StatusCode::from_u16(status), HeaderValue::from_str(content_type)) {
189 (Ok(status), Ok(content_type)) => {
190 *result.status_mut() = status;
191 result.headers_mut().insert(CONTENT_TYPE, content_type);
192 *result.body_mut() = Full::new(Bytes::from(body));
193 }
194 _ => *result.status_mut() = StatusCode::INTERNAL_SERVER_ERROR,
195 },
196 None => *result.status_mut() = StatusCode::NOT_FOUND,
197 }
198 Ok::<_, Infallible>(result)
199 }
200 });
201 let connection = http1::Builder::new().serve_connection(TokioIo::new(stream), service);
202 tokio::pin!(connection);
203 tokio::select! {
204 _ = &mut connection => {},
205 _ = tokio::time::sleep(TIMEOUT) => {},
207 _ = stop.changed() => {
208 connection.as_mut().graceful_shutdown();
209 let _ = connection.await;
210 }
211 }
212 });
213 }
214 }
215 };
216 drop(listener);
217 if tokio::time::timeout(TIMEOUT, async {
218 while connections.join_next().await.is_some() {}
219 })
220 .await
221 .is_err()
222 {
223 connections.shutdown().await;
224 }
225 result
226 });
227 Ok(Self {
228 address,
229 stop,
230 task: Some(task),
231 })
232 }
233
234 #[must_use]
235 pub fn local_addr(&self) -> SocketAddr {
236 self.address
237 }
238
239 pub async fn shutdown(mut self) -> io::Result<()> {
241 self.stop.send_replace(true);
242 match self.task.take() {
243 Some(task) => task.await.map_err(io::Error::other)?,
244 None => Ok(()),
245 }
246 }
247}
248
249fn dual_stack_listener(port: u16) -> io::Result<Option<std::net::TcpListener>> {
250 let socket = match Socket::new(Domain::IPV6, Type::STREAM, Some(Protocol::TCP)) {
251 Ok(socket) => socket,
252 Err(error) if ipv6_unavailable(&error) => return Ok(None),
253 Err(error) => return Err(error),
254 };
255 socket.set_only_v6(false)?;
256 socket.set_reuse_address(true)?;
257 socket.set_nonblocking(true)?;
258 if let Err(error) = socket.bind(&SocketAddr::from((Ipv6Addr::UNSPECIFIED, port)).into()) {
259 if ipv6_unavailable(&error) {
260 return Ok(None);
261 }
262 return Err(error);
263 }
264 socket.listen(1024)?;
265 Ok(Some(socket.into()))
266}
267
268fn ipv6_unavailable(error: &io::Error) -> bool {
269 error.kind() == io::ErrorKind::AddrNotAvailable
270 || error.raw_os_error() == Some(libc::EAFNOSUPPORT)
271}
272
273impl Drop for Listener {
274 fn drop(&mut self) {
275 self.stop.send_replace(true);
276 }
277}
278
279#[cfg(test)]
280mod tests {
281 use super::Listener;
282 use std::time::Duration;
283 use tokio::io::{AsyncReadExt, AsyncWriteExt};
284 use tokio::net::TcpStream;
285
286 async fn request(addr: std::net::SocketAddr, request: &str) -> String {
287 let mut stream = TcpStream::connect(addr).await.unwrap();
288 stream.write_all(request.as_bytes()).await.unwrap();
289 let mut response = String::new();
290 stream.read_to_string(&mut response).await.unwrap();
291 response
292 }
293
294 #[tokio::test]
295 async fn routes_can_return_service_unavailable() {
296 let listener = Listener::bind_with_status("127.0.0.1:0", |path| {
297 (path == "/ready").then(|| (503, "text/plain", "not-ready".into()))
298 })
299 .await
300 .unwrap();
301 let response = request(
302 listener.local_addr(),
303 "GET /ready HTTP/1.1\r\nHost: localhost\r\nConnection: close\r\n\r\n",
304 )
305 .await;
306 assert!(response.starts_with("HTTP/1.1 503 Service Unavailable\r\n"));
307 assert!(response.ends_with("not-ready"));
308 listener.shutdown().await.unwrap();
309 }
310
311 #[tokio::test]
312 async fn routes_get_and_returns_not_found_for_other_requests() {
313 let listener = Listener::bind("localhost:0", |path| {
314 (path == "/metrics").then(|| ("text/plain; version=0.0.4", "metric 1\n".into()))
315 })
316 .await
317 .unwrap();
318 let address = listener.local_addr();
319 let response = request(
320 address,
321 "GET /metrics?x=1 HTTP/1.1\r\nHost: localhost\r\nConnection: close\r\n\r\n",
322 )
323 .await;
324 assert!(response.starts_with("HTTP/1.1 200 OK\r\n"));
325 assert!(response.contains("content-type: text/plain; version=0.0.4\r\n"));
326 assert!(response.ends_with("metric 1\n"));
327 for req in ["GET /missing", "POST /metrics"] {
328 let response = request(
329 address,
330 &format!("{req} HTTP/1.1\r\nHost: localhost\r\nConnection: close\r\n\r\n"),
331 )
332 .await;
333 assert!(response.starts_with("HTTP/1.1 404 Not Found\r\n"));
334 }
335 listener.shutdown().await.unwrap();
336 assert!(TcpStream::connect(address).await.is_err());
337 }
338
339 #[tokio::test]
340 async fn wildcard_listener_serves_ipv4_and_ipv6() {
341 let ipv6_available = std::net::TcpListener::bind("[::1]:0").is_ok();
342 for address in [":0", "[::]:0"] {
343 let listener = Listener::bind(address, |path| {
344 (path == "/metrics").then(|| ("text/plain", "metric 1\n".into()))
345 })
346 .await
347 .unwrap();
348 let port = listener.local_addr().port();
349 let query = "GET /metrics HTTP/1.1\r\nHost: localhost\r\nConnection: close\r\n\r\n";
350 let response = request(([127, 0, 0, 1], port).into(), query).await;
351 assert!(response.ends_with("metric 1\n"));
352 if ipv6_available {
353 let response = request((std::net::Ipv6Addr::LOCALHOST, port).into(), query).await;
354 assert!(response.ends_with("metric 1\n"));
355 }
356 listener.shutdown().await.unwrap();
357 }
358 }
359
360 #[tokio::test]
361 async fn wildcard_listener_does_not_fallback_when_port_is_occupied() {
362 let occupied = Listener::bind(":0", |_| None).await.unwrap();
363 let port = occupied.local_addr().port();
364 let error = match Listener::bind(&format!(":{port}"), |_| None).await {
365 Ok(_) => panic!("occupied wildcard port was accepted"),
366 Err(error) => error,
367 };
368 assert_eq!(error.kind(), std::io::ErrorKind::AddrInUse);
369 occupied.shutdown().await.unwrap();
370 }
371
372 #[test]
373 fn only_ipv6_unavailable_errors_allow_fallback() {
374 use super::ipv6_unavailable;
375 use std::io::{Error, ErrorKind};
376
377 assert!(ipv6_unavailable(&Error::from_raw_os_error(
378 libc::EAFNOSUPPORT
379 )));
380 assert!(ipv6_unavailable(&Error::from(ErrorKind::AddrNotAvailable)));
381 assert!(!ipv6_unavailable(&Error::from(ErrorKind::AddrInUse)));
382 assert!(!ipv6_unavailable(&Error::from(ErrorKind::PermissionDenied)));
383 }
384
385 #[tokio::test]
386 async fn partial_request_times_out() {
387 let listener = Listener::bind("127.0.0.1:0", |_| None).await.unwrap();
388 let mut stream = TcpStream::connect(listener.local_addr()).await.unwrap();
389 stream.write_all(b"GET / HTTP/1.1\r\n").await.unwrap();
390 let mut bytes = Vec::new();
391 let result =
392 tokio::time::timeout(Duration::from_secs(6), stream.read_to_end(&mut bytes)).await;
393 assert!(result.is_ok(), "partial HTTP request did not time out");
394 listener.shutdown().await.unwrap();
395 }
396
397 #[tokio::test]
398 async fn shutdown_closes_idle_connections() {
399 let listener = Listener::bind("127.0.0.1:0", |_| None).await.unwrap();
400 let mut stream = TcpStream::connect(listener.local_addr()).await.unwrap();
401 stream
402 .write_all(b"GET / HTTP/1.1\r\nHost: localhost\r\n\r\n")
403 .await
404 .unwrap();
405 let mut response = [0; 1024];
406 assert!(stream.read(&mut response).await.unwrap() > 0);
407 stream.write_all(b"GET / HTTP/1.1\r\n").await.unwrap();
408 tokio::time::sleep(Duration::from_millis(20)).await;
409 tokio::time::timeout(Duration::from_secs(6), listener.shutdown())
410 .await
411 .unwrap()
412 .unwrap();
413 let mut bytes = Vec::new();
414 let result = stream.read_to_end(&mut bytes).await;
415 assert!(
416 result.is_ok() || result.unwrap_err().kind() == std::io::ErrorKind::ConnectionReset
417 );
418 }
419
420 #[tokio::test]
421 async fn shutdown_drains_response_in_progress() {
422 let body = "x".repeat(4 * 1024 * 1024);
423 let length = body.len();
424 let listener = Listener::bind("127.0.0.1:0", move |_| Some(("text/plain", body.clone())))
425 .await
426 .unwrap();
427 let mut stream = TcpStream::connect(listener.local_addr()).await.unwrap();
428 stream
429 .write_all(b"GET / HTTP/1.1\r\nHost: localhost\r\n\r\n")
430 .await
431 .unwrap();
432 let mut first = [0; 1024];
433 let n = stream.read(&mut first).await.unwrap();
434 assert!(n > 0);
435 let shutdown = tokio::spawn(listener.shutdown());
436 let mut response = first[..n].to_vec();
437 stream.read_to_end(&mut response).await.unwrap();
438 shutdown.await.unwrap().unwrap();
439 let header_end = response.windows(4).position(|w| w == b"\r\n\r\n").unwrap() + 4;
440 assert_eq!(response.len() - header_end, length);
441 }
442}