tako-rs-server 2.5.0

//! PROXY protocol address/header value types and TLV expansion.

use std::net::SocketAddr;

/// PROXY protocol version.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ProxyVersion {
  V1,
  V2,
}

/// Transport protocol from the PROXY header.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ProxyTransport {
  Tcp,
  Udp,
  Unknown,
}

/// Raw PROXY v2 TLV, kept for custom types; known ones also fill typed [`ProxyHeader`] fields.
#[derive(Debug, Clone)]
pub struct ProxyTlv {
  /// `PP2_TYPE_*` identifier byte.
  pub kind: u8,
  /// Raw TLV value bytes.
  pub value: Vec<u8>,
}

/// TLS-derived PROXY v2 sub-TLVs (`PP2_TYPE_SSL` container, type 0x20).
#[derive(Debug, Clone, Default)]
pub struct ProxyTlsInfo {
  /// `PP2_CLIENT_SSL` bitfield.
  pub client_flags: u8,
  /// rustls/openssl-style verify result code.
  pub verify: u32,
  /// `PP2_SUBTYPE_SSL_VERSION` (e.g. `"TLSv1.3"`).
  pub version: Option<String>,
  /// `PP2_SUBTYPE_SSL_CN` (peer common name).
  pub common_name: Option<String>,
  /// `PP2_SUBTYPE_SSL_CIPHER`.
  pub cipher: Option<String>,
  /// `PP2_SUBTYPE_SSL_SIG_ALG`.
  pub sig_alg: Option<String>,
  /// `PP2_SUBTYPE_SSL_KEY_ALG`.
  pub key_alg: Option<String>,
}

// Stops a peer from making each connection allocate 64 KiB; the largest address block (UNIX, 216
// bytes) plus a generous TLV stack stays below it.
pub(crate) const MAX_PROXY_ADDR_LEN: usize = 536;

const PP2_TYPE_ALPN: u8 = 0x01;
const PP2_TYPE_AUTHORITY: u8 = 0x02;
pub(crate) const PP2_TYPE_CRC32C: u8 = 0x03;
const PP2_TYPE_NOOP: u8 = 0x04;
const PP2_TYPE_UNIQUE_ID: u8 = 0x05;
const PP2_TYPE_SSL: u8 = 0x20;
const PP2_SUBTYPE_SSL_VERSION: u8 = 0x21;
const PP2_SUBTYPE_SSL_CN: u8 = 0x22;
const PP2_SUBTYPE_SSL_CIPHER: u8 = 0x23;
const PP2_SUBTYPE_SSL_SIG_ALG: u8 = 0x24;
const PP2_SUBTYPE_SSL_KEY_ALG: u8 = 0x25;
const PP2_TYPE_NETNS: u8 = 0x30;
const PP2_TYPE_AWS_VPC_ENDPOINT_ID: u8 = 0xEA;

/// Parsed PROXY protocol header: the real client (`source`), the address it connected to, and the
/// v2 TLVs as typed fields where they map cleanly.
#[derive(Debug, Clone)]
pub struct ProxyHeader {
  /// Protocol version (v1 text or v2 binary).
  pub version: ProxyVersion,
  /// Transport protocol.
  pub transport: ProxyTransport,
  /// Real client address (the original source).
  pub source: Option<SocketAddr>,
  /// Proxy/server address (the destination the client connected to).
  pub destination: Option<SocketAddr>,
  /// `AF_UNIX` source path, when the connection family is Unix.
  pub source_unix: Option<std::path::PathBuf>,
  /// `AF_UNIX` destination path, when the connection family is Unix.
  pub destination_unix: Option<std::path::PathBuf>,
  /// `PP2_TYPE_AUTHORITY` (a.k.a. SNI) value if present.
  pub authority: Option<String>,
  /// `PP2_TYPE_ALPN` protocol bytes if present.
  pub alpn: Option<Vec<u8>>,
  /// AWS VPC endpoint identifier (`PP2` type 0xEA) if present.
  pub aws_vpc_endpoint_id: Option<String>,
  /// Decoded `PP2_TYPE_SSL` sub-TLVs.
  pub tls: Option<ProxyTlsInfo>,
  /// Unique connection identifier (`PP2` type 0x05).
  pub unique_id: Option<Vec<u8>>,
  /// Every TLV, including types without a typed field.
  pub tlvs: Vec<ProxyTlv>,
  /// Whether the optional `PP2_TYPE_CRC32C` checksum matched; `None` without one and for v1. A
  /// mismatch is logged but does not reject the header.
  pub crc32c_verified: Option<bool>,
}

