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