|
| 1 | +// SPDX-FileCopyrightText: © 2024-2025 Phala Network <dstack@phala.network> |
| 2 | +// |
| 3 | +// SPDX-License-Identifier: Apache-2.0 |
| 4 | + |
| 5 | +use std::net::SocketAddr; |
| 6 | + |
| 7 | +use anyhow::{bail, Context, Result}; |
| 8 | +use proxy_protocol::{version1 as v1, version2 as v2, ProxyHeader}; |
| 9 | +use tokio::{ |
| 10 | + io::{AsyncRead, AsyncReadExt}, |
| 11 | + net::TcpStream, |
| 12 | +}; |
| 13 | + |
| 14 | +use crate::config::ProxyConfig; |
| 15 | + |
| 16 | +const V1_PROTOCOL_PREFIX: &str = "PROXY"; |
| 17 | +const V1_PREFIX_LEN: usize = 5; |
| 18 | +const V1_MAX_LENGTH: usize = 107; |
| 19 | +const V1_TERMINATOR: &[u8] = b"\r\n"; |
| 20 | + |
| 21 | +const V2_PROTOCOL_PREFIX: &[u8] = b"\r\n\r\n\0\r\nQUIT\n"; |
| 22 | +const V2_PREFIX_LEN: usize = 12; |
| 23 | +const V2_MINIMUM_LEN: usize = 16; |
| 24 | +const V2_LENGTH_INDEX: usize = 14; |
| 25 | +const READ_BUFFER_LEN: usize = 512; |
| 26 | +const V2_MAX_LENGTH: usize = 2048; |
| 27 | + |
| 28 | +/// Read or synthesize the inbound proxy protocol header. |
| 29 | +/// |
| 30 | +/// When `inbound_pp_enabled` is true, reads a PP header from the stream (e.g. from an upstream |
| 31 | +/// load balancer). When false, synthesizes one from the TCP peer address. |
| 32 | +pub(crate) async fn get_inbound_pp_header( |
| 33 | + inbound: TcpStream, |
| 34 | + config: &ProxyConfig, |
| 35 | +) -> Result<(TcpStream, ProxyHeader)> { |
| 36 | + if config.inbound_pp_enabled { |
| 37 | + read_proxy_header(inbound).await |
| 38 | + } else { |
| 39 | + let header = create_inbound_pp_header(&inbound); |
| 40 | + Ok((inbound, header)) |
| 41 | + } |
| 42 | +} |
| 43 | + |
| 44 | +pub struct DisplayAddr<'a>(pub &'a ProxyHeader); |
| 45 | + |
| 46 | +impl std::fmt::Display for DisplayAddr<'_> { |
| 47 | + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { |
| 48 | + match self.0 { |
| 49 | + ProxyHeader::Version2 { addresses, .. } => match addresses { |
| 50 | + v2::ProxyAddresses::Ipv4 { source, .. } => write!(f, "{}", source), |
| 51 | + v2::ProxyAddresses::Ipv6 { source, .. } => write!(f, "{}", source), |
| 52 | + v2::ProxyAddresses::Unix { .. } => write!(f, "<unix>"), |
| 53 | + v2::ProxyAddresses::Unspec => write!(f, "<unspec>"), |
| 54 | + }, |
| 55 | + ProxyHeader::Version1 { addresses, .. } => match addresses { |
| 56 | + v1::ProxyAddresses::Ipv4 { source, .. } => write!(f, "{}", source), |
| 57 | + v1::ProxyAddresses::Ipv6 { source, .. } => write!(f, "{}", source), |
| 58 | + v1::ProxyAddresses::Unknown => write!(f, "<unknown>"), |
| 59 | + }, |
| 60 | + _ => write!(f, "<unknown ver>"), |
| 61 | + } |
| 62 | + } |
| 63 | +} |
| 64 | + |
| 65 | +fn create_inbound_pp_header(inbound: &TcpStream) -> ProxyHeader { |
| 66 | + let peer_addr = inbound.peer_addr().ok(); |
| 67 | + let local_addr = inbound.local_addr().ok(); |
| 68 | + |
| 69 | + match (peer_addr, local_addr) { |
| 70 | + (Some(SocketAddr::V4(source)), Some(SocketAddr::V4(destination))) => { |
| 71 | + ProxyHeader::Version2 { |
| 72 | + command: v2::ProxyCommand::Proxy, |
| 73 | + transport_protocol: v2::ProxyTransportProtocol::Stream, |
| 74 | + addresses: v2::ProxyAddresses::Ipv4 { |
| 75 | + source, |
| 76 | + destination, |
| 77 | + }, |
| 78 | + } |
| 79 | + } |
| 80 | + (Some(SocketAddr::V6(source)), Some(SocketAddr::V6(destination))) => { |
| 81 | + ProxyHeader::Version2 { |
| 82 | + command: v2::ProxyCommand::Proxy, |
| 83 | + transport_protocol: v2::ProxyTransportProtocol::Stream, |
| 84 | + addresses: v2::ProxyAddresses::Ipv6 { |
| 85 | + source, |
| 86 | + destination, |
| 87 | + }, |
| 88 | + } |
| 89 | + } |
| 90 | + _ => ProxyHeader::Version2 { |
| 91 | + command: v2::ProxyCommand::Proxy, |
| 92 | + transport_protocol: v2::ProxyTransportProtocol::Stream, |
| 93 | + addresses: v2::ProxyAddresses::Unspec, |
| 94 | + }, |
| 95 | + } |
| 96 | +} |
| 97 | + |
| 98 | +async fn read_proxy_header<I>(mut stream: I) -> Result<(I, ProxyHeader)> |
| 99 | +where |
| 100 | + I: AsyncRead + Unpin, |
| 101 | +{ |
| 102 | + let mut buffer = [0; READ_BUFFER_LEN]; |
| 103 | + let mut dynamic_buffer = None; |
| 104 | + |
| 105 | + stream.read_exact(&mut buffer[..V1_PREFIX_LEN]).await?; |
| 106 | + |
| 107 | + if &buffer[..V1_PREFIX_LEN] == V1_PROTOCOL_PREFIX.as_bytes() { |
| 108 | + read_v1_header(&mut stream, &mut buffer).await?; |
| 109 | + } else { |
| 110 | + stream |
| 111 | + .read_exact(&mut buffer[V1_PREFIX_LEN..V2_MINIMUM_LEN]) |
| 112 | + .await?; |
| 113 | + if &buffer[..V2_PREFIX_LEN] == V2_PROTOCOL_PREFIX { |
| 114 | + dynamic_buffer = read_v2_header(&mut stream, &mut buffer).await?; |
| 115 | + } else { |
| 116 | + bail!("no valid proxy protocol header detected"); |
| 117 | + } |
| 118 | + } |
| 119 | + |
| 120 | + let mut buffer = dynamic_buffer.as_deref().unwrap_or(&buffer[..]); |
| 121 | + |
| 122 | + let header = |
| 123 | + proxy_protocol::parse(&mut buffer).context("failed to parse proxy protocol header")?; |
| 124 | + Ok((stream, header)) |
| 125 | +} |
| 126 | + |
| 127 | +async fn read_v2_header<I>( |
| 128 | + mut stream: I, |
| 129 | + buffer: &mut [u8; READ_BUFFER_LEN], |
| 130 | +) -> Result<Option<Vec<u8>>> |
| 131 | +where |
| 132 | + I: AsyncRead + Unpin, |
| 133 | +{ |
| 134 | + let length = |
| 135 | + u16::from_be_bytes([buffer[V2_LENGTH_INDEX], buffer[V2_LENGTH_INDEX + 1]]) as usize; |
| 136 | + let full_length = V2_MINIMUM_LEN + length; |
| 137 | + |
| 138 | + if full_length > V2_MAX_LENGTH { |
| 139 | + bail!("v2 proxy protocol header is too long"); |
| 140 | + } |
| 141 | + |
| 142 | + if full_length > READ_BUFFER_LEN { |
| 143 | + let mut dynamic_buffer = Vec::with_capacity(full_length); |
| 144 | + dynamic_buffer.extend_from_slice(&buffer[..V2_MINIMUM_LEN]); |
| 145 | + dynamic_buffer.resize(full_length, 0); |
| 146 | + stream |
| 147 | + .read_exact(&mut dynamic_buffer[V2_MINIMUM_LEN..full_length]) |
| 148 | + .await?; |
| 149 | + |
| 150 | + Ok(Some(dynamic_buffer)) |
| 151 | + } else { |
| 152 | + stream |
| 153 | + .read_exact(&mut buffer[V2_MINIMUM_LEN..full_length]) |
| 154 | + .await?; |
| 155 | + |
| 156 | + Ok(None) |
| 157 | + } |
| 158 | +} |
| 159 | + |
| 160 | +async fn read_v1_header<I>(mut stream: I, buffer: &mut [u8; READ_BUFFER_LEN]) -> Result<()> |
| 161 | +where |
| 162 | + I: AsyncRead + Unpin, |
| 163 | +{ |
| 164 | + let mut end_found = false; |
| 165 | + for i in V1_PREFIX_LEN..V1_MAX_LENGTH { |
| 166 | + buffer[i] = stream.read_u8().await?; |
| 167 | + |
| 168 | + if [buffer[i - 1], buffer[i]] == V1_TERMINATOR { |
| 169 | + end_found = true; |
| 170 | + break; |
| 171 | + } |
| 172 | + } |
| 173 | + if !end_found { |
| 174 | + bail!("no valid proxy protocol header detected"); |
| 175 | + } |
| 176 | + |
| 177 | + Ok(()) |
| 178 | +} |
0 commit comments