ngrok/
proxy_proto.rs

1use std::{
2    io,
3    mem,
4    pin::Pin,
5    task::{
6        Context,
7        Poll,
8        ready,
9    },
10};
11
12use bytes::{
13    Buf,
14    BytesMut,
15};
16use proxy_protocol::{
17    ParseError,
18    ProxyHeader,
19};
20use tokio::io::{
21    AsyncRead,
22    AsyncWrite,
23    ReadBuf,
24};
25use tracing::instrument;
26
27// 536 is the smallest possible TCP segment, which both v1 and v2 are guaranteed
28// to fit into.
29const MAX_HEADER_LEN: usize = 536;
30// v2 headers start with at least 16 bytes
31const MIN_HEADER_LEN: usize = 16;
32
33#[derive(Debug)]
34enum ReadState {
35    Reading(Option<ParseError>, BytesMut),
36    Error(proxy_protocol::ParseError, BytesMut),
37    Header(Option<proxy_protocol::ProxyHeader>, BytesMut),
38    None,
39}
40
41impl ReadState {
42    fn new() -> ReadState {
43        ReadState::Reading(None, BytesMut::with_capacity(MAX_HEADER_LEN))
44    }
45
46    fn header(&self) -> Result<Option<&ProxyHeader>, &ParseError> {
47        match self {
48            ReadState::Error(err, _) | ReadState::Reading(Some(err), _) => Err(err),
49            ReadState::None | ReadState::Reading(None, _) => Ok(None),
50            ReadState::Header(hdr, _) => Ok(hdr.as_ref()),
51        }
52    }
53
54    /// Read the header from the stream *once*. Once a header has been read, or
55    /// it's been determined that no header is coming, this will be a no-op.
56    #[instrument(level = "trace", skip(reader))]
57    fn poll_read_header_once(
58        &mut self,
59        cx: &mut Context,
60        mut reader: Pin<&mut impl AsyncRead>,
61    ) -> Poll<io::Result<()>> {
62        loop {
63            let read_state = mem::replace(self, ReadState::None);
64            let (last_err, mut hdr_buf) = match read_state {
65                // End states
66                ReadState::None | ReadState::Header(_, _) | ReadState::Error(_, _) => {
67                    *self = read_state;
68                    return Poll::Ready(Ok(()));
69                }
70                ReadState::Reading(err, hdr_buf) => (err, hdr_buf),
71            };
72
73            if hdr_buf.len() < MAX_HEADER_LEN {
74                let mut tmp_buf = ReadBuf::uninit(hdr_buf.spare_capacity_mut());
75                let read_res = reader.as_mut().poll_read(cx, &mut tmp_buf);
76                // Regardless of error, make sure we track the read bytes
77                let filled = tmp_buf.filled().len();
78                if filled > 0 {
79                    let len = hdr_buf.len();
80                    // Safety: the tmp_buf is backed by the uninitialized
81                    // portion of hdr_buf. Advancing the len to len + filled is
82                    // guaranteed to only cover the bytes initialized by the
83                    // read.
84                    unsafe { hdr_buf.set_len(len + filled) }
85                }
86                match read_res {
87                    // If we hit the end of the stream due to either an EOF or
88                    // an error, set the state to a terminal one and return the
89                    // result.
90                    Poll::Ready(ref res) if res.is_err() || filled == 0 => {
91                        *self = match last_err {
92                            Some(err) => ReadState::Error(err, hdr_buf),
93                            None => ReadState::Header(None, hdr_buf),
94                        };
95                        return read_res;
96                    }
97                    // Pending leaves the last error and buffer unchanged.
98                    Poll::Pending => {
99                        *self = ReadState::Reading(last_err, hdr_buf);
100                        return read_res;
101                    }
102                    _ => {}
103                }
104            }
105
106            // Create a view into the header buffer so that failed parse
107            // attempts don't consume it.
108            let mut hdr_view = &*hdr_buf;
109
110            // Don't try to parse unless we have a minimum number of bytes to
111            // avoid spurious "NotProxyHeader" errors.
112            // Also hack around a bug in the proxy_protocol crate that results
113            // in panics when the input ends in \r without the \n.
114            if hdr_view.len() < MIN_HEADER_LEN || matches!(hdr_view.last(), Some(b'\r')) {
115                *self = ReadState::Reading(last_err, hdr_buf);
116                continue;
117            }
118
119            match proxy_protocol::parse(&mut hdr_view) {
120                Ok(hdr) => {
121                    hdr_buf.advance(hdr_buf.len() - hdr_view.len());
122                    *self = ReadState::Header(Some(hdr), hdr_buf);
123                    return Poll::Ready(Ok(()));
124                }
125                Err(ParseError::NotProxyHeader) => {
126                    *self = ReadState::Header(None, hdr_buf);
127                    return Poll::Ready(Ok(()));
128                }
129
130                // Keep track of the last error - it might not be fatal if we
131                // simply haven't read enough
132                Err(err) => {
133                    // If we've read too much, consider the error fatal.
134                    if hdr_buf.len() >= MAX_HEADER_LEN {
135                        *self = ReadState::Error(err, hdr_buf);
136                    } else {
137                        *self = ReadState::Reading(Some(err), hdr_buf);
138                    }
139                    continue;
140                }
141            }
142        }
143    }
144}
145
146#[derive(Debug)]
147enum WriteState {
148    Writing(BytesMut),
149    None,
150}
151
152impl WriteState {
153    fn new(hdr: proxy_protocol::ProxyHeader) -> Result<WriteState, proxy_protocol::EncodeError> {
154        proxy_protocol::encode(hdr).map(WriteState::Writing)
155    }
156
157    /// Write the header *once*. After its written to the stream, this will be a
158    /// no-op.
159    #[instrument(level = "trace", skip(writer))]
160    fn poll_write_header_once(
161        &mut self,
162        cx: &mut Context,
163        mut writer: Pin<&mut impl AsyncWrite>,
164    ) -> Poll<io::Result<()>> {
165        loop {
166            let state = mem::replace(self, WriteState::None);
167            match state {
168                WriteState::None => return Poll::Ready(Ok(())),
169                WriteState::Writing(mut buf) => {
170                    let write_res = writer.as_mut().poll_write(cx, &buf);
171                    match write_res {
172                        Poll::Pending | Poll::Ready(Err(_)) => {
173                            *self = WriteState::Writing(buf);
174                            ready!(write_res)?;
175                            unreachable!(
176                                "ready! will return for us on either Pending or Ready(Err)"
177                            );
178                        }
179                        Poll::Ready(Ok(written)) => {
180                            buf.advance(written);
181                            if !buf.is_empty() {
182                                *self = WriteState::Writing(buf);
183                                continue;
184                            } else {
185                                return Ok(()).into();
186                            }
187                        }
188                    }
189                }
190            }
191        }
192    }
193}
194
195#[derive(Debug)]
196#[pin_project::pin_project]
197pub struct Stream<S> {
198    read_state: ReadState,
199    write_state: WriteState,
200    #[pin]
201    inner: S,
202}
203
204impl<S> Stream<S> {
205    pub fn outgoing(stream: S, header: ProxyHeader) -> Result<Self, proxy_protocol::EncodeError> {
206        Ok(Stream {
207            inner: stream,
208            write_state: WriteState::new(header)?,
209            read_state: ReadState::None,
210        })
211    }
212
213    pub fn incoming(stream: S) -> Self {
214        Stream {
215            inner: stream,
216            read_state: ReadState::new(),
217            write_state: WriteState::None,
218        }
219    }
220
221    pub fn disabled(stream: S) -> Self {
222        Stream {
223            inner: stream,
224            read_state: ReadState::None,
225            write_state: WriteState::None,
226        }
227    }
228}
229
230impl<S> Stream<S>
231where
232    S: AsyncRead,
233{
234    #[instrument(level = "debug", skip(self))]
235    pub async fn proxy_header(&mut self) -> io::Result<Result<Option<&ProxyHeader>, &ParseError>>
236    where
237        Self: Unpin,
238    {
239        let mut this = Pin::new(self);
240
241        futures::future::poll_fn(|cx| {
242            let this = this.as_mut().project();
243            this.read_state.poll_read_header_once(cx, this.inner)
244        })
245        .await?;
246
247        Ok(this.get_mut().read_state.header())
248    }
249}
250
251impl<S> AsyncRead for Stream<S>
252where
253    S: AsyncRead,
254{
255    #[instrument(level = "trace", skip(self), fields(read_state = ?self.read_state))]
256    fn poll_read(
257        self: Pin<&mut Self>,
258        cx: &mut Context<'_>,
259        buf: &mut ReadBuf<'_>,
260    ) -> Poll<io::Result<()>> {
261        let mut this = self.project();
262
263        ready!(
264            this.read_state
265                .poll_read_header_once(cx, this.inner.as_mut())
266        )?;
267
268        match this.read_state {
269            ReadState::Error(_, remainder) | ReadState::Header(_, remainder) => {
270                if !remainder.is_empty() {
271                    let available = std::cmp::min(remainder.len(), buf.remaining());
272                    buf.put_slice(&remainder.split_to(available));
273                    // Make sure Ready is returned regardless of inner's state
274                    return Poll::Ready(Ok(()));
275                }
276            }
277            ReadState::None => {}
278            _ => unreachable!(),
279        }
280
281        this.inner.poll_read(cx, buf)
282    }
283}
284
285impl<S> AsyncWrite for Stream<S>
286where
287    S: AsyncWrite,
288{
289    #[instrument(level = "trace", skip(self), fields(write_state = ?self.write_state))]
290    fn poll_write(
291        self: Pin<&mut Self>,
292        cx: &mut Context<'_>,
293        buf: &[u8],
294    ) -> Poll<Result<usize, io::Error>> {
295        let mut this = self.project();
296
297        ready!(
298            this.write_state
299                .poll_write_header_once(cx, this.inner.as_mut())
300        )?;
301
302        this.inner.poll_write(cx, buf)
303    }
304    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
305        self.project().inner.poll_flush(cx)
306    }
307    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
308        self.project().inner.poll_shutdown(cx)
309    }
310}
311
312#[cfg(feature = "hyper")]
313mod hyper {
314    use ::hyper::rt::{
315        Read as HyperRead,
316        Write as HyperWrite,
317    };
318
319    use super::*;
320
321    impl<S> HyperWrite for Stream<S>
322    where
323        S: AsyncWrite,
324    {
325        #[instrument(level = "trace", skip(self), fields(write_state = ?self.write_state))]
326        fn poll_write(
327            self: Pin<&mut Self>,
328            cx: &mut Context<'_>,
329            buf: &[u8],
330        ) -> Poll<Result<usize, io::Error>> {
331            <Self as AsyncWrite>::poll_write(self, cx, buf)
332        }
333        fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
334            <Self as AsyncWrite>::poll_flush(self, cx)
335        }
336        fn poll_shutdown(
337            self: Pin<&mut Self>,
338            cx: &mut Context<'_>,
339        ) -> Poll<Result<(), io::Error>> {
340            <Self as AsyncWrite>::poll_shutdown(self, cx)
341        }
342    }
343
344    impl<S> HyperRead for Stream<S>
345    where
346        S: AsyncRead,
347    {
348        fn poll_read(
349            self: Pin<&mut Self>,
350            cx: &mut Context<'_>,
351            mut buf: ::hyper::rt::ReadBufCursor<'_>,
352        ) -> Poll<Result<(), std::io::Error>> {
353            let mut tokio_buf = tokio::io::ReadBuf::uninit(unsafe { buf.as_mut() });
354            let res = ready!(<Self as AsyncRead>::poll_read(self, cx, &mut tokio_buf));
355            let filled = tokio_buf.filled().len();
356            unsafe { buf.advance(filled) };
357            Poll::Ready(res)
358        }
359    }
360}
361
362#[cfg(test)]
363mod test {
364    use std::{
365        cmp,
366        io,
367        pin::Pin,
368        task::{
369            Context,
370            Poll,
371            ready,
372        },
373        time::Duration,
374    };
375
376    use bytes::{
377        BufMut,
378        BytesMut,
379    };
380    use proxy_protocol::{
381        ProxyHeader,
382        version2::{
383            self,
384            ProxyCommand,
385        },
386    };
387    use tokio::io::{
388        AsyncRead,
389        AsyncReadExt,
390        AsyncWriteExt,
391        ReadBuf,
392    };
393
394    use super::Stream;
395
396    #[pin_project::pin_project]
397    struct ShortReader<S> {
398        #[pin]
399        inner: S,
400        min: usize,
401        max: usize,
402    }
403
404    impl<S> AsyncRead for ShortReader<S>
405    where
406        S: AsyncRead,
407    {
408        fn poll_read(
409            self: Pin<&mut Self>,
410            cx: &mut Context<'_>,
411            buf: &mut ReadBuf<'_>,
412        ) -> Poll<io::Result<()>> {
413            let mut this = self.project();
414            let max_bytes =
415                *this.min + cmp::max(1, rand::random::<usize>() % (*this.max - *this.min));
416            let mut tmp = vec![0; max_bytes];
417            let mut tmp_buf = ReadBuf::new(&mut tmp);
418            let res = ready!(this.inner.as_mut().poll_read(cx, &mut tmp_buf));
419
420            buf.put_slice(tmp_buf.filled());
421
422            res?;
423
424            Poll::Ready(Ok(()))
425        }
426    }
427
428    impl<S> ShortReader<S> {
429        fn new(inner: S, min: usize, max: usize) -> Self {
430            ShortReader { inner, min, max }
431        }
432    }
433
434    const INPUT: &str = "PROXY TCP4 192.168.0.1 192.168.0.11 56324 443\r\n";
435    const PARTIAL_INPUT: &str = "PROXY TCP4 192.168.0.1";
436    const FINAL_INPUT: &str = " 192.168.0.11 56324 443\r\n";
437
438    // Smoke test to ensure that the proxy protocol parser works as expected.
439    // Not actually testing our code.
440    #[test]
441    fn test_proxy_protocol() {
442        let mut buf = BytesMut::from(INPUT);
443
444        assert!(proxy_protocol::parse(&mut buf).is_ok());
445
446        buf = BytesMut::from(PARTIAL_INPUT);
447
448        assert!(proxy_protocol::parse(&mut &*buf).is_err());
449
450        buf.put_slice(FINAL_INPUT.as_bytes());
451
452        assert!(proxy_protocol::parse(&mut &*buf).is_ok());
453    }
454
455    #[tokio::test]
456    #[tracing_test::traced_test]
457    async fn test_header_stream_v2() {
458        let (left, mut right) = tokio::io::duplex(1024);
459
460        let header = ProxyHeader::Version2 {
461            command: ProxyCommand::Proxy,
462            transport_protocol: version2::ProxyTransportProtocol::Stream,
463            addresses: version2::ProxyAddresses::Ipv4 {
464                source: "127.0.0.1:1".parse().unwrap(),
465                destination: "127.0.0.2:2".parse().unwrap(),
466            },
467        };
468
469        let input = proxy_protocol::encode(header).unwrap();
470
471        let mut proxy_stream = Stream::incoming(ShortReader::new(left, 2, 5));
472
473        // Chunk our writes to ensure that our reader is resilient across split inputs.
474        tokio::spawn(async move {
475            tokio::time::sleep(Duration::from_millis(50)).await;
476
477            right.write_all(&input).await.expect("write header");
478
479            right
480                .write_all(b"Hello, world!")
481                .await
482                .expect("write hello");
483
484            right.shutdown().await.expect("shutdown");
485        });
486
487        let hdr = proxy_stream
488            .proxy_header()
489            .await
490            .expect("read header")
491            .expect("decode header")
492            .expect("header exists");
493
494        assert!(matches!(hdr, ProxyHeader::Version2 { .. }));
495
496        let mut buf = String::new();
497
498        proxy_stream
499            .read_to_string(&mut buf)
500            .await
501            .expect("read rest");
502
503        assert_eq!(buf, "Hello, world!");
504
505        // Get the header again - should be the same.
506        let hdr = proxy_stream
507            .proxy_header()
508            .await
509            .expect("read header")
510            .expect("decode header")
511            .expect("header exists");
512
513        assert!(matches!(hdr, ProxyHeader::Version2 { .. }));
514    }
515
516    #[tokio::test]
517    #[tracing_test::traced_test]
518    async fn test_header_stream() {
519        let (left, mut right) = tokio::io::duplex(1024);
520
521        let mut proxy_stream = Stream::incoming(ShortReader::new(left, 2, 5));
522
523        // Chunk our writes to ensure that our reader is resilient across split inputs.
524        tokio::spawn(async move {
525            tokio::time::sleep(Duration::from_millis(50)).await;
526
527            right
528                .write_all(INPUT.as_bytes())
529                .await
530                .expect("write header");
531
532            right
533                .write_all(b"Hello, world!")
534                .await
535                .expect("write hello");
536
537            right.shutdown().await.expect("shutdown");
538        });
539
540        let hdr = proxy_stream
541            .proxy_header()
542            .await
543            .expect("read header")
544            .expect("decode header")
545            .expect("header exists");
546
547        assert!(matches!(hdr, ProxyHeader::Version1 { .. }));
548
549        let mut buf = String::new();
550
551        proxy_stream
552            .read_to_string(&mut buf)
553            .await
554            .expect("read rest");
555
556        assert_eq!(buf, "Hello, world!");
557
558        // Get the header again - should be the same.
559        let hdr = proxy_stream
560            .proxy_header()
561            .await
562            .expect("read header")
563            .expect("decode header")
564            .expect("header exists");
565
566        assert!(matches!(hdr, ProxyHeader::Version1 { .. }));
567    }
568
569    #[tokio::test]
570    #[tracing_test::traced_test]
571    async fn test_noheader() {
572        let (left, mut right) = tokio::io::duplex(1024);
573
574        let mut proxy_stream = Stream::incoming(left);
575
576        right
577            .write_all(b"Hello, world!")
578            .await
579            .expect("write stream");
580
581        right.shutdown().await.expect("shutdown");
582        drop(right);
583
584        assert!(
585            proxy_stream
586                .proxy_header()
587                .await
588                .unwrap()
589                .unwrap()
590                .is_none()
591        );
592
593        let mut buf = String::new();
594
595        proxy_stream
596            .read_to_string(&mut buf)
597            .await
598            .expect("read stream");
599
600        assert_eq!(buf, "Hello, world!");
601    }
602}
ngrok/proxy_proto.rs

ngrok/
proxy_proto.rs
