rama_http/layer/dns/dns_resolve/
layer.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
use super::DnsResolveModeService;
use crate::HeaderName;
use rama_core::Layer;

/// Layer which can extend `Dns` (see `rama_core`) overwrites with mappings.
///
/// See [the module level documentation](crate::layer::dns) for more information.
#[derive(Debug, Clone)]
pub struct DnsResolveModeLayer {
    header_name: HeaderName,
}

impl DnsResolveModeLayer {
    /// Creates a new [`DnsResolveModeLayer`].
    pub const fn new(name: HeaderName) -> Self {
        Self { header_name: name }
    }
}

impl<S> Layer<S> for DnsResolveModeLayer {
    type Service = DnsResolveModeService<S>;

    fn layer(&self, inner: S) -> Self::Service {
        DnsResolveModeService::new(inner, self.header_name.clone())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{layer::dns::DnsResolveMode, Request};
    use rama_core::{service::service_fn, Context, Service};
    use std::convert::Infallible;

    #[tokio::test]
    async fn test_dns_resolve_mode_layer() {
        let svc = DnsResolveModeLayer::new(HeaderName::from_static("x-dns-resolve")).layer(
            service_fn(|ctx: Context<()>, _req: Request<()>| async move {
                assert_eq!(
                    ctx.get::<DnsResolveMode>().unwrap(),
                    &DnsResolveMode::eager()
                );
                Ok::<_, Infallible>(())
            }),
        );

        let ctx = Context::default();
        let req = Request::builder()
            .header("x-dns-resolve", "eager")
            .uri("http://example.com")
            .body(())
            .unwrap();

        svc.serve(ctx, req).await.unwrap();
    }
}