rust_web_server/early_hints/
mod.rs1use std::cell::RefCell;
31use std::sync::atomic::{AtomicUsize, Ordering};
32use std::sync::Arc;
33
34use crate::application::Application;
35use crate::middleware::Middleware;
36use crate::request::Request;
37use crate::response::Response;
38use crate::server::ConnectionInfo;
39
40pub const MAX_PER_REQUEST: usize = 4;
42
43pub(crate) struct Sink {
45 deliver: Box<dyn Fn(Vec<String>) + Send + Sync>,
46 sent: AtomicUsize,
47}
48
49impl Sink {
50 pub(crate) fn new(deliver: impl Fn(Vec<String>) + Send + Sync + 'static) -> Arc<Sink> {
51 Arc::new(Sink { deliver: Box::new(deliver), sent: AtomicUsize::new(0) })
52 }
53
54 #[cfg_attr(not(feature = "http2"), allow(dead_code))]
56 pub(crate) fn reset(&self) {
57 self.sent.store(0, Ordering::Relaxed);
58 }
59
60 fn send(&self, links: Vec<String>) -> bool {
61 if self.sent.fetch_add(1, Ordering::Relaxed) >= MAX_PER_REQUEST {
62 return false;
63 }
64 (self.deliver)(links);
65 true
66 }
67}
68
69thread_local! {
70 static CURRENT: RefCell<Option<Arc<Sink>>> = const { RefCell::new(None) };
71}
72
73#[cfg(feature = "http2")]
74tokio::task_local! {
75 static TASK: Option<Arc<Sink>>;
76}
77
78pub(crate) fn with_sink<R>(sink: Option<Arc<Sink>>, f: impl FnOnce() -> R) -> R {
81 struct Restore(Option<Arc<Sink>>);
82 impl Drop for Restore {
83 fn drop(&mut self) {
84 let previous = self.0.take();
85 CURRENT.with(|c| *c.borrow_mut() = previous);
86 }
87 }
88 let previous = CURRENT.with(|c| std::mem::replace(&mut *c.borrow_mut(), sink));
89 let _restore = Restore(previous);
90 f()
91}
92
93#[cfg(feature = "http2")]
96pub(crate) async fn scope<F: std::future::Future>(sink: Option<Arc<Sink>>, future: F) -> F::Output {
97 TASK.scope(sink, future).await
98}
99
100fn current() -> Option<Arc<Sink>> {
101 #[cfg(feature = "http2")]
102 if let Ok(Some(sink)) = TASK.try_with(|s| s.clone()) {
103 return Some(sink);
104 }
105 CURRENT.with(|c| c.borrow().clone())
106}
107
108pub fn available() -> bool {
110 current().is_some()
111}
112
113pub fn send<S: AsRef<str>>(links: &[S]) -> bool {
118 send_from(Source::Handler, links)
119}
120
121#[derive(Clone, Copy, Debug, PartialEq, Eq)]
123pub(crate) enum Source {
124 Handler,
126 Configured,
128 Upstream,
130}
131
132impl Source {
133 fn label(self) -> &'static str {
134 match self {
135 Source::Handler => "handler",
136 Source::Configured => "configured",
137 Source::Upstream => "upstream",
138 }
139 }
140}
141
142static SENT: [AtomicUsize; 3] = [AtomicUsize::new(0), AtomicUsize::new(0), AtomicUsize::new(0)];
143
144pub(crate) fn send_from<S: AsRef<str>>(source: Source, links: &[S]) -> bool {
145 if links.is_empty() || !links.iter().all(|l| valid(l.as_ref())) {
146 return false;
147 }
148 let Some(sink) = current() else { return false };
149 let sent = sink.send(links.iter().map(|l| l.as_ref().to_string()).collect());
150 if sent {
151 SENT[source as usize].fetch_add(1, Ordering::Relaxed);
152 }
153 sent
154}
155
156pub(crate) fn upstream_links(head: &[u8]) -> Vec<String> {
158 let text = String::from_utf8_lossy(head);
159 let mut lines = text.split("\r\n");
160 let is_103 = lines.next().and_then(|l| l.split(' ').nth(1)) == Some("103");
161 if !is_103 {
162 return Vec::new();
163 }
164 lines
165 .filter_map(|line| line.split_once(':'))
166 .filter(|(name, _)| name.trim().eq_ignore_ascii_case("link"))
167 .map(|(_, value)| value.trim().to_string())
168 .filter(|value| valid(value))
169 .collect()
170}
171
172pub(crate) fn forward_upstream(head: &[u8]) {
175 let links = upstream_links(head);
176 if !links.is_empty() {
177 send_from(Source::Upstream, &links);
178 }
179}
180
181pub(crate) fn prometheus_text() -> String {
183 let counts: Vec<usize> = SENT.iter().map(|c| c.load(Ordering::Relaxed)).collect();
184 if counts.iter().all(|n| *n == 0) {
185 return String::new();
186 }
187 let mut out = String::from("# HELP rws_early_hints_total 103 Early Hints responses sent, by where the hints came from\n# TYPE rws_early_hints_total counter\n");
188 for source in [Source::Handler, Source::Configured, Source::Upstream] {
189 out.push_str(&format!("rws_early_hints_total{{source=\"{}\"}} {}\n", source.label(), counts[source as usize]));
190 }
191 out
192}
193
194pub fn valid(link: &str) -> bool {
196 !link.trim().is_empty() && !link.bytes().any(|b| matches!(b, b'\r' | b'\n' | 0))
197}
198
199pub(crate) fn http1_wanted(request: &Request) -> bool {
202 request.http_version == "HTTP/1.1" && crate::server_config::ServerConfig::current().early_hints_http1
203}
204
205pub(crate) fn http1_head(links: &[String]) -> Vec<u8> {
207 let mut head = b"HTTP/1.1 103 Early Hints\r\n".to_vec();
208 for link in links {
209 head.extend_from_slice(b"Link: ");
210 head.extend_from_slice(link.as_bytes());
211 head.extend_from_slice(b"\r\n");
212 }
213 head.extend_from_slice(b"\r\n");
214 head
215}
216
217#[derive(Clone, Debug)]
220pub struct EarlyHintsLayer {
221 links: Arc<Vec<String>>,
222}
223
224impl EarlyHintsLayer {
225 pub fn new<I, S>(links: I) -> EarlyHintsLayer
228 where
229 I: IntoIterator<Item = S>,
230 S: Into<String>,
231 {
232 EarlyHintsLayer { links: Arc::new(links.into_iter().map(Into::into).filter(|l: &String| valid(l)).collect()) }
233 }
234
235 pub fn links(&self) -> &[String] {
236 &self.links
237 }
238}
239
240impl Middleware for EarlyHintsLayer {
241 fn handle(&self, request: &Request, connection: &ConnectionInfo, next: &dyn Application) -> Result<Response, String> {
242 if !self.links.is_empty() {
243 send_from(Source::Configured, &self.links);
244 }
245 next.execute(request, connection)
246 }
247
248 fn runs_inline(&self, _request: &Request, _connection: &ConnectionInfo) -> bool {
250 true
251 }
252
253 fn accepts_streamed_body(&self, _request: &Request, _connection: &ConnectionInfo) -> bool {
254 true
255 }
256}
257
258#[cfg(test)]
259mod tests;