actix_proxy_protocol/service/
stream.rs

1use std::{
2    io,
3    pin::Pin,
4    task::{Context, Poll},
5};
6
7use actix_rt::net::{ActixStream, Ready};
8use proxyproto::{Header, ParseError, v1, v2};
9use tokio::io::{AsyncRead, AsyncReadExt as _, AsyncWrite, ReadBuf};
10
11use super::{HeaderPolicy, ProxyProtocolError};
12
13pin_project_lite::pin_project! {
14    /// Stream wrapper that consumes a leading PROXY protocol header and then behaves like `IO`.
15    #[derive(Debug)]
16    pub struct ProxyStream<IO> {
17        #[pin]
18        io: IO,
19        header: Option<Header>,
20        pending: Vec<u8>,
21    }
22}
23
24impl<IO> ProxyStream<IO> {
25    /// Constructs a wrapper from an already parsed header and stream.
26    pub fn new(io: IO, header: Option<Header>) -> Self {
27        Self {
28            io,
29            header,
30            pending: Vec::new(),
31        }
32    }
33
34    /// Returns the parsed PROXY protocol header, if one was present.
35    pub const fn header(&self) -> Option<&Header> {
36        self.header.as_ref()
37    }
38
39    /// Removes and returns the parsed PROXY protocol header, if one was present.
40    pub fn take_header(&mut self) -> Option<Header> {
41        self.header.take()
42    }
43
44    /// Returns a shared reference to the wrapped stream.
45    pub const fn get_ref(&self) -> &IO {
46        &self.io
47    }
48
49    /// Returns a mutable reference to the wrapped stream.
50    pub fn get_mut(&mut self) -> &mut IO {
51        &mut self.io
52    }
53
54    /// Consumes the wrapper and returns the wrapped stream and parsed header.
55    pub fn into_parts(self) -> (IO, Option<Header>) {
56        (self.io, self.header)
57    }
58}
59
60impl<IO> ProxyStream<IO>
61where
62    IO: AsyncRead + AsyncWrite + Unpin,
63{
64    /// Reads and consumes a required PROXY protocol header from `io`.
65    pub async fn accept(io: IO) -> Result<Self, ProxyProtocolError> {
66        Self::accept_with_policy(io, HeaderPolicy::Required).await
67    }
68
69    /// Reads and consumes an optional PROXY protocol header from `io`.
70    pub async fn accept_optional(io: IO) -> Result<Self, ProxyProtocolError> {
71        Self::accept_with_policy(io, HeaderPolicy::Optional).await
72    }
73
74    /// Reads and consumes a PROXY protocol header according to `policy`.
75    pub async fn accept_with_policy(
76        mut io: IO,
77        policy: HeaderPolicy,
78    ) -> Result<Self, ProxyProtocolError> {
79        let mut prefix = Vec::with_capacity(v2::SIGNATURE.len());
80
81        loop {
82            let Some(byte) = read_byte(&mut io).await? else {
83                return if policy.is_required() {
84                    Err(ProxyProtocolError::MissingHeader)
85                } else {
86                    Ok(Self {
87                        io,
88                        header: None,
89                        pending: prefix,
90                    })
91                };
92            };
93
94            prefix.push(byte);
95
96            if prefix == v1::SIGNATURE.as_bytes() {
97                let header = read_v1_header(&mut io, prefix).await?;
98                return Ok(Self {
99                    io,
100                    header: Some(Header::V1(header)),
101                    pending: Vec::new(),
102                });
103            }
104
105            if prefix == v2::SIGNATURE {
106                let header = read_v2_header(&mut io, prefix).await?;
107                return Ok(Self {
108                    io,
109                    header: Some(Header::V2(header)),
110                    pending: Vec::new(),
111                });
112            }
113
114            let could_be_v1 = v1::SIGNATURE.as_bytes().starts_with(&prefix);
115            let could_be_v2 = v2::SIGNATURE.starts_with(&prefix);
116
117            if !could_be_v1 && !could_be_v2 {
118                return if policy.is_required() {
119                    Err(ProxyProtocolError::MissingHeader)
120                } else {
121                    Ok(Self {
122                        io,
123                        header: None,
124                        pending: prefix,
125                    })
126                };
127            }
128        }
129    }
130}
131
132impl<IO> AsyncRead for ProxyStream<IO>
133where
134    IO: AsyncRead,
135{
136    fn poll_read(
137        self: Pin<&mut Self>,
138        cx: &mut Context<'_>,
139        buf: &mut ReadBuf<'_>,
140    ) -> Poll<io::Result<()>> {
141        let mut this = self.project();
142
143        if !this.pending.is_empty() {
144            let len = buf.remaining().min(this.pending.len());
145            buf.put_slice(&this.pending[..len]);
146            this.pending.drain(..len);
147            return Poll::Ready(Ok(()));
148        }
149
150        this.io.as_mut().poll_read(cx, buf)
151    }
152}
153
154impl<IO> AsyncWrite for ProxyStream<IO>
155where
156    IO: AsyncWrite,
157{
158    fn poll_write(
159        self: Pin<&mut Self>,
160        cx: &mut Context<'_>,
161        buf: &[u8],
162    ) -> Poll<io::Result<usize>> {
163        self.project().io.poll_write(cx, buf)
164    }
165
166    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
167        self.project().io.poll_flush(cx)
168    }
169
170    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
171        self.project().io.poll_shutdown(cx)
172    }
173
174    fn poll_write_vectored(
175        self: Pin<&mut Self>,
176        cx: &mut Context<'_>,
177        bufs: &[io::IoSlice<'_>],
178    ) -> Poll<io::Result<usize>> {
179        self.project().io.poll_write_vectored(cx, bufs)
180    }
181
182    fn is_write_vectored(&self) -> bool {
183        self.io.is_write_vectored()
184    }
185}
186
187impl<IO> ActixStream for ProxyStream<IO>
188where
189    IO: ActixStream,
190{
191    fn poll_read_ready(&self, cx: &mut Context<'_>) -> Poll<io::Result<Ready>> {
192        if !self.pending.is_empty() {
193            return Poll::Ready(Ok(Ready::READABLE));
194        }
195
196        self.io.poll_read_ready(cx)
197    }
198
199    fn poll_write_ready(&self, cx: &mut Context<'_>) -> Poll<io::Result<Ready>> {
200        self.io.poll_write_ready(cx)
201    }
202}
203
204async fn read_v1_header<IO>(io: &mut IO, mut bytes: Vec<u8>) -> Result<v1::Header, ParseError>
205where
206    IO: AsyncRead + Unpin,
207{
208    while !bytes.ends_with(b"\r\n") {
209        if bytes.len() == v1::MAX_HEADER_SIZE {
210            return Err(ParseError::invalid(
211                "PROXY v1 header exceeds maximum length",
212            ));
213        }
214
215        let Some(byte) = read_byte(io).await? else {
216            return Err(ParseError::invalid("stream ended inside PROXY v1 header"));
217        };
218
219        bytes.push(byte);
220    }
221
222    v1::Header::try_from_bytes(&bytes)
223        .map(|(_, header)| header)
224        .map_err(|_| ParseError::invalid("invalid PROXY v1 header"))
225}
226
227async fn read_v2_header<IO>(io: &mut IO, mut bytes: Vec<u8>) -> Result<v2::Header, ParseError>
228where
229    IO: AsyncRead + Unpin,
230{
231    let mut fixed = [0; 4];
232    io.read_exact(&mut fixed).await?;
233    bytes.extend_from_slice(&fixed);
234
235    let len = u16::from_be_bytes([fixed[2], fixed[3]]) as usize;
236    let mut payload = vec![0; len];
237    io.read_exact(&mut payload).await?;
238    bytes.extend_from_slice(&payload);
239
240    v2::Header::try_from_bytes(&bytes).map(|(_, header)| header)
241}
242
243async fn read_byte<IO>(io: &mut IO) -> io::Result<Option<u8>>
244where
245    IO: AsyncRead + Unpin,
246{
247    let mut byte = [0];
248
249    match io.read(&mut byte).await? {
250        0 => Ok(None),
251        1 => Ok(Some(byte[0])),
252        _ => unreachable!("read with one-byte buffer returned more than one byte"),
253    }
254}
255
256#[cfg(test)]
257mod tests {
258    use std::net::SocketAddr;
259
260    use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _, duplex};
261
262    use super::*;
263    use crate::{AddressFamily, Command, TransportProtocol};
264
265    #[actix_rt::test]
266    async fn consumes_v1_header_and_preserves_stream_body() {
267        let (mut client, server) = duplex(1024);
268        let header = v1::Header::new_inet(
269            SocketAddr::from(([192, 0, 2, 1], 12345)),
270            SocketAddr::from(([198, 51, 100, 2], 443)),
271        );
272
273        actix_rt::spawn(async move {
274            header.write_to_tokio(&mut client).await.unwrap();
275            client.write_all(b"GET / HTTP/1.1\r\n\r\n").await.unwrap();
276            client.shutdown().await.unwrap();
277        });
278
279        let mut stream = ProxyStream::accept(server).await.unwrap();
280
281        assert_eq!(
282            stream.header().unwrap().source_addr().unwrap(),
283            SocketAddr::from(([192, 0, 2, 1], 12345))
284        );
285
286        let mut body = String::new();
287        stream.read_to_string(&mut body).await.unwrap();
288        assert_eq!(body, "GET / HTTP/1.1\r\n\r\n");
289    }
290
291    #[actix_rt::test]
292    async fn consumes_v2_header_and_preserves_stream_body() {
293        let (mut client, server) = duplex(1024);
294        let mut header = v2::Header::new(
295            Command::Proxy,
296            TransportProtocol::Stream,
297            AddressFamily::Inet,
298            SocketAddr::from(([192, 0, 2, 1], 12345)),
299            SocketAddr::from(([198, 51, 100, 2], 443)),
300        );
301        header.add_tlv(0x05, b"abc123");
302
303        actix_rt::spawn(async move {
304            header.write_to_tokio(&mut client).await.unwrap();
305            client.write_all(b"payload").await.unwrap();
306            client.shutdown().await.unwrap();
307        });
308
309        let mut stream = ProxyStream::accept(server).await.unwrap();
310
311        let Header::V2(header) = stream.header().unwrap() else {
312            panic!("expected v2 header");
313        };
314
315        assert_eq!(
316            header.tlvs().collect::<Vec<_>>(),
317            vec![(0x05, b"abc123".as_slice())]
318        );
319
320        let mut body = String::new();
321        stream.read_to_string(&mut body).await.unwrap();
322        assert_eq!(body, "payload");
323    }
324
325    #[actix_rt::test]
326    async fn optional_mode_replays_bytes_when_no_header_is_present() {
327        let (mut client, server) = duplex(1024);
328
329        actix_rt::spawn(async move {
330            client.write_all(b"GET / HTTP/1.1\r\n\r\n").await.unwrap();
331            client.shutdown().await.unwrap();
332        });
333
334        let mut stream = ProxyStream::accept_optional(server).await.unwrap();
335
336        assert!(stream.header().is_none());
337
338        let mut body = String::new();
339        stream.read_to_string(&mut body).await.unwrap();
340        assert_eq!(body, "GET / HTTP/1.1\r\n\r\n");
341    }
342}
actix_proxy_protocol/service/stream.rs

actix_proxy_protocol/service/
stream.rs
