sozu-lib 2.1.0

//! PROXY protocol v1 / v2 header model.
//!
//! Owns the `ProxyProtocolHeader` enum, the `into_bytes` serializers for
//! both wire versions, and the address helpers (`ProxyAddr`,
//! `ProtocolSupportedV1`, `Command`, `HeaderV1`, `HeaderV2`) shared by the
//! `expect`, `relay`, and `send` roles. No I/O here; pure data shapes.

use std::{
    fmt,
    net::{SocketAddr, SocketAddrV4, SocketAddrV6},
};

#[derive(PartialEq, Debug)]
pub enum ProxyProtocolHeader {
    V1(HeaderV1),
    V2(HeaderV2),
}

impl ProxyProtocolHeader {
    // Use this method to writte the header in the backend socket
    pub fn into_bytes(&self) -> Vec<u8> {
        match *self {
            ProxyProtocolHeader::V1(ref header) => header.into_bytes(),
            ProxyProtocolHeader::V2(ref header) => header.into_bytes(),
        }
    }
}

/// Indicate the proxied INET protocol and family
#[derive(Debug, PartialEq, Eq)]
pub enum ProtocolSupportedV1 {
    TCP4,    // TCP over IPv4
    TCP6,    // TCP over IPv6
    UNKNOWN, // unsupported,or unknown protocols
}

impl fmt::Display for ProtocolSupportedV1 {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            ProtocolSupportedV1::TCP4 => write!(f, "TCP4"),
            ProtocolSupportedV1::TCP6 => write!(f, "TCP6"),
            ProtocolSupportedV1::UNKNOWN => write!(f, "UNKNOWN"),
        }
    }
}

/// WARNING: proxy protocol v1 is never used in Sōzu, only v2.
/// The tests have been commented out, the whole code may be meant for deletion.
///
/// Proxy Protocol header for version 1 (text version)
/// Example:
/// - TCP/IPv4: `PROXY TCP4 255.255.255.255 255.255.255.255 65535 65535\r\n`
/// - TCP/IPv6: `PROXY TCP6 ffff:f...f:ffff ffff:f...f:ffff 65535 65535\r\n`
/// - Unknown: `PROXY UNKNOWN\r\n`
#[derive(Debug, PartialEq, Eq)]
pub struct HeaderV1 {
    pub protocol: ProtocolSupportedV1,
    pub addr_src: SocketAddr,
    pub addr_dst: SocketAddr,
}

const PROXY_PROTO_IDENTIFIER: &str = "PROXY";

impl HeaderV1 {
    pub fn new(addr_src: SocketAddr, addr_dst: SocketAddr) -> Self {
        let protocol = if addr_dst.is_ipv6() {
            ProtocolSupportedV1::TCP6
        } else if addr_dst.is_ipv4() {
            ProtocolSupportedV1::TCP4
        } else {
            ProtocolSupportedV1::UNKNOWN
        };

        HeaderV1 {
            protocol,
            addr_src,
            addr_dst,
        }
    }

    pub fn into_bytes(&self) -> Vec<u8> {
        let bytes = if self.protocol.eq(&ProtocolSupportedV1::UNKNOWN) {
            format!("{} {}\r\n", PROXY_PROTO_IDENTIFIER, self.protocol,).into_bytes()
        } else {
            format!(
                "{} {} {} {} {} {}\r\n",
                PROXY_PROTO_IDENTIFIER,
                self.protocol,
                self.addr_src.ip(),
                self.addr_dst.ip(),
                self.addr_src.port(),
                self.addr_dst.port(),
            )
            .into_bytes()
        };
        // The v1 text header is framed `PROXY ...\r\n`; both the identifier
        // prefix and the CRLF terminator are load-bearing for the peer parser.
        debug_assert!(
            bytes.starts_with(PROXY_PROTO_IDENTIFIER.as_bytes()),
            "v1 header must start with the PROXY identifier"
        );
        debug_assert!(
            bytes.ends_with(b"\r\n"),
            "v1 header must be CRLF-terminated"
        );
        bytes
    }
}

/*
#[cfg(test)]
mod test {

    use super::*;
    use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};

    #[test]
    fn it_should_return_a_correct_header_with_ipv4() {
        let header = HeaderV1::new(
            SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 80),
            SocketAddr::new(IpAddr::V4(Ipv4Addr::new(172, 17, 40, 59)), 80),
        );

        let header_to_cmp = "PROXY TCP4 127.0.0.1 172.17.40.59 80 80\r\n".as_bytes();

        assert_eq!(header_to_cmp, &header.into_bytes()[..]);
    }

    #[test]
    fn it_should_return_a_correct_header_with_ipv6() {
        let header = HeaderV1::new(
            SocketAddr::new(IpAddr::V6(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 0xffff)), 80),
            SocketAddr::new(IpAddr::V6(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0x9c, 0x76)), 80),
        );

        let header_to_cmp = "PROXY TCP6 ::ffff ::9c:76 80 80\r\n".as_bytes();

        assert_eq!(header_to_cmp, &header.into_bytes()[..]);
    }
}
*/

