playit_agent_core/network/
proxy_protocol.rs

1use std::{io::Write, net::{Ipv4Addr, Ipv6Addr}};
2
3use byteorder::{BigEndian, ReadBytesExt};
4use playit_agent_proto::udp_proto::UdpFlow;
5use tokio::io::{AsyncWrite, AsyncWriteExt};
6
7use crate::utils::ip_bytes::ReadIpBytesExt;
8
9/*
10 DOCS: https://www.haproxy.org/download/1.8/doc/proxy-protocol.txt
11*/
12
13
14#[derive(PartialEq, Eq, Debug)]
15pub enum ProxyProtocolHeader {
16    AfInet {
17        client_ip: Ipv4Addr,
18        proxy_ip: Ipv4Addr,
19        client_port: u16,
20        proxy_port: u16,
21    },
22    AfInet6 {
23        client_ip: Ipv6Addr,
24        proxy_ip: Ipv6Addr,
25        client_port: u16,
26        proxy_port: u16,
27    },
28}
29
30impl ProxyProtocolHeader {
31    pub fn from_udp_flow(flow: &UdpFlow) -> Self {
32        match flow {
33            UdpFlow::V4 { src, dst, .. } => {
34                ProxyProtocolHeader::AfInet {
35                    client_ip: *src.ip(),
36                    proxy_ip: *dst.ip(),
37                    client_port: src.port(),
38                    proxy_port: dst.port(),
39                }
40            }
41            UdpFlow::V6 { src, dst, .. } => {
42                ProxyProtocolHeader::AfInet6 {
43                    client_ip: src.0,
44                    proxy_ip: dst.0,
45                    client_port: src.1,
46                    proxy_port: dst.1,
47                }
48            }
49        }
50    }
51}
52
53impl std::fmt::Display for ProxyProtocolHeader {
54    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
55        match self {
56            Self::AfInet { client_ip, proxy_ip, client_port, proxy_port } => {
57                write!(f, "PROXY TCP4 {client_ip} {proxy_ip} {client_port} {proxy_port}\r\n")
58            }
59            Self::AfInet6 { client_ip, proxy_ip, client_port, proxy_port } => {
60                write!(f, "PROXY TCP6 {client_ip} {proxy_ip} {client_port} {proxy_port}\r\n")
61            }
62        }
63    }
64}
65
66const PROXY_PROTOCOL_V2_HEADER: &[u8] = &[
67    0x0D, 0x0A, 0x0D, 0x0A, 0x00, 0x0D, 0x0A, 0x51, 0x55, 0x49, 0x54, 0x0A,
68    /* version 2 + proxy connection byte */ 0x21
69];
70
71pub const UDP_PROXY_PROTOCOL_LEN_V6: usize = 16 + 36;
72pub const UDP_PROXY_PROTOCOL_LEN_V4: usize = 16 + 12;
73pub const UDP_PROXY_PROTOCOL_MAX_LEN: usize = UDP_PROXY_PROTOCOL_LEN_V6;
74
75impl ProxyProtocolHeader {
76    pub async fn write_v1_tcp<W: AsyncWrite + Unpin>(&self, out: &mut W) -> Result<(), std::io::Error> {
77        out.write_all(self.to_string().as_bytes()).await
78    }
79
80    pub async fn write_v2_tcp<W: AsyncWrite + Unpin>(&self, out: &mut W) -> Result<(), std::io::Error> {
81        out.write_all(PROXY_PROTOCOL_V2_HEADER).await?;
82
83        match self {
84            Self::AfInet { client_ip, proxy_ip, client_port, proxy_port } => {
85                out.write_all(&[ /* TCP4: AF_INET + STREAM */ 0x11 ]).await?;
86                out.write_all(/* length */ &12u16.to_be_bytes()).await?;
87                out.write_all(&client_ip.octets()).await?;
88                out.write_all(&proxy_ip.octets()).await?;
89                out.write_all(&client_port.to_be_bytes()).await?;
90                out.write_all(&proxy_port.to_be_bytes()).await?;
91            }
92            Self::AfInet6 { client_ip, proxy_ip, client_port, proxy_port } => {
93                out.write_all(&[ /* TCP6: AF_INET6 + STREAM */ 0x21 ]).await?;
94                out.write_all(/* length */ &36u16.to_be_bytes()).await?;
95                out.write_all(&client_ip.octets()).await?;
96                out.write_all(&proxy_ip.octets()).await?;
97                out.write_all(&client_port.to_be_bytes()).await?;
98                out.write_all(&proxy_port.to_be_bytes()).await?;
99            }
100        }
101
102        Ok(())
103    }
104
105    pub fn write_v2_udp<W: Write>(&self, out: &mut W) -> Result<(), std::io::Error> {
106        out.write_all(PROXY_PROTOCOL_V2_HEADER)?;
107
108        match self {
109            Self::AfInet { client_ip, proxy_ip, client_port, proxy_port } => {
110                out.write_all(&[ /* UDP4: AF_INET + DGRAM */ 0x12 ])?;
111                out.write_all(/* length */ &12u16.to_be_bytes())?;
112                out.write_all(&client_ip.octets())?;
113                out.write_all(&proxy_ip.octets())?;
114                out.write_all(&client_port.to_be_bytes())?;
115                out.write_all(&proxy_port.to_be_bytes())?;
116            }
117            Self::AfInet6 { client_ip, proxy_ip, client_port, proxy_port } => {
118                out.write_all(&[ /* UDP6: AF_INET6 + DGRAM */ 0x22 ])?;
119                out.write_all(/* length */ &36u16.to_be_bytes())?;
120                out.write_all(&client_ip.octets())?;
121                out.write_all(&proxy_ip.octets())?;
122                out.write_all(&client_port.to_be_bytes())?;
123                out.write_all(&proxy_port.to_be_bytes())?;
124            }
125        }
126
127        Ok(())
128    }
129
130    pub fn parse_v2_udp<R: std::io::Read>(buffer: &mut R) -> Option<Self> {
131        let mut header = [0u8; PROXY_PROTOCOL_V2_HEADER.len()];
132        buffer.read_exact(&mut header).ok()?;
133
134        if !header.eq(PROXY_PROTOCOL_V2_HEADER) {
135            return None;
136        }
137
138        let proto_type = buffer.read_u8().ok()?;
139        match proto_type {
140            /* AF_INET */
141            0x12 => {
142                let mut slab = [0u8; 14];
143                buffer.read_exact(&mut slab).ok()?;
144                let mut reader = &slab[..];
145
146                /* length */
147                if reader.read_u16::<BigEndian>().unwrap() != 12 {
148                    return None;
149                }
150
151                let client_ip = reader.read_ip4().unwrap();
152                let proxy_ip = reader.read_ip4().unwrap();
153                let client_port = reader.read_u16::<BigEndian>().unwrap();
154                let proxy_port = reader.read_u16::<BigEndian>().unwrap();
155
156                Some(ProxyProtocolHeader::AfInet {
157                    client_ip,
158                    proxy_ip,
159                    client_port,
160                    proxy_port,
161                })
162            }
163            /* AF_INET6 */
164            0x22 => {
165                let mut slab = [0u8; 38];
166                buffer.read_exact(&mut slab).ok()?;
167                let mut reader = &slab[..];
168
169                if reader.read_u16::<BigEndian>().unwrap() != 36 {
170                    return None;
171                }
172
173                let client_ip = reader.read_ip6().unwrap();
174                let proxy_ip = reader.read_ip6().unwrap();
175                let client_port = reader.read_u16::<BigEndian>().unwrap();
176                let proxy_port = reader.read_u16::<BigEndian>().unwrap();
177
178                Some(ProxyProtocolHeader::AfInet6 {
179                    client_ip,
180                    proxy_ip,
181                    client_port,
182                    proxy_port,
183                })
184            }
185            _ => None,
186        }
187    }
188}
189
190#[cfg(test)]
191mod test {
192    use super::ProxyProtocolHeader;
193
194    #[test]
195    fn test_parse_header() {
196        let mut buffer = Vec::new();
197        let header = ProxyProtocolHeader::AfInet {
198            client_ip: "123.45.12.34".parse().unwrap(),
199            proxy_ip: "5.6.7.8".parse().unwrap(),
200            client_port: 421,
201            proxy_port: 662
202        };
203
204        header.write_v2_udp(&mut buffer).unwrap();
205
206        let mut reader = &buffer[..];
207        let parsed = ProxyProtocolHeader::parse_v2_udp(&mut reader).unwrap();
208        assert_eq!(header, parsed);
209    }
210}
playit_agent_core/network/proxy_protocol.rs

playit_agent_core/network/
proxy_protocol.rs