impl ProxyHeader {
  pub(crate) fn empty(version: ProxyVersion, transport: ProxyTransport) -> Self {
    Self {
      version,
      transport,
      source: None,
      destination: None,
      source_unix: None,
      destination_unix: None,
      authority: None,
      alpn: None,
      aws_vpc_endpoint_id: None,
      tls: None,
      unique_id: None,
      tlvs: Vec::new(),
      crc32c_verified: None,
    }
  }
}

/// Walks a PROXY v2 TLV byte stream and applies each entry to a [`ProxyHeader`].
pub(crate) fn apply_tlvs(header: &mut ProxyHeader, mut buf: &[u8]) {
  while buf.len() >= 3 {
    let kind = buf[0];
    let len = u16::from_be_bytes([buf[1], buf[2]]) as usize;
    if buf.len() < 3 + len {
      break;
    }
    let value = &buf[3..3 + len];
    match kind {
      PP2_TYPE_ALPN => header.alpn = Some(value.to_vec()),
      PP2_TYPE_AUTHORITY => {
        if let Ok(s) = std::str::from_utf8(value) {
          header.authority = Some(s.to_string());
        }
      }
      PP2_TYPE_AWS_VPC_ENDPOINT_ID => {
        if let Ok(s) = std::str::from_utf8(value) {
          header.aws_vpc_endpoint_id = Some(s.to_string());
        }
      }
      PP2_TYPE_UNIQUE_ID => header.unique_id = Some(value.to_vec()),
      PP2_TYPE_SSL => {
        // 1 byte of client flags, a 4-byte big-endian verify result, then sub-TLVs.
        if value.len() >= 5 {
          let mut tls = ProxyTlsInfo {
            client_flags: value[0],
            verify: u32::from_be_bytes([value[1], value[2], value[3], value[4]]),
            ..Default::default()
          };
          let mut sub = &value[5..];
          while sub.len() >= 3 {
            let sk = sub[0];
            let slen = u16::from_be_bytes([sub[1], sub[2]]) as usize;
            if sub.len() < 3 + slen {
              break;
            }
            let sval = &sub[3..3 + slen];
            match sk {
              PP2_SUBTYPE_SSL_VERSION => {
                tls.version = std::str::from_utf8(sval).ok().map(str::to_string);
              }
              PP2_SUBTYPE_SSL_CN => {
                tls.common_name = std::str::from_utf8(sval).ok().map(str::to_string);
              }
              PP2_SUBTYPE_SSL_CIPHER => {
                tls.cipher = std::str::from_utf8(sval).ok().map(str::to_string);
              }
              PP2_SUBTYPE_SSL_SIG_ALG => {
                tls.sig_alg = std::str::from_utf8(sval).ok().map(str::to_string);
              }
              PP2_SUBTYPE_SSL_KEY_ALG => {
                tls.key_alg = std::str::from_utf8(sval).ok().map(str::to_string);
              }
              _ => {}
            }
            sub = &sub[3 + slen..];
          }
          header.tls = Some(tls);
        }
      }
      // `verify_v2_crc32c` already checked the CRC; NOOP and NETNS carry nothing to surface.
      PP2_TYPE_CRC32C | PP2_TYPE_NOOP | PP2_TYPE_NETNS => {}
      _ => {}
    }
    header.tlvs.push(ProxyTlv {
      kind,
      value: value.to_vec(),
    });
    buf = &buf[3 + len..];
  }
}