#[derive(Debug, PartialEq, Eq)]
pub enum Command {
    Local,
    Proxy,
}

#[derive(Debug, PartialEq)]
pub struct HeaderV2 {
    pub command: Command,
    pub family: u8, // protocol family and address
    pub addr: ProxyAddr,
}

impl HeaderV2 {
    pub fn new(command: Command, addr_src: SocketAddr, addr_dst: SocketAddr) -> Self {
        let addr = ProxyAddr::from(addr_src, addr_dst);
        let family = get_family(&addr);

        // Invariant: the cached `family` byte is exactly the one derived from
        // `addr`. Serialization writes both independently, so they must agree
        // or the wire header would self-contradict (family vs. address block).
        debug_assert_eq!(
            family,
            get_family(&addr),
            "cached family must match the address it describes"
        );
        debug_assert!(
            matches!(addr, ProxyAddr::AfUnspec) == (family == 0x00),
            "AfUnspec iff zero family byte"
        );

        HeaderV2 {
            command,
            family,
            addr,
        }
    }

    pub fn into_bytes(&self) -> Vec<u8> {
        let expected_len = self.len();
        let addr_len = self.addr.len() as usize;
        let mut header = Vec::with_capacity(expected_len);

        let signature = [
            0x0D, 0x0A, 0x0D, 0x0A, 0x00, 0x0D, 0x0A, 0x51, 0x55, 0x49, 0x54, 0x0A,
        ];
        header.extend_from_slice(&signature);
        debug_assert_eq!(
            header.len(),
            signature.len(),
            "v2 header must open with exactly the 12-byte signature"
        );

        let command = match self.command {
            Command::Local => 0,
            Command::Proxy => 1,
        };
        let ver_and_cmd = 0x20 | command;
        header.push(ver_and_cmd);

        header.push(self.family);
        header.extend_from_slice(&u16_to_array_of_u8(self.addr.len()));
        // Fixed 12-byte signature + 1 (ver/cmd) + 1 (family) + 2 (length) = 16
        // bytes precede the variable address block.
        debug_assert_eq!(
            header.len(),
            16,
            "v2 fixed prefix (signature + ver/cmd + family + length) must be 16 bytes"
        );
        self.addr.write_bytes_to(&mut header);
        // Postcondition: the serialized header length reconciles with the
        // declared `len()` and with the 16-byte prefix plus the address block.
        debug_assert_eq!(
            header.len(),
            expected_len,
            "serialized v2 header length must match HeaderV2::len()"
        );
        debug_assert_eq!(
            header.len(),
            16 + addr_len,
            "serialized v2 header must be the 16-byte prefix plus the address block"
        );
        header
    }

    pub fn len(&self) -> usize {
        // signature + ver_and_cmd + family + len + addr
        let total = 12 + 1 + 1 + 2 + self.addr.len() as usize;
        // The fixed prefix is always 16 bytes; the only variable part is the
        // address block, whose size is one of the enumerated ProxyAddr lengths.
        debug_assert!(
            total >= 16,
            "v2 header is at least its 16-byte fixed prefix"
        );
        debug_assert!(
            total <= 16 + 216,
            "v2 header never exceeds the 16-byte prefix plus the largest (unix) address block"
        );
        total
    }

    pub fn is_empty(&self) -> bool {
        0 == self.len()
    }
}

pub enum ProxyAddr {
    Ipv4Addr {
        src_addr: SocketAddrV4,
        dst_addr: SocketAddrV4,
    },
    Ipv6Addr {
        src_addr: SocketAddrV6,
        dst_addr: SocketAddrV6,
    },
    UnixAddr {
        src_addr: [u8; 108],
        dst_addr: [u8; 108],
    },
    AfUnspec,
}

