astelm/telemt
1
0
Fork 0
mirror of https://github.com/telemt/telemt.git synced 2026-04-27 15:34:09 +03:00
Code Issues Packages Projects Releases Wiki Activity

Format

Browse Source
This commit is contained in:
Alexey
2026-03-21 15:45:29 +03:00
parent 7a8f946029
commit d7bbb376c9
154 changed files with 6194 additions and 3775 deletions
Download Patch File Download Diff File Expand all files Collapse all files

187
src/transport/proxy_protocol.rs
View File

@@ -1,16 +1,15 @@
//! HAProxy PROXY protocol V1/V2
use crate::error::{ProxyError, Result};
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
use tokio::io::{AsyncRead, AsyncReadExt};
use crate::error::{ProxyError, Result};
/// PROXY protocol v1 signature
const PROXY_V1_SIGNATURE: &[u8] = b"PROXY ";
/// PROXY protocol v2 signature
const PROXY_V2_SIGNATURE: &[u8] = &[
0x0d, 0x0a, 0x0d, 0x0a, 0x00, 0x0d, 0x0a,
0x51, 0x55, 0x49, 0x54, 0x0a
0x0d, 0x0a, 0x0d, 0x0a, 0x00, 0x0d, 0x0a, 0x51, 0x55, 0x49, 0x54, 0x0a,
];
/// Minimum length for v1 detection
@@ -51,7 +50,7 @@ impl ProxyProtocolInfo {
}
/// Parse PROXY protocol header from a stream
///
///
/// Returns the parsed info or an error if the header is invalid.
/// The stream position is advanced past the header.
pub async fn parse_proxy_protocol<R: AsyncRead + Unpin>(
@@ -60,23 +59,27 @@ pub async fn parse_proxy_protocol<R: AsyncRead + Unpin>(
) -> Result<ProxyProtocolInfo> {
// Read enough bytes to detect version
let mut header = [0u8; PROXY_V2_MIN_LEN];
reader.read_exact(&mut header[..PROXY_V1_MIN_LEN]).await
reader
.read_exact(&mut header[..PROXY_V1_MIN_LEN])
.await
.map_err(|_| ProxyError::InvalidProxyProtocol)?;
// Check for v1
if header[..PROXY_V1_MIN_LEN] == PROXY_V1_SIGNATURE[..] {
return parse_v1(reader, default_peer).await;
}
// Read rest for v2 detection
reader.read_exact(&mut header[PROXY_V1_MIN_LEN..]).await
reader
.read_exact(&mut header[PROXY_V1_MIN_LEN..])
.await
.map_err(|_| ProxyError::InvalidProxyProtocol)?;
// Check for v2
if header[..12] == PROXY_V2_SIGNATURE[..] {
return parse_v2(reader, &header, default_peer).await;
}
Err(ProxyError::InvalidProxyProtocol)
}
@@ -88,43 +91,49 @@ async fn parse_v1<R: AsyncRead + Unpin>(
// Read until CRLF (max 107 bytes total for v1)
let mut line = Vec::with_capacity(128);
line.extend_from_slice(PROXY_V1_SIGNATURE);
loop {
let mut byte = [0u8];
reader.read_exact(&mut byte).await
reader
.read_exact(&mut byte)
.await
.map_err(|_| ProxyError::InvalidProxyProtocol)?;
line.push(byte[0]);
if line.ends_with(b"\r\n") {
break;
}
if line.len() > 256 {
return Err(ProxyError::InvalidProxyProtocol);
}
}
// Parse the line: PROXY TCP4/TCP6/UNKNOWN src_ip dst_ip src_port dst_port
let line_str = std::str::from_utf8(&line[PROXY_V1_MIN_LEN..line.len() - 2])
.map_err(|_| ProxyError::InvalidProxyProtocol)?;
let parts: Vec<&str> = line_str.split_whitespace().collect();
if parts.is_empty() {
return Err(ProxyError::InvalidProxyProtocol);
}
match parts[0] {
"TCP4" | "TCP6" if parts.len() >= 5 => {
let src_ip: IpAddr = parts[1].parse()
let src_ip: IpAddr = parts[1]
.parse()
.map_err(|_| ProxyError::InvalidProxyProtocol)?;
let dst_ip: IpAddr = parts[2].parse()
let dst_ip: IpAddr = parts[2]
.parse()
.map_err(|_| ProxyError::InvalidProxyProtocol)?;
let src_port: u16 = parts[3].parse()
let src_port: u16 = parts[3]
.parse()
.map_err(|_| ProxyError::InvalidProxyProtocol)?;
let dst_port: u16 = parts[4].parse()
let dst_port: u16 = parts[4]
.parse()
.map_err(|_| ProxyError::InvalidProxyProtocol)?;
Ok(ProxyProtocolInfo {
src_addr: SocketAddr::new(src_ip, src_port),
dst_addr: Some(SocketAddr::new(dst_ip, dst_port)),
@@ -152,22 +161,24 @@ async fn parse_v2<R: AsyncRead + Unpin>(
let version_command = header[12];
let version = version_command >> 4;
let command = version_command & 0x0f;
// Must be version 2
if version != 2 {
return Err(ProxyError::InvalidProxyProtocol);
}
let family_protocol = header[13];
let addr_len = u16::from_be_bytes([header[14], header[15]]) as usize;
// Read address data
let mut addr_data = vec![0u8; addr_len];
if addr_len > 0 {
reader.read_exact(&mut addr_data).await
reader
.read_exact(&mut addr_data)
.await
.map_err(|_| ProxyError::InvalidProxyProtocol)?;
}
// LOCAL command (0x0) - use default peer
if command == 0 {
return Ok(ProxyProtocolInfo {
@@ -176,28 +187,22 @@ async fn parse_v2<R: AsyncRead + Unpin>(
version: 2,
});
}
// PROXY command (0x1) - parse addresses
if command != 1 {
return Err(ProxyError::InvalidProxyProtocol);
}
let family = family_protocol >> 4;
match family {
address_family::INET if addr_len >= 12 => {
// IPv4: 4 + 4 + 2 + 2 = 12 bytes
let src_ip = Ipv4Addr::new(
addr_data[0], addr_data[1],
addr_data[2], addr_data[3]
);
let dst_ip = Ipv4Addr::new(
addr_data[4], addr_data[5],
addr_data[6], addr_data[7]
);
let src_ip = Ipv4Addr::new(addr_data[0], addr_data[1], addr_data[2], addr_data[3]);
let dst_ip = Ipv4Addr::new(addr_data[4], addr_data[5], addr_data[6], addr_data[7]);
let src_port = u16::from_be_bytes([addr_data[8], addr_data[9]]);
let dst_port = u16::from_be_bytes([addr_data[10], addr_data[11]]);
Ok(ProxyProtocolInfo {
src_addr: SocketAddr::new(IpAddr::V4(src_ip), src_port),
dst_addr: Some(SocketAddr::new(IpAddr::V4(dst_ip), dst_port)),
@@ -206,28 +211,22 @@ async fn parse_v2<R: AsyncRead + Unpin>(
}
address_family::INET6 if addr_len >= 36 => {
// IPv6: 16 + 16 + 2 + 2 = 36 bytes
let src_ip = Ipv6Addr::from(
<[u8; 16]>::try_from(&addr_data[0..16]).unwrap()
);
let dst_ip = Ipv6Addr::from(
<[u8; 16]>::try_from(&addr_data[16..32]).unwrap()
);
let src_ip = Ipv6Addr::from(<[u8; 16]>::try_from(&addr_data[0..16]).unwrap());
let dst_ip = Ipv6Addr::from(<[u8; 16]>::try_from(&addr_data[16..32]).unwrap());
let src_port = u16::from_be_bytes([addr_data[32], addr_data[33]]);
let dst_port = u16::from_be_bytes([addr_data[34], addr_data[35]]);
Ok(ProxyProtocolInfo {
src_addr: SocketAddr::new(IpAddr::V6(src_ip), src_port),
dst_addr: Some(SocketAddr::new(IpAddr::V6(dst_ip), dst_port)),
version: 2,
})
}
address_family::UNSPEC => {
Ok(ProxyProtocolInfo {
src_addr: default_peer,
dst_addr: None,
version: 2,
})
}
address_family::UNSPEC => Ok(ProxyProtocolInfo {
src_addr: default_peer,
dst_addr: None,
version: 2,
}),
_ => Err(ProxyError::InvalidProxyProtocol),
}
}
@@ -247,33 +246,32 @@ impl ProxyProtocolV1Builder {
dst_addr: None,
}
}
pub fn tcp4(mut self, src: SocketAddr, dst: SocketAddr) -> Self {
self.family = "TCP4";
self.src_addr = Some(src);
self.dst_addr = Some(dst);
self
}
pub fn tcp6(mut self, src: SocketAddr, dst: SocketAddr) -> Self {
self.family = "TCP6";
self.src_addr = Some(src);
self.dst_addr = Some(dst);
self
}
pub fn build(&self) -> Vec<u8> {
match (self.src_addr, self.dst_addr) {
(Some(src), Some(dst)) => {
format!(
"PROXY {} {} {} {} {}\r\n",
self.family,
src.ip(),
dst.ip(),
src.port(),
dst.port()
).into_bytes()
}
(Some(src), Some(dst)) => format!(
"PROXY {} {} {} {} {}\r\n",
self.family,
src.ip(),
dst.ip(),
src.port(),
dst.port()
)
.into_bytes(),
_ => b"PROXY UNKNOWN\r\n".to_vec(),
}
}
@@ -299,7 +297,10 @@ impl Default for ProxyProtocolV2Builder {
impl ProxyProtocolV2Builder {
pub fn new() -> Self {
Self { src: None, dst: None }
Self {
src: None,
dst: None,
}
}
pub fn with_addrs(mut self, src: SocketAddr, dst: SocketAddr) -> Self {
@@ -347,34 +348,34 @@ impl ProxyProtocolV2Builder {
mod tests {
use super::*;
use std::io::Cursor;
#[tokio::test]
async fn test_parse_v1_tcp4() {
let header = b"PROXY TCP4 192.168.1.1 10.0.0.1 12345 443\r\n";
let mut cursor = Cursor::new(&header[PROXY_V1_MIN_LEN..]);
let default = "0.0.0.0:0".parse().unwrap();
// Simulate that we've already read the signature
let info = parse_v1(&mut cursor, default).await.unwrap();
assert_eq!(info.version, 1);
assert_eq!(info.src_addr.ip().to_string(), "192.168.1.1");
assert_eq!(info.src_addr.port(), 12345);
assert!(info.dst_addr.is_some());
}
#[tokio::test]
async fn test_parse_v1_unknown() {
let header = b"PROXY UNKNOWN\r\n";
let mut cursor = Cursor::new(&header[PROXY_V1_MIN_LEN..]);
let default: SocketAddr = "1.2.3.4:5678".parse().unwrap();
let info = parse_v1(&mut cursor, default).await.unwrap();
assert_eq!(info.version, 1);
assert_eq!(info.src_addr, default);
}
#[tokio::test]
async fn test_parse_v2_tcp4() {
// v2 header for TCP4
@@ -384,24 +385,24 @@ mod tests {
header[13] = 0x11; // AF_INET, STREAM
header[14] = 0x00;
header[15] = 0x0c; // 12 bytes of address data
let addr_data = [
192, 168, 1, 1, // src IP
10, 0, 0, 1, // dst IP
0x30, 0x39, // src port (12345)
0x01, 0xbb, // dst port (443)
192, 168, 1, 1, // src IP
10, 0, 0, 1, // dst IP
0x30, 0x39, // src port (12345)
0x01, 0xbb, // dst port (443)
];
let mut cursor = Cursor::new(addr_data.to_vec());
let default = "0.0.0.0:0".parse().unwrap();
let info = parse_v2(&mut cursor, &header, default).await.unwrap();
assert_eq!(info.version, 2);
assert_eq!(info.src_addr.ip().to_string(), "192.168.1.1");
assert_eq!(info.src_addr.port(), 12345);
}
#[tokio::test]
async fn test_parse_v2_local() {
let mut header = [0u8; 16];
@@ -410,29 +411,27 @@ mod tests {
header[13] = 0x00;
header[14] = 0x00;
header[15] = 0x00; // 0 bytes of address data
let mut cursor = Cursor::new(Vec::new());
let default: SocketAddr = "1.2.3.4:5678".parse().unwrap();
let info = parse_v2(&mut cursor, &header, default).await.unwrap();
assert_eq!(info.version, 2);
assert_eq!(info.src_addr, default);
}
#[test]
fn test_v1_builder() {
let src: SocketAddr = "192.168.1.1:12345".parse().unwrap();
let dst: SocketAddr = "10.0.0.1:443".parse().unwrap();
let header = ProxyProtocolV1Builder::new()
.tcp4(src, dst)
.build();
let header = ProxyProtocolV1Builder::new().tcp4(src, dst).build();
let expected = b"PROXY TCP4 192.168.1.1 10.0.0.1 12345 443\r\n";
assert_eq!(header, expected);
}
#[test]
fn test_v1_builder_unknown() {
let header = ProxyProtocolV1Builder::new().build();