actix_proxy_protocol/service/
acceptor.rs

1use std::convert::Infallible;
2
3use actix_rt::net::ActixStream;
4use actix_service::{Service, ServiceFactory};
5use actix_utils::future::{Ready, ready};
6use futures_core::future::LocalBoxFuture;
7
8use super::{HeaderPolicy, ProxyProtocolError, ProxyStream};
9
10/// Actix service factory that wraps accepted streams in [`ProxyStream`].
11#[derive(Debug, Clone, Copy, Default)]
12pub struct Acceptor {
13    policy: HeaderPolicy,
14}
15
16impl Acceptor {
17    /// Constructs an acceptor that requires a PROXY protocol header.
18    pub const fn new() -> Self {
19        Self {
20            policy: HeaderPolicy::Required,
21        }
22    }
23
24    /// Constructs an acceptor using `policy`.
25    pub const fn with_policy(policy: HeaderPolicy) -> Self {
26        Self { policy }
27    }
28
29    /// Constructs an acceptor that allows streams without a PROXY protocol header.
30    pub const fn optional() -> Self {
31        Self {
32            policy: HeaderPolicy::Optional,
33        }
34    }
35}
36
37impl<IO> ServiceFactory<IO> for Acceptor
38where
39    IO: ActixStream + 'static,
40{
41    type Response = ProxyStream<IO>;
42    type Error = ProxyProtocolError<Infallible>;
43    type Config = ();
44    type Service = AcceptorService;
45    type InitError = ();
46    type Future = Ready<Result<Self::Service, Self::InitError>>;
47
48    fn new_service(&self, _: ()) -> Self::Future {
49        ready(Ok(AcceptorService {
50            policy: self.policy,
51        }))
52    }
53}
54
55/// Actix service that wraps streams in [`ProxyStream`].
56#[derive(Debug, Clone, Copy)]
57pub struct AcceptorService {
58    policy: HeaderPolicy,
59}
60
61impl<IO> Service<IO> for AcceptorService
62where
63    IO: ActixStream + 'static,
64{
65    type Response = ProxyStream<IO>;
66    type Error = ProxyProtocolError<Infallible>;
67    type Future = LocalBoxFuture<'static, Result<Self::Response, Self::Error>>;
68
69    actix_service::always_ready!();
70
71    fn call(&self, io: IO) -> Self::Future {
72        let policy = self.policy;
73
74        Box::pin(async move { ProxyStream::accept_with_policy(io, policy).await })
75    }
76}
77
78#[cfg(test)]
79mod tests {
80    use actix_service::Service as _;
81    use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _};
82
83    use super::*;
84    use crate::{Header, Version, v1};
85
86    #[actix_rt::test]
87    async fn acceptor_service_wraps_streams() {
88        let listener = actix_rt::net::TcpListener::bind(("127.0.0.1", 0))
89            .await
90            .unwrap();
91        let local_addr = listener.local_addr().unwrap();
92        let factory = Acceptor::new();
93        let acceptor =
94            <Acceptor as ServiceFactory<actix_rt::net::TcpStream>>::new_service(&factory, ())
95                .await
96                .unwrap();
97
98        actix_rt::spawn(async move {
99            let mut client = actix_rt::net::TcpStream::connect(local_addr).await.unwrap();
100            v1::Header::unknown()
101                .write_to_tokio(&mut client)
102                .await
103                .unwrap();
104            client.write_all(b"hello").await.unwrap();
105            client.shutdown().await.unwrap();
106        });
107
108        let (server, _) = listener.accept().await.unwrap();
109        let mut stream = acceptor.call(server).await.unwrap();
110
111        assert_eq!(stream.header().map(Header::version), Some(Version::V1));
112
113        let mut body = String::new();
114        stream.read_to_string(&mut body).await.unwrap();
115        assert_eq!(body, "hello");
116    }
117}
actix_proxy_protocol/service/acceptor.rs

actix_proxy_protocol/service/
acceptor.rs