impl ProxyAddr {
    pub fn from(addr_src: SocketAddr, addr_dst: SocketAddr) -> Self {
        let addr = match (addr_src, addr_dst) {
            (SocketAddr::V4(addr_ipv4_src), SocketAddr::V4(addr_ipv4_dst)) => ProxyAddr::Ipv4Addr {
                src_addr: addr_ipv4_src,
                dst_addr: addr_ipv4_dst,
            },
            (SocketAddr::V6(addr_ipv6_src), SocketAddr::V6(addr_ipv6_dst)) => ProxyAddr::Ipv6Addr {
                src_addr: addr_ipv6_src,
                dst_addr: addr_ipv6_dst,
            },
            _ => ProxyAddr::AfUnspec,
        };
        // Postcondition: a same-family pair maps to a concrete variant, a
        // mixed v4/v6 pair collapses to AfUnspec (PROXY-v2 has no cross-family
        // encoding). The chosen variant must agree with both inputs' family.
        debug_assert_eq!(
            matches!(addr, ProxyAddr::Ipv4Addr { .. }),
            addr_src.is_ipv4() && addr_dst.is_ipv4(),
            "Ipv4Addr variant iff both endpoints are IPv4"
        );
        debug_assert_eq!(
            matches!(addr, ProxyAddr::Ipv6Addr { .. }),
            addr_src.is_ipv6() && addr_dst.is_ipv6(),
            "Ipv6Addr variant iff both endpoints are IPv6"
        );
        addr
    }

    fn len(&self) -> u16 {
        match *self {
            ProxyAddr::Ipv4Addr { .. } => 12,
            ProxyAddr::Ipv6Addr { .. } => 36,
            ProxyAddr::UnixAddr { .. } => 216,
            ProxyAddr::AfUnspec => 0,
        }
    }

    pub fn source(&self) -> Option<SocketAddr> {
        match self {
            ProxyAddr::Ipv4Addr { src_addr: src, .. } => Some(SocketAddr::V4(*src)),
            ProxyAddr::Ipv6Addr { src_addr: src, .. } => Some(SocketAddr::V6(*src)),
            _ => None,
        }
    }

    pub fn destination(&self) -> Option<SocketAddr> {
        match self {
            ProxyAddr::Ipv4Addr { dst_addr: dst, .. } => Some(SocketAddr::V4(*dst)),
            ProxyAddr::Ipv6Addr { dst_addr: dst, .. } => Some(SocketAddr::V6(*dst)),
            _ => None,
        }
    }

    // TODO: rename to a less ambiguous name, like "write bytes to buffer"
    fn write_bytes_to(&self, buf: &mut Vec<u8>) {
        let before = buf.len();
        let declared = self.len() as usize;
        match *self {
            ProxyAddr::Ipv4Addr { src_addr, dst_addr } => {
                buf.extend_from_slice(&src_addr.ip().octets());
                buf.extend_from_slice(&dst_addr.ip().octets());
                buf.extend_from_slice(&u16_to_array_of_u8(src_addr.port()));
                buf.extend_from_slice(&u16_to_array_of_u8(dst_addr.port()));
            }
            ProxyAddr::Ipv6Addr { src_addr, dst_addr } => {
                buf.extend_from_slice(&src_addr.ip().octets());
                buf.extend_from_slice(&dst_addr.ip().octets());
                buf.extend_from_slice(&u16_to_array_of_u8(src_addr.port()));
                buf.extend_from_slice(&u16_to_array_of_u8(dst_addr.port()));
            }
            ProxyAddr::UnixAddr { src_addr, dst_addr } => {
                buf.extend_from_slice(&src_addr);
                buf.extend_from_slice(&dst_addr);
            }
            ProxyAddr::AfUnspec => {}
        };
        // Postcondition: the bytes appended must be exactly the address
        // block size advertised by `len()` — the wire length field is
        // derived from `len()`, so any divergence would desynchronize the
        // serialized header from its declared length.
        debug_assert!(
            buf.len() >= before,
            "write_bytes_to must never shrink the buffer"
        );
        debug_assert_eq!(
            buf.len() - before,
            declared,
            "appended address bytes must equal the declared ProxyAddr::len()"
        );
    }
}

// Implemented because we don't have the Debug for [u8; 108] (UnixAddr case)
impl fmt::Debug for ProxyAddr {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            ProxyAddr::Ipv4Addr { src_addr, dst_addr } => {
                write!(f, "{dst_addr:?} {src_addr:?}")
            }
            ProxyAddr::Ipv6Addr { src_addr, dst_addr } => {
                write!(f, "{dst_addr:?} {src_addr:?}")
            }
            ProxyAddr::UnixAddr { src_addr, dst_addr } => {
                write!(f, "{:?} {:?}", &dst_addr[..], &src_addr[..])
            }
            ProxyAddr::AfUnspec => write!(f, "AFUNSPEC"),
        }
    }
}

