1use std::future::Future;
22use std::pin::Pin;
23use std::task::{Context, Poll};
24
25use http::header::HeaderValue;
26use http::{HeaderMap, Response};
27use tower_layer::Layer;
28use tower_service::Service;
29
30use crate::error::Error;
31use crate::link::{self, ManifestLink};
32
33pub const LINK: http::header::HeaderName = http::header::LINK;
35
36pub fn append_to(headers: &mut HeaderMap, uri: &str) -> Result<(), Error> {
43 let value = link::format(uri)?;
44 let header = HeaderValue::from_str(&value)
45 .map_err(|_| Error::Malformed("target URI is not a valid header value"))?;
46 headers.append(LINK, header);
47 Ok(())
48}
49
50pub fn extract_from(headers: &HeaderMap) -> Result<ManifestLink, Error> {
55 link::extract(headers.get_all(LINK).iter().filter_map(|v| v.to_str().ok()))
56}
57
58#[derive(Debug, Clone)]
63pub struct ManifestLinkLayer {
64 header: HeaderValue,
65}
66
67impl ManifestLinkLayer {
68 pub fn new(uri: &str) -> Result<Self, Error> {
73 Self::from_value(link::format(uri)?)
74 }
75
76 pub fn new_strict(uri: &str) -> Result<Self, Error> {
83 Self::from_value(link::format_strict(uri)?)
84 }
85
86 fn from_value(value: String) -> Result<Self, Error> {
87 let header = HeaderValue::from_str(&value)
88 .map_err(|_| Error::Malformed("target URI is not a valid header value"))?;
89 Ok(Self { header })
90 }
91}
92
93impl<S> Layer<S> for ManifestLinkLayer {
94 type Service = ManifestLinkService<S>;
95
96 fn layer(&self, inner: S) -> Self::Service {
97 ManifestLinkService {
98 inner,
99 header: self.header.clone(),
100 }
101 }
102}
103
104#[derive(Debug, Clone)]
106pub struct ManifestLinkService<S> {
107 inner: S,
108 header: HeaderValue,
109}
110
111impl<S, Request, B> Service<Request> for ManifestLinkService<S>
112where
113 S: Service<Request, Response = Response<B>>,
114{
115 type Response = S::Response;
116 type Error = S::Error;
117 type Future = ResponseFuture<S::Future>;
118
119 fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
120 self.inner.poll_ready(cx)
121 }
122
123 fn call(&mut self, request: Request) -> Self::Future {
124 ResponseFuture {
125 inner: self.inner.call(request),
126 header: self.header.clone(),
127 }
128 }
129}
130
131pin_project_lite::pin_project! {
132 #[derive(Debug)]
134 pub struct ResponseFuture<F> {
135 #[pin]
136 inner: F,
137 header: HeaderValue,
138 }
139}
140
141impl<F, B, E> Future for ResponseFuture<F>
142where
143 F: Future<Output = Result<Response<B>, E>>,
144{
145 type Output = Result<Response<B>, E>;
146
147 fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
148 let this = self.project();
149 let mut response = match this.inner.poll(cx) {
150 Poll::Pending => return Poll::Pending,
151 Poll::Ready(Err(e)) => return Poll::Ready(Err(e)),
152 Poll::Ready(Ok(r)) => r,
153 };
154 response.headers_mut().append(LINK, this.header.clone());
155 Poll::Ready(Ok(response))
156 }
157}
158
159#[cfg(test)]
160mod tests {
161 use super::*;
162 use tower::{ServiceBuilder, ServiceExt};
163
164 const URI: &str = "https://a.example/m.c2pa";
165
166 async fn ok(_: http::Request<()>) -> Result<Response<()>, std::convert::Infallible> {
167 Ok(Response::new(()))
168 }
169
170 async fn with_existing_link(
171 _: http::Request<()>,
172 ) -> Result<Response<()>, std::convert::Infallible> {
173 let mut r = Response::new(());
174 r.headers_mut()
175 .append(LINK, HeaderValue::from_static("</s.css>; rel=preload"));
176 Ok(r)
177 }
178
179 #[tokio::test]
180 async fn the_layer_advertises_the_manifest() {
181 let svc = ServiceBuilder::new()
182 .layer(ManifestLinkLayer::new(URI).unwrap())
183 .service_fn(ok);
184 let response = svc.oneshot(http::Request::new(())).await.unwrap();
185 assert_eq!(extract_from(response.headers()).unwrap().uri, URI);
186 }
187
188 #[tokio::test]
189 async fn an_existing_link_header_survives() {
190 let svc = ServiceBuilder::new()
192 .layer(ManifestLinkLayer::new(URI).unwrap())
193 .service_fn(with_existing_link);
194 let response = svc.oneshot(http::Request::new(())).await.unwrap();
195 assert_eq!(response.headers().get_all(LINK).iter().count(), 2);
196 assert_eq!(extract_from(response.headers()).unwrap().uri, URI);
197 }
198
199 #[test]
200 fn append_and_extract_round_trip() {
201 let mut headers = HeaderMap::new();
202 assert_eq!(extract_from(&headers), Err(Error::NotFound));
203 append_to(&mut headers, URI).unwrap();
204 assert_eq!(extract_from(&headers).unwrap().uri, URI);
205 }
206
207 #[test]
208 fn append_preserves_unrelated_links() {
209 let mut headers = HeaderMap::new();
210 headers.append(LINK, HeaderValue::from_static("</a>; rel=next"));
211 append_to(&mut headers, URI).unwrap();
212 assert_eq!(headers.get_all(LINK).iter().count(), 2);
213 assert_eq!(extract_from(&headers).unwrap().uri, URI);
214 }
215
216 #[test]
217 fn a_hostile_target_cannot_inject_a_header() {
218 let mut headers = HeaderMap::new();
221 append_to(&mut headers, "https://a.example/\r\nX-Evil: 1").unwrap();
222 assert_eq!(headers.len(), 1, "a second header was injected");
223 assert!(headers.get("x-evil").is_none());
224 assert!(extract_from(&headers).unwrap().uri.contains("%0D%0A"));
225 }
226
227 #[tokio::test]
228 async fn a_hostile_target_stays_inert_through_the_layer() {
229 let svc = ServiceBuilder::new()
230 .layer(ManifestLinkLayer::new("https://a.example/\r\nX-Evil: 1").unwrap())
231 .service_fn(ok);
232 let response = svc.oneshot(http::Request::new(())).await.unwrap();
233 assert_eq!(response.headers().len(), 1);
234 assert!(response.headers().get("x-evil").is_none());
235 }
236
237 #[test]
238 fn strict_construction_rejects_what_lenient_repairs() {
239 assert!(ManifestLinkLayer::new_strict("https://a.example/a b").is_err());
242 assert!(ManifestLinkLayer::new("https://a.example/a b").is_ok());
243 assert!(ManifestLinkLayer::new_strict("https://a.example/m.c2pa").is_ok());
244 }
245
246 #[test]
247 fn a_jumbf_target_survives_the_header_round_trip() {
248 let mut headers = HeaderMap::new();
249 append_to(&mut headers, "https://a.example/i.jpg#jumbf=c2pa").unwrap();
250 let found = extract_from(&headers).unwrap();
251 assert!(found.is_embedded());
252 assert_eq!(found.jumbf.as_deref(), Some("c2pa"));
253 }
254
255 #[test]
256 fn competing_targets_across_two_fields_are_rejected() {
257 let mut headers = HeaderMap::new();
258 append_to(&mut headers, "https://a.example/a.c2pa").unwrap();
259 append_to(&mut headers, "https://a.example/b.c2pa").unwrap();
260 assert_eq!(extract_from(&headers), Err(Error::MultipleLinks));
261 }
262
263 #[test]
264 fn a_non_utf8_header_value_does_not_hide_a_good_one() {
265 let mut headers = HeaderMap::new();
266 headers.append(
267 LINK,
268 HeaderValue::from_bytes(b"</\xFF\xFE>; rel=next").unwrap(),
269 );
270 append_to(&mut headers, URI).unwrap();
271 assert_eq!(extract_from(&headers).unwrap().uri, URI);
272 }
273}