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

rama_haproxy/client/
layer.rs
