rama_haproxy/client/
layer.rs

1use rama_core::error::BoxErrorExt as _;
2use std::{fmt, marker::PhantomData, net::IpAddr};
3
4use crate::protocol::{v1, v2};
5use rama_core::{Layer, Service, bytes::Bytes, error::BoxError, extensions::ExtensionsRef, io::Io};
6use rama_net::{
7    client::{ConnectionError, ConnectionErrorKind, ConnectorService, EstablishedClientConnection},
8    forwarded::Forwarded,
9    stream::{Socket, SocketInfo},
10};
11use tokio::io::AsyncWriteExt;
12
13/// Layer to encode and write the HaProxy Protocol,
14/// as a client on the connected stream.
15///
16/// This connector should in most cases
17/// happen as the first thing after establishing the connection.
18#[derive(Debug, Clone)]
19pub struct HaProxyLayer<P = protocol::Tcp, V = version::Two> {
20    version: V,
21    _phantom: PhantomData<fn(P)>,
22}
23
24impl HaProxyLayer {
25    /// Create a new [`HaProxyLayer`] for the TCP protocol (default).
26    ///
27    /// This is in the PROXY spec referred to as:
28    ///
29    /// - TCP4 (for IPv4, v1)
30    /// - TCP6 (for IPv6, v1)
31    /// - Stream (v2)
32    #[must_use]
33    pub fn tcp() -> Self {
34        Self {
35            version: Default::default(),
36            _phantom: PhantomData,
37        }
38    }
39
40    /// Use version one of PROXY protocol, instead of the
41    /// default version two.
42    ///
43    /// Version one makes use of a less advanced text protocol,
44    /// instead the more advanced binary v2 protocol.
45    ///
46    /// Use this only if you have no control over a v1-only server.
47    pub fn v1(self) -> HaProxyLayer<protocol::Tcp, version::One> {
48        HaProxyLayer {
49            version: Default::default(),
50            _phantom: PhantomData,
51        }
52    }
53}
54
55impl HaProxyLayer<protocol::Udp> {
56    /// Create a new [`HaProxyLayer`] for the UDP protocol,
57    /// instead of the default TCP protocol.
58    ///
59    /// This is in the PROXY spec referred to as:
60    ///
61    /// - Datagram (v2)
62    #[must_use]
63    pub fn udp() -> Self {
64        Self {
65            version: Default::default(),
66            _phantom: PhantomData,
67        }
68    }
69}
70
71impl<P> HaProxyLayer<P> {
72    rama_utils::macros::generate_set_and_with! {
73        /// Attach a raw bytes payload to the PROXY v2 header, written
74        /// verbatim after the address block.
75        ///
76        /// **Wire-format hazard.** Spec section 2.2 says everything after the
77        /// address block is a sequence of TLVs. If your bytes are not valid
78        /// TLV encoding, conforming receivers (including rama's server) will
79        /// either reject the connection or mis-parse the data. Prefer
80        /// [`Self::tlv`] for anything spec-compliant; treat `payload` as a
81        /// raw escape hatch for testing.
82        ///
83        /// In particular, **combining `payload` with [`Self::crc32c`] is
84        /// rejected at send time**: the CRC32C TLV would be appended after
85        /// the raw payload, which means the receiver would try to interpret
86        /// the payload as TLVs, fail, and the CRC would be computed over a
87        /// header the receiver can't validate. Use TLVs instead.
88        ///
89        /// NOTE this is only possible in Version two of the PROXY Protocol.
90        /// In case you downgrade this [`HaProxyLayer`] to version one later
91        /// using [`Self::v1`] this payload will be dropped.
92        pub fn payload(mut self, payload: impl Into<Bytes>) -> Self {
93            self.version.payload = Some(payload.into());
94            self
95        }
96    }
97
98    rama_utils::macros::generate_set_and_with! {
99        /// Attach a Type-Length-Value entry to the emitted v2 header.
100        ///
101        /// Call this once per TLV; entries are written in insertion order.
102        /// Queuing a [`v2::Type::CRC32C`] entry here is rejected at send
103        /// time — a CRC value must be computed over the final header bytes,
104        /// so use [`Self::crc32c`] instead.
105        ///
106        /// NOTE this is only possible in Version two of the PROXY Protocol.
107        /// On downgrade to v1 via [`Self::v1`] all TLVs are dropped.
108        pub fn tlv(mut self, kind: v2::Type, value: impl Into<Bytes>) -> Self {
109            self.version.tlvs.push((kind, value.into()));
110            self
111        }
112    }
113
114    rama_utils::macros::generate_set_and_with! {
115        /// Toggle automatic emission of a valid `PP2_TYPE_CRC32C` TLV.
116        ///
117        /// When enabled, the CRC is computed over the final header (with the
118        /// CRC field zeroed during computation, per spec section 2.2.5) and
119        /// appended as the last TLV. Receivers that enforce CRC verification
120        /// (rama's server default) will accept the header.
121        ///
122        /// Incompatible with [`Self::payload`]: setting both is rejected at
123        /// send time because the raw payload would sit between the TLVs and
124        /// the CRC32C TLV, which standards-conforming receivers cannot parse.
125        ///
126        /// NOTE this is only possible in Version two of the PROXY Protocol.
127        pub fn crc32c(mut self, enabled: bool) -> Self {
128            self.version.crc32c = enabled;
129            self
130        }
131    }
132}
133
134impl<S, P, V: Clone> Layer<S> for HaProxyLayer<P, V> {
135    type Service = HaProxyService<S, P, V>;
136
137    fn layer(&self, inner: S) -> Self::Service {
138        HaProxyService {
139            inner,
140            version: self.version.clone(),
141            _phantom: PhantomData,
142        }
143    }
144
145    fn into_layer(self, inner: S) -> Self::Service {
146        HaProxyService {
147            inner,
148            version: self.version,
149            _phantom: PhantomData,
150        }
151    }
152}
153
154/// Service to encode and write the HaProxy Protocol
155/// as a client on the connected stream.
156///
157/// This connector should in most cases
158/// happen as the first thing after establishing the connection.
159pub struct HaProxyService<S, P = protocol::Tcp, V = version::Two> {
160    inner: S,
161    version: V,
162    _phantom: PhantomData<fn(P)>,
163}
164
165impl<S> HaProxyService<S> {
166    /// Create a new [`HaProxyService`] for the TCP protocol (default).
167    ///
168    /// This is in the PROXY spec referred to as:
169    ///
170    /// - TCP4 (for IPv4, v1)
171    /// - TCP6 (for IPv6, v1)
172    /// - Stream (v2)
173    pub fn tcp(inner: S) -> Self {
174        Self {
175            inner,
176            version: Default::default(),
177            _phantom: PhantomData,
178        }
179    }
180
181    /// Use version one of PROXY protocol, instead of the
182    /// default version two.
183    ///
184    /// Version one makes use of a less advanced text protocol,
185    /// instead the more advanced binary v2 protocol.
186    ///
187    /// Use this only if you have no control over a v1-only server.
188    pub fn v1(self) -> HaProxyService<S, protocol::Tcp, version::One> {
189        HaProxyService {
190            inner: self.inner,
191            version: Default::default(),
192            _phantom: PhantomData,
193        }
194    }
195}
196
197impl<S> HaProxyService<S, protocol::Udp> {
198    /// Create a new [`HaProxyService`] for the UDP protocol,
199    /// instead of the default TCP protocol.
200    ///
201    /// This is in the PROXY spec referred to as:
202    ///
203    /// - Datagram (v2)
204    pub fn udp(inner: S) -> Self {
205        Self {
206            inner,
207            version: Default::default(),
208            _phantom: PhantomData,
209        }
210    }
211}
212
213impl<S, P> HaProxyService<S, P> {
214    rama_utils::macros::generate_set_and_with! {
215        /// Attach a custom bytes payload to the PROXY header.
216        ///
217        /// NOTE this is only possible in Version two of the PROXY Protocol.
218        /// In case you downgrade this [`HaProxyLayer`] to version one later
219        /// using [`Self::v1`] this payload will be dropped.
220        pub fn payload(mut self, payload: impl Into<Bytes>) -> Self {
221            self.version.payload = Some(payload.into());
222            self
223        }
224    }
225
226    rama_utils::macros::generate_set_and_with! {
227        /// Attach a Type-Length-Value entry to the emitted v2 header.
228        /// See [`HaProxyLayer::with_tlv`] for details.
229        pub fn tlv(mut self, kind: v2::Type, value: impl Into<Bytes>) -> Self {
230            self.version.tlvs.push((kind, value.into()));
231            self
232        }
233    }
234
235    rama_utils::macros::generate_set_and_with! {
236        /// Toggle automatic emission of a valid `PP2_TYPE_CRC32C` TLV.
237        /// See [`HaProxyLayer::crc32c`] for details.
238        pub fn crc32c(mut self, enabled: bool) -> Self {
239            self.version.crc32c = enabled;
240            self
241        }
242    }
243}
244
245impl<S: fmt::Debug, P, V: fmt::Debug> fmt::Debug for HaProxyService<S, P, V> {
246    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
247        f.debug_struct("HaProxyService")
248            .field("inner", &self.inner)
249            .field("version", &self.version)
250            .field(
251                "_phantom",
252                &format_args!("{}", std::any::type_name::<fn(P)>()),
253            )
254            .finish()
255    }
256}
257
258impl<S: Clone, P, V: Clone> Clone for HaProxyService<S, P, V> {
259    fn clone(&self) -> Self {
260        Self {
261            inner: self.inner.clone(),
262            version: self.version.clone(),
263            _phantom: PhantomData,
264        }
265    }
266}
267
268impl<S, P, Input> Service<Input> for HaProxyService<S, P, version::One>
269where
270    S: ConnectorService<Input, Connection: Io + Socket + Unpin>,
271    P: Send + 'static,
272    Input: Send + ExtensionsRef + 'static,
273{
274    type Output = EstablishedClientConnection<S::Connection, Input>;
275    type Error = ConnectionError;
276
277    async fn serve(&self, input: Input) -> Result<Self::Output, Self::Error> {
278        let EstablishedClientConnection { input, mut conn } = self.inner.connect(input).await?;
279
280        let src = input
281            .extensions()
282            .clone_to_if_absent::<Forwarded>(conn.extensions())
283            .and_then(|f| f.client_socket_addr())
284            .or_else(|| {
285                input
286                    .extensions()
287                    .clone_to_if_absent::<SocketInfo>(conn.extensions())
288                    .map(|info| info.peer_addr())
289            })
290            .ok_or_else(|| {
291                ConnectionError::local(
292                    BoxError::from_static_str("PROXY client (v1): missing src socket address"),
293                    ConnectionErrorKind::InvalidInput,
294                )
295            })?;
296
297        let peer_addr = conn.peer_addr().map_err(|error| {
298            ConnectionError::local(error, ConnectionErrorKind::Internal)
299                .context("PROXY client (v1): read peer address")
300        })?;
301        let addresses = match (src.ip_addr, peer_addr.ip_addr) {
302            (IpAddr::V4(src_ip), IpAddr::V4(dst_ip)) => {
303                v1::Addresses::new_tcp4(src_ip, dst_ip, src.port, peer_addr.port)
304            }
305            (IpAddr::V6(src_ip), IpAddr::V6(dst_ip)) => {
306                v1::Addresses::new_tcp6(src_ip, dst_ip, src.port, peer_addr.port)
307            }
308            (_, _) => {
309                return Err(ConnectionError::local(
310                    BoxError::from_static_str(
311                        "PROXY client (v1): IP version mismatch between src and dest",
312                    ),
313                    ConnectionErrorKind::InvalidInput,
314                ));
315            }
316        };
317
318        conn.write_all(addresses.to_string().as_bytes())
319            .await
320            .map_err(|error| {
321                ConnectionError::application(error, ConnectionErrorKind::Protocol)
322                    .context("PROXY client (v1): write addresses")
323            })?;
324
325        Ok(EstablishedClientConnection { input, conn })
326    }
327}
328
329impl<S, P, Input> Service<Input> for HaProxyService<S, P, version::Two>
330where
331    S: ConnectorService<Input, Connection: Io + Socket + Unpin>,
332    P: protocol::Protocol + Send + 'static,
333    Input: Send + ExtensionsRef + 'static,
334{
335    type Output = EstablishedClientConnection<S::Connection, Input>;
336    type Error = ConnectionError;
337
338    async fn serve(&self, input: Input) -> Result<Self::Output, Self::Error> {
339        let EstablishedClientConnection { input, mut conn } = self.inner.connect(input).await?;
340
341        let src = {
342            input
343                .extensions()
344                .clone_to_if_absent::<Forwarded>(conn.extensions())
345                .and_then(|f| f.client_socket_addr())
346                .or_else(|| {
347                    input
348                        .extensions()
349                        .clone_to_if_absent::<SocketInfo>(conn.extensions())
350                        .map(|info| info.peer_addr())
351                })
352                .ok_or_else(|| {
353                    ConnectionError::local(
354                        BoxError::from_static_str("PROXY client (v2): missing src socket address"),
355                        ConnectionErrorKind::InvalidInput,
356                    )
357                })?
358        };
359
360        let peer_addr = conn.peer_addr().map_err(|error| {
361            ConnectionError::local(error, ConnectionErrorKind::Internal)
362                .context("PROXY client (v2): read peer address")
363        })?;
364        let builder = match (src.ip_addr, peer_addr.ip_addr) {
365            (IpAddr::V4(src_ip), IpAddr::V4(dst_ip)) => v2::Builder::with_addresses(
366                v2::Version::Two | v2::Command::Proxy,
367                P::v2_protocol(),
368                v2::IPv4::new(src_ip, dst_ip, src.port, peer_addr.port),
369            ),
370            (IpAddr::V6(src_ip), IpAddr::V6(dst_ip)) => v2::Builder::with_addresses(
371                v2::Version::Two | v2::Command::Proxy,
372                P::v2_protocol(),
373                v2::IPv6::new(src_ip, dst_ip, src.port, peer_addr.port),
374            ),
375            (_, _) => {
376                return Err(ConnectionError::local(
377                    BoxError::from_static_str(
378                        "PROXY client (v2): IP version mismatch between src and dest",
379                    ),
380                    ConnectionErrorKind::InvalidInput,
381                ));
382            }
383        };
384
385        if self.version.crc32c && self.version.payload.is_some() {
386            // The raw `payload` bytes sit between the TLVs and the CRC32C TLV
387            // — receivers parse the entire post-address area as TLVs, so the
388            // payload would be mis-interpreted (or rejected) and the CRC
389            // would be computed over a header the peer can't validate.
390            return Err(ConnectionError::local(
391                BoxError::from_static_str(
392                    "PROXY client (v2): `payload` and `crc32c` cannot be combined; \
393                     use typed TLVs instead",
394                ),
395                ConnectionErrorKind::InvalidInput,
396            ));
397        }
398
399        // A CRC32C TLV value is never something a sender should compute by
400        // hand: it has to be computed over the *final* header bytes, which
401        // only the builder can see. If the user manually queued one via
402        // `with_tlv(Type::CRC32C, ...)`, reject — `crc32c(true)` is the
403        // only correct way to emit a CRC32C TLV.
404        if self
405            .version
406            .tlvs
407            .iter()
408            .any(|(kind, _)| *kind == v2::Type::CRC32C)
409        {
410            return Err(ConnectionError::local(
411                BoxError::from_static_str(
412                    "PROXY client (v2): manual `with_tlv(Type::CRC32C, ...)` is not supported; \
413                     use `with_crc32c(true)` instead",
414                ),
415                ConnectionErrorKind::InvalidInput,
416            ));
417        }
418
419        let mut builder = builder;
420        for (kind, value) in &self.version.tlvs {
421            builder = builder.write_tlv(*kind, value.as_ref()).map_err(|error| {
422                ConnectionError::local(error, ConnectionErrorKind::InvalidInput)
423                    .context("PROXY client (v2): write TLV")
424            })?;
425        }
426        if let Some(payload) = self.version.payload.as_deref() {
427            builder = builder.write_payload(payload).map_err(|error| {
428                ConnectionError::local(error, ConnectionErrorKind::InvalidInput)
429                    .context("PROXY client (v2): write custom binary payload to header")
430            })?;
431        }
432        if self.version.crc32c {
433            // Spec section 2.2.5: build() will append the CRC32C TLV last
434            // (covering every other TLV) and patch in the computed value.
435            builder = builder.with_crc32c(true);
436        }
437
438        let header = builder.build().map_err(|error| {
439            ConnectionError::local(error, ConnectionErrorKind::InvalidInput)
440                .context("PROXY client (v2): encode header")
441        })?;
442        conn.write_all(&header[..]).await.map_err(|error| {
443            ConnectionError::application(error, ConnectionErrorKind::Protocol)
444                .context("PROXY client (v2): write header")
445        })?;
446
447        Ok(EstablishedClientConnection { input, conn })
448    }
449}
450
451pub mod version {
452    //! Marker traits for the HaProxy (PROXY) version to be used by client layer (service).
453
454    use super::*;
455
456    #[derive(Debug, Clone, Default)]
457    /// Use version 1 of the PROXY protocol.
458    ///
459    /// See [`crate::protocol`] for more information.
460    #[non_exhaustive]
461    pub struct One;
462
463    #[derive(Debug, Clone, Default)]
464    /// Use version 2 of the PROXY protocol.
465    ///
466    /// See [`crate::protocol`] for more information.
467    pub struct Two {
468        pub(crate) payload: Option<Bytes>,
469        pub(crate) tlvs: Vec<(crate::protocol::v2::Type, Bytes)>,
470        pub(crate) crc32c: bool,
471    }
472}
473
474pub mod protocol {
475    //! Marker traits for the HaProxy (PROXY) protocol to be used by client layer (service).
476
477    use crate::protocol::v2;
478
479    #[derive(Debug, Clone)]
480    /// Encode the data for the TCP protocol (possible in [`super::version::One`] and [`super::version::Two`]).
481    ///
482    /// See [`crate::protocol`] for more information.
483    pub struct Tcp;
484
485    #[derive(Debug, Clone)]
486    /// Encode the data for the UDP protocol (possible only in [`super::version::Two`]).
487    ///
488    /// See [`crate::protocol`] for more information.
489    pub struct Udp;
490
491    pub(super) trait Protocol {
492        /// Return the v2 PROXY protocol linked to the protocol implementation.
493        fn v2_protocol() -> v2::Protocol;
494    }
495
496    impl Protocol for Tcp {
497        fn v2_protocol() -> v2::Protocol {
498            v2::Protocol::Stream
499        }
500    }
501
502    impl Protocol for Udp {
503        fn v2_protocol() -> v2::Protocol {
504            v2::Protocol::Datagram
505        }
506    }
507}
508
509#[cfg(test)]
510mod tests {
511    use super::*;
512    use rama_core::{
513        Layer, ServiceInput, extensions::Extensions, extensions::ExtensionsRef, service::service_fn,
514    };
515    use rama_net::{
516        address::SocketAddress,
517        forwarded::{ForwardedElement, NodeId},
518    };
519    use std::{convert::Infallible, pin::Pin};
520    use tokio::io::{AsyncRead, AsyncWrite};
521    use tokio_test::io::{Builder, Mock};
522
523    struct SocketConnection {
524        conn: Mock,
525        extensions: Extensions,
526        socket: SocketAddress,
527    }
528
529    impl ExtensionsRef for SocketConnection {
530        fn extensions(&self) -> &Extensions {
531            &self.extensions
532        }
533    }
534
535    impl Socket for SocketConnection {
536        fn local_addr(&self) -> std::io::Result<SocketAddress> {
537            Ok(self.socket)
538        }
539
540        fn peer_addr(&self) -> std::io::Result<SocketAddress> {
541            Ok(self.socket)
542        }
543    }
544
545    #[warn(clippy::missing_trait_methods)]
546    impl AsyncWrite for SocketConnection {
547        fn poll_write(
548            mut self: std::pin::Pin<&mut Self>,
549            cx: &mut std::task::Context<'_>,
550            buf: &[u8],
551        ) -> std::task::Poll<Result<usize, std::io::Error>> {
552            Pin::new(&mut self.conn).poll_write(cx, buf)
553        }
554
555        fn poll_flush(
556            mut self: std::pin::Pin<&mut Self>,
557            cx: &mut std::task::Context<'_>,
558        ) -> std::task::Poll<Result<(), std::io::Error>> {
559            Pin::new(&mut self.conn).poll_flush(cx)
560        }
561
562        fn poll_shutdown(
563            mut self: std::pin::Pin<&mut Self>,
564            cx: &mut std::task::Context<'_>,
565        ) -> std::task::Poll<Result<(), std::io::Error>> {
566            Pin::new(&mut self.conn).poll_shutdown(cx)
567        }
568
569        fn is_write_vectored(&self) -> bool {
570            self.conn.is_write_vectored()
571        }
572
573        fn poll_write_vectored(
574            mut self: Pin<&mut Self>,
575            cx: &mut std::task::Context<'_>,
576            bufs: &[std::io::IoSlice<'_>],
577        ) -> std::task::Poll<Result<usize, std::io::Error>> {
578            Pin::new(&mut self.conn).poll_write_vectored(cx, bufs)
579        }
580    }
581
582    #[warn(clippy::missing_trait_methods)]
583    impl AsyncRead for SocketConnection {
584        fn poll_read(
585            mut self: Pin<&mut Self>,
586            cx: &mut std::task::Context<'_>,
587            buf: &mut tokio::io::ReadBuf<'_>,
588        ) -> std::task::Poll<std::io::Result<()>> {
589            Pin::new(&mut self.conn).poll_read(cx, buf)
590        }
591    }
592
593    #[tokio::test]
594    async fn test_v1_tcp() {
595        for (expected_line, ext, target_addr) in [
596            (
597                "PROXY TCP4 127.0.1.2 192.168.1.101 80 443\r\n",
598                {
599                    let ext = Extensions::new();
600                    ext.insert(SocketInfo::new(None, "127.0.1.2:80".parse().unwrap()));
601                    ext
602                },
603                "192.168.1.101:443",
604            ),
605            (
606                "PROXY TCP4 127.0.1.2 192.168.1.101 80 443\r\n",
607                {
608                    let ext = Extensions::new();
609                    ext.insert(SocketInfo::new(
610                        None,
611                        "[1234:5678:90ab:cdef:fedc:ba09:8765:4321]:443"
612                            .parse()
613                            .unwrap(),
614                    ));
615                    ext.insert(Forwarded::new(ForwardedElement::new_forwarded_for(
616                        NodeId::try_from("127.0.1.2:80").unwrap(),
617                    )));
618                    ext
619                },
620                "192.168.1.101:443",
621            ),
622            (
623                "PROXY TCP6 1234:5678:90ab:cdef:fedc:ba09:8765:4321 4321:8765:ba09:fedc:cdef:90ab:5678:1234 443 65535\r\n",
624                {
625                    let ext = Extensions::new();
626                    ext.insert(SocketInfo::new(
627                        None,
628                        "[1234:5678:90ab:cdef:fedc:ba09:8765:4321]:443"
629                            .parse()
630                            .unwrap(),
631                    ));
632                    ext
633                },
634                "[4321:8765:ba09:fedc:cdef:90ab:5678:1234]:65535",
635            ),
636            (
637                "PROXY TCP6 1234:5678:90ab:cdef:fedc:ba09:8765:4321 4321:8765:ba09:fedc:cdef:90ab:5678:1234 443 65535\r\n",
638                {
639                    let ext = Extensions::new();
640                    ext.insert(SocketInfo::new(None, "127.0.1.2:80".parse().unwrap()));
641                    ext.insert(Forwarded::new(ForwardedElement::new_forwarded_for(
642                        NodeId::try_from("[1234:5678:90ab:cdef:fedc:ba09:8765:4321]:443").unwrap(),
643                    )));
644                    ext
645                },
646                "[4321:8765:ba09:fedc:cdef:90ab:5678:1234]:65535",
647            ),
648        ] {
649            let svc =
650                HaProxyLayer::tcp()
651                    .v1()
652                    .layer(service_fn(async move |input: ServiceInput<()>| {
653                        Ok::<_, Infallible>(EstablishedClientConnection {
654                            input,
655                            conn: SocketConnection {
656                                socket: target_addr.parse().unwrap(),
657                                conn: Builder::new().write(expected_line.as_bytes()).build(),
658                                extensions: Extensions::new(),
659                            },
660                        })
661                    }));
662
663            let input = ServiceInput::new(());
664            input.extensions().extend(&ext);
665            svc.serve(input).await.unwrap();
666        }
667    }
668
669    #[tokio::test]
670    async fn test_v1_tcp_ip_version_mismatch() {
671        for (ext, target_addr) in [
672            (
673                {
674                    let ext = Extensions::new();
675                    ext.insert(SocketInfo::new(
676                        None,
677                        "[1234:5678:90ab:cdef:fedc:ba09:8765:4321]:80"
678                            .parse()
679                            .unwrap(),
680                    ));
681                    ext
682                },
683                "192.168.1.101:443",
684            ),
685            (
686                {
687                    let ext = Extensions::new();
688                    ext.insert(SocketInfo::new(None, "127.0.1.2:80".parse().unwrap()));
689                    ext.insert(Forwarded::new(ForwardedElement::new_forwarded_for(
690                        NodeId::try_from("[1234:5678:90ab:cdef:fedc:ba09:8765:4321]:80").unwrap(),
691                    )));
692                    ext
693                },
694                "192.168.1.101:443",
695            ),
696            (
697                {
698                    let ext = Extensions::new();
699                    ext.insert(SocketInfo::new(None, "127.0.1.2:80".parse().unwrap()));
700                    ext
701                },
702                "[4321:8765:ba09:fedc:cdef:90ab:5678:1234]:65535",
703            ),
704            (
705                {
706                    let ext = Extensions::new();
707                    ext.insert(SocketInfo::new(
708                        None,
709                        "[1234:5678:90ab:cdef:fedc:ba09:8765:4321]:80"
710                            .parse()
711                            .unwrap(),
712                    ));
713                    ext.insert(Forwarded::new(ForwardedElement::new_forwarded_for(
714                        NodeId::try_from("127.0.1.2:80").unwrap(),
715                    )));
716                    ext
717                },
718                "[4321:8765:ba09:fedc:cdef:90ab:5678:1234]:65535",
719            ),
720        ] {
721            let svc =
722                HaProxyLayer::tcp()
723                    .v1()
724                    .layer(service_fn(async move |input: ServiceInput<()>| {
725                        Ok::<_, Infallible>(EstablishedClientConnection {
726                            input,
727                            conn: SocketConnection {
728                                socket: target_addr.parse().unwrap(),
729                                conn: Builder::new().build(),
730                                extensions: Extensions::new(),
731                            },
732                        })
733                    }));
734
735            let input = ServiceInput::new(());
736            input.extensions().extend(&ext);
737            assert!(svc.serve(input).await.is_err());
738        }
739    }
740
741    #[tokio::test]
742    async fn test_v1_tcp_missing_src() {
743        for target_addr in [
744            "192.168.1.101:443",
745            "[1234:5678:90ab:cdef:fedc:ba09:8765:4321]:443",
746        ] {
747            let svc = HaProxyLayer::tcp()
748                .v1()
749                .layer(service_fn(async move |input| {
750                    Ok::<_, Infallible>(EstablishedClientConnection {
751                        input,
752                        conn: SocketConnection {
753                            socket: target_addr.parse().unwrap(),
754                            conn: Builder::new().build(),
755                            extensions: Extensions::new(),
756                        },
757                    })
758                }));
759            assert!(svc.serve(ServiceInput::new(())).await.is_err());
760        }
761    }
762
763    #[tokio::test]
764    async fn test_v2_tcp4() {
765        for ext in [
766            {
767                let ext = Extensions::new();
768                ext.insert(SocketInfo::new(None, "127.0.0.1:80".parse().unwrap()));
769                ext
770            },
771            {
772                let ext = Extensions::new();
773                ext.insert(SocketInfo::new(
774                    None,
775                    "[1234:5678:90ab:cdef:fedc:ba09:8765:4321]:443"
776                        .parse()
777                        .unwrap(),
778                ));
779                ext.insert(Forwarded::new(ForwardedElement::new_forwarded_for(
780                    NodeId::try_from("127.0.0.1:80").unwrap(),
781                )));
782                ext
783            },
784        ] {
785            let svc = HaProxyLayer::tcp().with_payload(vec![42]).layer(service_fn(
786                async move |input: ServiceInput<()>| {
787                    Ok::<_, Infallible>(EstablishedClientConnection {
788                        input,
789                        conn: SocketConnection {
790                            socket: "192.168.1.1:443".parse().unwrap(),
791                            extensions: Extensions::new(),
792                            conn: Builder::new()
793                                .write(&[
794                                    b'\r', b'\n', b'\r', b'\n', b'\0', b'\r', b'\n', b'Q', b'U',
795                                    b'I', b'T', b'\n', 0x21, 0x11, 0, 13, 127, 0, 0, 1, 192, 168,
796                                    1, 1, 0, 80, 1, 187, 42,
797                                ])
798                                .build(),
799                        },
800                    })
801                },
802            ));
803
804            let input = ServiceInput::new(());
805            input.extensions().extend(&ext);
806            svc.serve(input).await.unwrap();
807        }
808    }
809
810    #[tokio::test]
811    async fn test_v2_udp4() {
812        for ext in [
813            {
814                let ext = Extensions::new();
815                ext.insert(SocketInfo::new(None, "127.0.0.1:80".parse().unwrap()));
816                ext
817            },
818            {
819                let ext = Extensions::new();
820                ext.insert(SocketInfo::new(
821                    None,
822                    "[1234:5678:90ab:cdef:fedc:ba09:8765:4321]:443"
823                        .parse()
824                        .unwrap(),
825                ));
826                ext.insert(Forwarded::new(ForwardedElement::new_forwarded_for(
827                    NodeId::try_from("127.0.0.1:80").unwrap(),
828                )));
829                ext
830            },
831        ] {
832            let svc = HaProxyLayer::udp().with_payload(vec![42]).layer(service_fn(
833                async move |input: ServiceInput<()>| {
834                    Ok::<_, Infallible>(EstablishedClientConnection {
835                        input,
836                        conn: SocketConnection {
837                            socket: "192.168.1.1:443".parse().unwrap(),
838                            extensions: Extensions::new(),
839                            conn: Builder::new()
840                                .write(&[
841                                    b'\r', b'\n', b'\r', b'\n', b'\0', b'\r', b'\n', b'Q', b'U',
842                                    b'I', b'T', b'\n', 0x21, 0x12, 0, 13, 127, 0, 0, 1, 192, 168,
843                                    1, 1, 0, 80, 1, 187, 42,
844                                ])
845                                .build(),
846                        },
847                    })
848                },
849            ));
850
851            let input = ServiceInput::new(());
852            input.extensions().extend(&ext);
853            svc.serve(input).await.unwrap();
854        }
855    }
856
857    #[tokio::test]
858    async fn test_v2_tcp6() {
859        for ext in [
860            {
861                let ext = Extensions::new();
862                ext.insert(SocketInfo::new(
863                    None,
864                    "[1234:5678:90ab:cdef:fedc:ba09:8765:4321]:80"
865                        .parse()
866                        .unwrap(),
867                ));
868                ext
869            },
870            {
871                let ext = Extensions::new();
872                ext.insert(SocketInfo::new(None, "127.0.0.1:80".parse().unwrap()));
873                ext.insert(Forwarded::new(ForwardedElement::new_forwarded_for(
874                    NodeId::try_from("[1234:5678:90ab:cdef:fedc:ba09:8765:4321]:80").unwrap(),
875                )));
876                ext
877            },
878        ] {
879            let svc = HaProxyLayer::tcp().with_payload(vec![42]).layer(service_fn(
880                async move |input: ServiceInput<()>| {
881                    Ok::<_, Infallible>(EstablishedClientConnection {
882                        input,
883                        conn: SocketConnection {
884                            socket: "[4321:8765:ba09:fedc:cdef:90ab:5678:1234]:443"
885                                .parse()
886                                .unwrap(),
887                            extensions: Extensions::new(),
888                            conn: Builder::new()
889                                .write(&[
890                                    b'\r', b'\n', b'\r', b'\n', b'\0', b'\r', b'\n', b'Q', b'U',
891                                    b'I', b'T', b'\n', 0x21, 0x21, 0, 37, 0x12, 0x34, 0x56, 0x78,
892                                    0x90, 0xab, 0xcd, 0xef, 0xfe, 0xdc, 0xba, 0x09, 0x87, 0x65,
893                                    0x43, 0x21, 0x43, 0x21, 0x87, 0x65, 0xba, 0x09, 0xfe, 0xdc,
894                                    0xcd, 0xef, 0x90, 0xab, 0x56, 0x78, 0x12, 0x34, 0, 80, 1, 187,
895                                    42,
896                                ])
897                                .build(),
898                        },
899                    })
900                },
901            ));
902
903            let input = ServiceInput::new(());
904            input.extensions().extend(&ext);
905            svc.serve(input).await.unwrap();
906        }
907    }
908
909    #[tokio::test]
910    async fn test_v2_udp6() {
911        for ext in [
912            {
913                let ext = Extensions::new();
914                ext.insert(SocketInfo::new(
915                    None,
916                    "[1234:5678:90ab:cdef:fedc:ba09:8765:4321]:80"
917                        .parse()
918                        .unwrap(),
919                ));
920                ext
921            },
922            {
923                let ext = Extensions::new();
924                ext.insert(SocketInfo::new(None, "127.0.0.1:80".parse().unwrap()));
925                ext.insert(Forwarded::new(ForwardedElement::new_forwarded_for(
926                    NodeId::try_from("[1234:5678:90ab:cdef:fedc:ba09:8765:4321]:80").unwrap(),
927                )));
928                ext
929            },
930        ] {
931            let svc = HaProxyLayer::udp().with_payload(vec![42]).layer(service_fn(
932                async move |input: ServiceInput<()>| {
933                    Ok::<_, Infallible>(EstablishedClientConnection {
934                        input,
935                        conn: SocketConnection {
936                            socket: "[4321:8765:ba09:fedc:cdef:90ab:5678:1234]:443"
937                                .parse()
938                                .unwrap(),
939                            extensions: Extensions::new(),
940                            conn: Builder::new()
941                                .write(&[
942                                    b'\r', b'\n', b'\r', b'\n', b'\0', b'\r', b'\n', b'Q', b'U',
943                                    b'I', b'T', b'\n', 0x21, 0x22, 0, 37, 0x12, 0x34, 0x56, 0x78,
944                                    0x90, 0xab, 0xcd, 0xef, 0xfe, 0xdc, 0xba, 0x09, 0x87, 0x65,
945                                    0x43, 0x21, 0x43, 0x21, 0x87, 0x65, 0xba, 0x09, 0xfe, 0xdc,
946                                    0xcd, 0xef, 0x90, 0xab, 0x56, 0x78, 0x12, 0x34, 0, 80, 1, 187,
947                                    42,
948                                ])
949                                .build(),
950                        },
951                    })
952                },
953            ));
954
955            let input = ServiceInput::new(());
956            input.extensions().extend(&ext);
957            svc.serve(input).await.unwrap();
958        }
959    }
960
961    #[tokio::test]
962    async fn test_v2_ip_version_mismatch() {
963        for (ext, target_addr) in [
964            (
965                {
966                    let ext = Extensions::new();
967                    ext.insert(SocketInfo::new(
968                        None,
969                        "[1234:5678:90ab:cdef:fedc:ba09:8765:4321]:80"
970                            .parse()
971                            .unwrap(),
972                    ));
973                    ext
974                },
975                "192.168.1.101:443",
976            ),
977            (
978                {
979                    let ext = Extensions::new();
980                    ext.insert(SocketInfo::new(None, "127.0.1.2:80".parse().unwrap()));
981                    ext.insert(Forwarded::new(ForwardedElement::new_forwarded_for(
982                        NodeId::try_from("[1234:5678:90ab:cdef:fedc:ba09:8765:4321]:80").unwrap(),
983                    )));
984                    ext
985                },
986                "192.168.1.101:443",
987            ),
988            (
989                {
990                    let ext = Extensions::new();
991                    ext.insert(SocketInfo::new(None, "127.0.1.2:80".parse().unwrap()));
992                    ext
993                },
994                "[4321:8765:ba09:fedc:cdef:90ab:5678:1234]:65535",
995            ),
996            (
997                {
998                    let ext = Extensions::new();
999                    ext.insert(SocketInfo::new(
1000                        None,
1001                        "[1234:5678:90ab:cdef:fedc:ba09:8765:4321]:80"
1002                            .parse()
1003                            .unwrap(),
1004                    ));
1005                    ext.insert(Forwarded::new(ForwardedElement::new_forwarded_for(
1006                        NodeId::try_from("127.0.1.2:80").unwrap(),
1007                    )));
1008                    ext
1009                },
1010                "[4321:8765:ba09:fedc:cdef:90ab:5678:1234]:65535",
1011            ),
1012        ] {
1013            // TCP
1014
1015            let svc = HaProxyLayer::tcp().layer(service_fn(async move |input: ServiceInput<()>| {
1016                Ok::<_, Infallible>(EstablishedClientConnection {
1017                    input,
1018                    conn: SocketConnection {
1019                        socket: target_addr.parse().unwrap(),
1020                        extensions: Extensions::new(),
1021                        conn: Builder::new().build(),
1022                    },
1023                })
1024            }));
1025
1026            let input = ServiceInput::new(());
1027            input.extensions().extend(&ext);
1028            assert!(svc.serve(input).await.is_err());
1029
1030            // UDP
1031
1032            let svc = HaProxyLayer::udp().layer(service_fn(async move |input: ServiceInput<()>| {
1033                Ok::<_, Infallible>(EstablishedClientConnection {
1034                    input,
1035                    conn: SocketConnection {
1036                        socket: target_addr.parse().unwrap(),
1037                        extensions: Extensions::new(),
1038                        conn: Builder::new().build(),
1039                    },
1040                })
1041            }));
1042
1043            let input = ServiceInput::new(());
1044            input.extensions().extend(&ext);
1045            assert!(svc.serve(input).await.is_err());
1046        }
1047    }
1048
1049    #[tokio::test]
1050    async fn test_v2_missing_src() {
1051        for target_addr in [
1052            "192.168.1.101:443",
1053            "[1234:5678:90ab:cdef:fedc:ba09:8765:4321]:443",
1054        ] {
1055            // TCP
1056
1057            let svc = HaProxyLayer::tcp().layer(service_fn(async move |input| {
1058                Ok::<_, Infallible>(EstablishedClientConnection {
1059                    input,
1060                    conn: SocketConnection {
1061                        socket: target_addr.parse().unwrap(),
1062                        extensions: Extensions::new(),
1063                        conn: Builder::new().build(),
1064                    },
1065                })
1066            }));
1067            assert!(svc.serve(ServiceInput::new(())).await.is_err());
1068
1069            // UDP
1070
1071            let svc = HaProxyLayer::udp().layer(service_fn(async move |input| {
1072                Ok::<_, Infallible>(EstablishedClientConnection {
1073                    input,
1074                    conn: SocketConnection {
1075                        socket: target_addr.parse().unwrap(),
1076                        extensions: Extensions::new(),
1077                        conn: Builder::new().build(),
1078                    },
1079                })
1080            }));
1081            assert!(svc.serve(ServiceInput::new(())).await.is_err());
1082        }
1083    }
1084}
rama_haproxy/client/layer.rs

rama_haproxy/client/
layer.rs
