rama_ua/layer/
classifier.rs

1//! User-Agent http layer support
2//!
3//! # Example
4//!
5//! ```
6//! use rama_http::{
7//!     Request, Response, StatusCode, service::client::HttpClientExt,
8//!     service::web::response::IntoResponse,
9//! };
10//! use rama_ua::{
11//!     PlatformKind, UserAgent, UserAgentInfo, UserAgentKind,
12//!     layer::classifier::UserAgentClassifierLayer,
13//! };
14//! use rama_core::{extensions::ExtensionsRef, Layer, service::service_fn};
15//! use std::convert::Infallible;
16//!
17//! const UA: &str = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/124.0.0.0 Safari/537.36 Edg/124.0.2478.67";
18//!
19//! async fn handle(req: Request) -> Result<Response, Infallible> {
20//!     let ua: &UserAgent = req.extensions().get_ref().unwrap();
21//!
22//!     assert_eq!(ua.header_str(), UA);
23//!     assert_eq!(ua.info(), Some(UserAgentInfo{ kind: UserAgentKind::Chromium, version: Some(124) }));
24//!     assert_eq!(ua.platform(), Some(PlatformKind::Windows));
25//!
26//!     Ok(StatusCode::OK.into_response())
27//! }
28//!
29//! # #[tokio::main]
30//! # async fn main() {
31//! let service = UserAgentClassifierLayer::new().into_layer(service_fn(handle));
32//!
33//! _ = service
34//!     .get("http://www.example.com")
35//!     .typed_header(rama_http::headers::UserAgent::from_static(UA))
36//!     .send()
37//!     .await
38//!     .unwrap();
39//! # }
40//! ```
41
42use rama_core::{Layer, Service, extensions::ExtensionsRef};
43use rama_http::{
44    HeaderName, Request,
45    headers::{self, HeaderMapExt},
46};
47use rama_utils::macros::define_inner_service_accessors;
48use std::fmt::Debug;
49
50use crate::{UserAgent, UserAgentOverwrites};
51
52/// A [`Service`] that classifies the [`UserAgent`] of incoming [`Request`]s.
53///
54/// The [`Extensions`] is updated with the [`UserAgent`]
55/// if the [`Request`] contains a valid [`UserAgent`] header.
56///
57/// [`Extensions`]: rama_core::extensions::Extensions
58#[derive(Debug, Clone, Default)]
59pub struct UserAgentClassifier<S> {
60    inner: S,
61    overwrite_header: Option<HeaderName>,
62}
63
64impl<S> UserAgentClassifier<S> {
65    /// Create a new [`UserAgentClassifier`] [`Service`].
66    pub const fn new(inner: S, overwrite_header: Option<HeaderName>) -> Self {
67        Self {
68            inner,
69            overwrite_header,
70        }
71    }
72
73    define_inner_service_accessors!();
74}
75
76impl<S, Body> Service<Request<Body>> for UserAgentClassifier<S>
77where
78    S: Service<Request<Body>>,
79{
80    type Output = S::Output;
81    type Error = S::Error;
82
83    fn serve(
84        &self,
85        req: Request<Body>,
86    ) -> impl Future<Output = Result<Self::Output, Self::Error>> + Send + '_ {
87        let overwrites = self
88            .overwrite_header
89            .as_ref()
90            .and_then(|header| req.headers().get(header))
91            .map(|header| header.as_bytes())
92            .and_then(|value| serde_html_form::from_bytes::<UserAgentOverwrites>(value).ok());
93
94        let mut user_agent = overwrites
95            .as_ref()
96            .and_then(|o| o.ua.as_deref())
97            .map(UserAgent::new)
98            .or_else(|| {
99                req.headers()
100                    .typed_get::<headers::UserAgent>()
101                    .map(|ua| UserAgent::new(ua.to_string()))
102            });
103
104        if let Some(mut ua) = user_agent.take() {
105            if let Some(overwrites) = overwrites {
106                if let Some(http_agent) = overwrites.http {
107                    ua.set_http_agent(http_agent);
108                }
109                if let Some(tls_agent) = overwrites.tls {
110                    ua.set_tls_agent(tls_agent);
111                }
112            }
113
114            req.extensions().insert(ua);
115        }
116
117        self.inner.serve(req)
118    }
119}
120
121#[derive(Debug, Clone, Default)]
122/// A [`Layer`] that wraps a [`Service`] with a [`UserAgentClassifier`].
123///
124/// This [`Layer`] is used to classify the [`UserAgent`] of incoming [`Request`]s.
125pub struct UserAgentClassifierLayer {
126    overwrite_header: Option<HeaderName>,
127}
128
129impl UserAgentClassifierLayer {
130    /// Create a new [`UserAgentClassifierLayer`].
131    #[must_use]
132    pub const fn new() -> Self {
133        Self {
134            overwrite_header: None,
135        }
136    }
137
138    rama_utils::macros::generate_set_and_with! {
139        /// Define a custom header to allow overwriting certain
140        /// [`UserAgent`] information.
141        pub fn overwrite_header(mut self, header: HeaderName) -> Self {
142            self.overwrite_header = Some(header);
143            self
144        }
145    }
146}
147
148impl<S> Layer<S> for UserAgentClassifierLayer {
149    type Service = UserAgentClassifier<S>;
150
151    fn layer(&self, inner: S) -> Self::Service {
152        UserAgentClassifier::new(inner, self.overwrite_header.clone())
153    }
154
155    fn into_layer(self, inner: S) -> Self::Service {
156        UserAgentClassifier::new(inner, self.overwrite_header)
157    }
158}
159
160#[cfg(test)]
161mod tests {
162    use crate::{HttpAgent, PlatformKind, TlsAgent, UserAgentKind};
163
164    use super::*;
165    use rama_core::extensions::ExtensionsRef;
166    use rama_core::service::service_fn;
167    use rama_http::layer::required_header::AddRequiredRequestHeadersLayer;
168    use rama_http::service::client::HttpClientExt;
169    use rama_http::service::web::response::IntoResponse;
170    use rama_http::{Response, StatusCode, headers};
171    use std::convert::Infallible;
172
173    #[tokio::test]
174    async fn test_user_agent_classifier_layer_ua_rama() {
175        async fn handle(req: Request) -> Result<Response, Infallible> {
176            let ua: &UserAgent = req.extensions().get_ref().unwrap();
177
178            assert_eq!(
179                ua.header_str(),
180                format!("{}/{}", rama_utils::info::NAME, rama_utils::info::VERSION).as_str(),
181            );
182            assert!(ua.info().is_none());
183            assert!(ua.platform().is_none());
184
185            Ok(StatusCode::OK.into_response())
186        }
187
188        let service = (
189            AddRequiredRequestHeadersLayer::default(),
190            UserAgentClassifierLayer::new(),
191        )
192            .into_layer(service_fn(handle));
193
194        _ = service.get("http://www.example.com").send().await.unwrap();
195    }
196
197    #[tokio::test]
198    async fn test_user_agent_classifier_layer_ua_iphone_app() {
199        const UA: &str = "iPhone App/1.0";
200
201        async fn handle(req: Request) -> Result<Response, Infallible> {
202            let ua: &UserAgent = req.extensions().get_ref().unwrap();
203
204            assert_eq!(ua.header_str(), UA);
205            assert!(ua.info().is_none());
206            assert_eq!(ua.platform(), Some(PlatformKind::IOS));
207            assert_eq!(ua.http_agent(), None);
208            assert_eq!(ua.tls_agent(), None);
209
210            Ok(StatusCode::OK.into_response())
211        }
212
213        let service = UserAgentClassifierLayer::new().into_layer(service_fn(handle));
214
215        _ = service
216            .get("http://www.example.com")
217            .typed_header(headers::UserAgent::from_static(UA))
218            .send()
219            .await
220            .unwrap();
221    }
222
223    #[tokio::test]
224    async fn test_user_agent_classifier_layer_ua_chrome() {
225        const UA: &str = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/124.0.0.0 Safari/537.36 Edg/124.0.2478.67";
226
227        async fn handle(req: Request) -> Result<Response, Infallible> {
228            let ua: &UserAgent = req.extensions().get_ref().unwrap();
229
230            assert_eq!(ua.header_str(), UA);
231            let ua_info = ua.info().unwrap();
232            assert_eq!(ua_info.kind, UserAgentKind::Chromium);
233            assert_eq!(ua_info.version, Some(124));
234            assert_eq!(ua.platform(), Some(PlatformKind::Windows));
235
236            Ok(StatusCode::OK.into_response())
237        }
238
239        let service = UserAgentClassifierLayer::new().into_layer(service_fn(handle));
240
241        _ = service
242            .get("http://www.example.com")
243            .typed_header(headers::UserAgent::from_static(UA))
244            .send()
245            .await
246            .unwrap();
247    }
248
249    #[tokio::test]
250    async fn test_user_agent_classifier_layer_overwrite_ua() {
251        const UA: &str = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/124.0.0.0 Safari/537.36 Edg/124.0.2478.67";
252
253        async fn handle(req: Request) -> Result<Response, Infallible> {
254            let ua: &UserAgent = req.extensions().get_ref().unwrap();
255
256            assert_eq!(ua.header_str(), UA);
257            let ua_info = ua.info().unwrap();
258            assert_eq!(ua_info.kind, UserAgentKind::Chromium);
259            assert_eq!(ua_info.version, Some(124));
260            assert_eq!(ua.platform(), Some(PlatformKind::Windows));
261
262            Ok(StatusCode::OK.into_response())
263        }
264
265        let service = UserAgentClassifierLayer::new()
266            .with_overwrite_header(HeaderName::from_static("x-proxy-ua"))
267            .into_layer(service_fn(handle));
268
269        _ = service
270            .get("http://www.example.com")
271            .header(
272                "x-proxy-ua",
273                serde_html_form::to_string(&UserAgentOverwrites {
274                    ua: Some(UA.to_owned()),
275                    ..Default::default()
276                })
277                .unwrap(),
278            )
279            .send()
280            .await
281            .unwrap();
282    }
283
284    #[tokio::test]
285    async fn test_user_agent_classifier_layer_overwrite_ua_all() {
286        const UA: &str = "iPhone App/1.0";
287
288        async fn handle(req: Request) -> Result<Response, Infallible> {
289            let ua: &UserAgent = req.extensions().get_ref().unwrap();
290
291            assert_eq!(ua.header_str(), UA);
292            assert!(ua.info().is_none());
293            assert_eq!(ua.platform(), Some(PlatformKind::IOS));
294            assert_eq!(ua.http_agent(), Some(HttpAgent::Firefox));
295            assert_eq!(ua.tls_agent(), Some(TlsAgent::Boringssl));
296
297            Ok(StatusCode::OK.into_response())
298        }
299
300        let service = UserAgentClassifierLayer::new()
301            .with_overwrite_header(HeaderName::from_static("x-proxy-ua"))
302            .into_layer(service_fn(handle));
303
304        _ = service
305            .get("http://www.example.com")
306            .header(
307                "x-proxy-ua",
308                serde_html_form::to_string(&UserAgentOverwrites {
309                    ua: Some(UA.to_owned()),
310                    http: Some(HttpAgent::Firefox),
311                    tls: Some(TlsAgent::Boringssl),
312                })
313                .unwrap(),
314            )
315            .send()
316            .await
317            .unwrap();
318    }
319}
rama_ua/layer/classifier.rs

rama_ua/layer/
classifier.rs