// Implemented because we don't have the PartialEq for [u8; 108] (UnixAddr case)
impl PartialEq for ProxyAddr {
    fn eq(&self, other: &ProxyAddr) -> bool {
        match *self {
            ProxyAddr::Ipv4Addr { src_addr, dst_addr } => match other {
                ProxyAddr::Ipv4Addr {
                    src_addr: src_other,
                    dst_addr: dst_other,
                } => *src_other == src_addr && *dst_other == dst_addr,
                _ => false,
            },
            ProxyAddr::Ipv6Addr { src_addr, dst_addr } => match other {
                ProxyAddr::Ipv6Addr {
                    src_addr: src_other,
                    dst_addr: dst_other,
                } => *src_other == src_addr && *dst_other == dst_addr,
                _ => false,
            },
            ProxyAddr::UnixAddr { src_addr, dst_addr } => match other {
                ProxyAddr::UnixAddr {
                    src_addr: src_other,
                    dst_addr: dst_other,
                } => src_other[..] == src_addr[..] && dst_other[..] == dst_addr[..],
                _ => false,
            },
            ProxyAddr::AfUnspec => {
                if let ProxyAddr::AfUnspec = other {
                    return true;
                }
                false
            }
        }
    }
}

fn get_family(addr: &ProxyAddr) -> u8 {
    let family = match *addr {
        ProxyAddr::Ipv4Addr { .. } => 0x10 | 0x01, // AF_INET  = 1 + STREAM = 1
        ProxyAddr::Ipv6Addr { .. } => 0x20 | 0x01, // AF_INET6 = 2 + STREAM = 1
        ProxyAddr::UnixAddr { .. } => 0x30 | 0x01, // AF_UNIX  = 3 + STREAM = 1
        ProxyAddr::AfUnspec => 0x00,               // AF_UNSPEC + UNSPEC
    };
    // Postcondition: the high nibble (address family) is within the
    // enumerated set the parser accepts (0..=3), and the low nibble
    // (transport) is STREAM for the concrete families, UNSPEC for AfUnspec.
    debug_assert!(
        (family >> 4) <= 0x03,
        "address family nibble must be one of AF_UNSPEC/INET/INET6/UNIX"
    );
    debug_assert!(
        matches!(addr, ProxyAddr::AfUnspec) == (family == 0x00),
        "only AfUnspec maps to the all-zero family byte"
    );
    debug_assert!(
        matches!(addr, ProxyAddr::AfUnspec) || (family & 0x0f) == 0x01,
        "concrete address families must advertise the STREAM transport"
    );
    family
}

fn u16_to_array_of_u8(x: u16) -> [u8; 2] {
    let b1: u8 = ((x >> 8) & 0xff) as u8;
    let b2: u8 = (x & 0xff) as u8;
    let out = [b1, b2];
    // Postcondition: this is a pure big-endian split — reassembling the two
    // bytes must round-trip to the input exactly. Port/length fields on the
    // wire depend on this being byte-exact.
    debug_assert_eq!(
        u16::from_be_bytes(out),
        x,
        "big-endian split must round-trip the input u16"
    );
    out
}

#[cfg(test)]
mod test_v2 {

    use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};

    use super::*;

    #[test]
    fn test_u16_to_array_of_u8() {
        let val_u16: u16 = 65534;
        let expected = [0xff, 0xfe];
        assert_eq!(expected, u16_to_array_of_u8(val_u16));
    }

    #[test]
    fn test_deserialize_tcp_ipv4_proxy_protocol_header() {
        let src_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(125, 25, 10, 1)), 8080);
        let dst_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(10, 4, 5, 8)), 4200);

        let header = HeaderV2::new(Command::Local, src_addr, dst_addr);
        let expected = &[
            0x0D, 0x0A, 0x0D, 0x0A, 0x00, 0x0D, 0x0A, 0x51, 0x55, 0x49, 0x54,
            0x0A, // MAGIC header
            0x20, // Version 2 and command LOCAL
            0x11, // family AF_UNIX with IPv4
            0x00, 0x0C, // address sizes = 12
            0x7D, 0x19, 0x0A, 0x01, // source address
            0x0A, 0x04, 0x05, 0x08, // destination address
            0x1F, 0x90, // source port
            0x10, 0x68, // destination port
        ];

        assert_eq!(expected, &header.into_bytes()[..]);
    }

    #[test]
    fn test_deserialize_tcp_ipv6_proxy_protocol_header() {
        let src_addr = SocketAddr::new(IpAddr::V6(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 1)), 8080);
        let dst_addr = SocketAddr::new(IpAddr::V6(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 1)), 4200);

        let header = HeaderV2::new(Command::Proxy, src_addr, dst_addr);
        let expected = [
            0x0D, 0x0A, 0x0D, 0x0A, 0x00, 0x0D, 0x0A, 0x51, 0x55, 0x49, 0x54,
            0x0A, // MAGIC header
            0x21, // Version 2 and command PROXY
            0x21, // family AF_UNIX with IPv6
            0x00, 0x24, // address sizes = 36
            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
            0x00, 0x01, // source address
            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
            0x00, 0x01, // destination address
            0x1F, 0x90, // source port
            0x10, 0x68,
        ];

        assert_eq!(&expected[..], &header.into_bytes()[..]);
    }
}