//! PROXY protocol header definition.

use std::{borrow::Cow, fmt, net::SocketAddr, str::from_utf8};

use thiserror::Error;
use tracing::instrument;

/// Protocol type
#[derive(Debug, PartialEq, Eq, Clone, Copy, Hash)]
pub enum Protocol {
    /// Stream protocol (TCP)
    Stream,
    /// Datagram protocol (UDP)
    Datagram,
}

/// Address information from a PROXY protocol header
#[derive(Debug, PartialEq, Eq, Clone, Hash)]
pub struct Address {
    /// Protocol type
    pub protocol: Protocol,
    /// Source address (of the actual client)
    pub source: SocketAddr,
    /// Destination address (of the proxy)
    pub destination: SocketAddr,
}

macro_rules! tlv {
    ($self:expr, $kind:ident) => {{
        $self.tlvs().find_map(|f| match f {
            Ok(Tlv::$kind(v)) => Some(v),
            _ => None,
        })
    }};
}

macro_rules! tlv_borrowed {
    ($self:expr, $kind:ident) => {{
        $self.tlvs().find_map(|f| match f {
            Ok(Tlv::$kind(v)) => match v {
                // It is more ergonomic to return the borrowed value directly rather
                // than it wrapped in a `Cow::Borrowed`. We know that tlvs always borrows
                // so we can safely unwrap the `Cow::Borrowed` and return the borrowed value.
                Cow::Owned(_) => unreachable!(),
                Cow::Borrowed(v) => Some(v),
            },
            _ => None,
        })
    }};
}

/// Iterator over PROXY protocol TLV fields
pub struct Tlvs<'a> {
    buf: &'a [u8],
}

#[expect(
    clippy::missing_trait_methods,
    reason = "we don't need to implement the other methods here"
)]
impl<'a> Iterator for Tlvs<'a> {
    type Item = Result<Tlv<'a>, ProxyProtocolError>;

    fn next(&mut self) -> Option<Self::Item> {
        if self.buf.is_empty() {
            return None;
        }

        let kind = self.buf[0];
        match self
            .buf
            .get(1..3)
            .map(|s| -> usize { u16::from_be_bytes(s.try_into().expect("infallible")).into() })
        {
            Some(u) if u + 3 <= self.buf.len() => {
                let (ret, new) = self.buf.split_at(3 + u);
                self.buf = new;

                Some(Tlv::decode(kind, &ret[3..]))
            }
            _ => {
                // Malformed TLV, cannot continue
                self.buf = &[];
                Some(Err(ProxyProtocolError::Invalid))
            }
        }
    }
}

/// Typed TLV field
///
/// Represents the currently known types of TLV fields from the PROXY protocol specification.
/// Non-recognized TLV fields are represented as [`Tlv::Custom`].
#[non_exhaustive]
#[derive(Debug, PartialEq, Eq, Clone)]
pub enum Tlv<'a> {
    /// Application-Layer Protocol Negotiation (ALPN). It is a byte sequence defining the upper
    /// layer protocol in use over the connection. The most common use case will be to pass the
    /// exact copy of the ALPN extension of the Transport Layer Security (TLS) protocol as defined
    /// by RFC7301.
    Alpn(Cow<'a, [u8]>),

    /// Contains the host name value passed by the client, as an UTF-8 encoded string. In case of
    /// TLS being used on the client connection, this is the exact copy of the `server_name`
    /// extension as defined by RFC3546, section 3.1, often referred to as SNI. There are probably
    /// other situations where an authority can be mentionned on a connection without TLS being
    /// involved at all.
    Authority(Cow<'a, str>),

    /// The value of the type `PP2_TYPE_CRC32C` is a 32-bit number storing the `CRC32c` checksum of
    /// the PROXY protocol header.
    ///
    /// When the checksum is supported by the sender after constructing the header the sender MUST:
    ///
    /// - initialize the checksum field to '0's.
    ///
    /// - calculate the `CRC32c` checksum of the PROXY header as described in RFC4960, Appendix B.
    ///
    /// - put the resultant value into the checksum field, and leave the rest of the bits unchanged.
    ///
    /// If the checksum is provided as part of the PROXY header and the checksum functionality is
    /// supported by the receiver, the receiver MUST:
    ///
    /// - store the received `CRC32c` checksum value aside.
    ///
    /// - replace the 32 bits of the checksum field in the received PROXY header with all '0's and
    ///   calculate a `CRC32c` checksum value of the whole PROXY header.
    ///
    /// - verify that the calculated `CRC32c` checksum is the same as the received `CRC32c`
    ///   checksum. If it is not, the receiver MUST treat the TCP connection providing the header as
    ///   invalid.
    ///
    /// The default procedure for handling an invalid TCP connection is to abort it.
    Crc32c(u32),

    /// The TLV of this type should be ignored when parsed. The value is zero or more bytes. Can be
    /// used for data padding or alignment. Note that it can be used to align only by 3 or more
    /// bytes because a TLV can not be smaller than that.
    Noop(usize),

    /// The value of the type `PP2_TYPE_UNIQUE_ID` is an opaque byte sequence of up to
    /// 128 bytes generated by the upstream proxy that uniquely identifies the connection.
    ///
    /// The unique ID can be used to easily correlate connections across multiple layers of
    /// proxies, without needing to look up IP addresses and port numbers.
    UniqueId(Cow<'a, [u8]>),

    /// SSL (TLS) information
    ///
    /// See [`SslInfo`] for more information.
    Ssl(SslInfo<'a>),

    /// The type `PP2_TYPE_NETNS` defines the value as the US-ASCII string representation of the
    /// namespace's name.
    Netns(Cow<'a, str>),

    // The following can only appear as a sub-TLV of SslInfo
    /// SSL/TLS version
    SslVersion(Cow<'a, str>),
    /// In all cases, the string representation (in UTF8) of the Common Name field (OID: 2.5.4.3)
    /// of the client certificate's Distinguished Name, is appended using the TLV format and the
    /// type `PP2_SUBTYPE_SSL_CN`. E.g. "example.com".
    SslCn(Cow<'a, str>),
    /// The second level TLV `PP2_SUBTYPE_SSL_CIPHER` provides the US-ASCII string name of the used
    /// cipher, for example "ECDHE-RSA-AES128-GCM-SHA256".
    SslCipher(Cow<'a, str>),
    /// The second level TLV `PP2_SUBTYPE_SSL_SIG_ALG` provides the US-ASCII string name of the
    /// algorithm used to sign the certificate presented by the frontend when the incoming
    /// connection was made over an SSL/TLS transport layer, for example "SHA256".
    SslSigAlg(Cow<'a, str>),
    /// The second level TLV `PP2_SUBTYPE_SSL_KEY_ALG` provides the US-ASCII string name of the
    /// algorithm used to generate the key of the certificate presented by the frontend when the
    /// incoming connection was made over an SSL/TLS transport layer, for example "RSA2048".
    SslKeyAlg(Cow<'a, str>),

    /// Unrecognized or custom TLV field
    Custom(u8, Cow<'a, [u8]>),
}

impl<'a> Tlv<'a> {
    fn decode(kind: u8, data: &'a [u8]) -> Result<Self, ProxyProtocolError> {
        match kind {
            0x01 => Ok(Self::Alpn(data.into())),
            0x02 => Ok(Self::Authority(
                from_utf8(data)
                    .map_err(|_| ProxyProtocolError::Invalid)?
                    .into(),
            )),
            0x03 => Ok(Self::Crc32c(u32::from_be_bytes(
                data.try_into().map_err(|_| ProxyProtocolError::Invalid)?,
            ))),
            0x04 => Ok(Self::Noop(data.len())),
            0x05 => Ok(Self::UniqueId(data.into())),
            0x20 => Ok(Tlv::Ssl(SslInfo(
                *data.first().ok_or(ProxyProtocolError::Invalid)?,
                u32::from_be_bytes(
                    data.get(1..5)
                        .ok_or(ProxyProtocolError::Invalid)?
                        .try_into()
                        .map_err(|_| ProxyProtocolError::Invalid)?,
                ),
                data.get(5..).ok_or(ProxyProtocolError::Invalid)?.into(),
            ))),
            0x21 => Ok(Self::SslVersion(
                from_utf8(data)
                    .map_err(|_| ProxyProtocolError::Invalid)?
                    .into(),
            )),
            0x22 => Ok(Self::SslCn(
                from_utf8(data)
                    .map_err(|_| ProxyProtocolError::Invalid)?
                    .into(),
            )),
            0x23 => Ok(Self::SslCipher(
                from_utf8(data)
                    .map_err(|_| ProxyProtocolError::Invalid)?
                    .into(),
            )),
            0x24 => Ok(Self::SslSigAlg(
                from_utf8(data)
                    .map_err(|_| ProxyProtocolError::Invalid)?
                    .into(),
            )),
            0x25 => Ok(Self::SslKeyAlg(
                from_utf8(data)
                    .map_err(|_| ProxyProtocolError::Invalid)?
                    .into(),
            )),
            0x30 => Ok(Self::Netns(
                from_utf8(data)
                    .map_err(|_| ProxyProtocolError::Invalid)?
                    .into(),
            )),
            t => Ok(Self::Custom(t, data.into())),
        }
    }

    pub fn into_owned(self) -> Tlv<'static> {
        match self {
            Self::Alpn(v) => Tlv::Alpn(Cow::Owned(v.into_owned())),
            Self::Authority(v) => Tlv::Authority(Cow::Owned(v.into_owned())),
            Self::Crc32c(v) => Tlv::Crc32c(v),
            Self::Noop(v) => Tlv::Noop(v),
            Self::UniqueId(v) => Tlv::UniqueId(Cow::Owned(v.into_owned())),
            Self::Ssl(v) => Tlv::Ssl(v.into_owned()),
            Self::Netns(v) => Tlv::Netns(Cow::Owned(v.into_owned())),
            Self::SslVersion(v) => Tlv::SslVersion(Cow::Owned(v.into_owned())),
            Self::SslCn(v) => Tlv::SslCn(Cow::Owned(v.into_owned())),
            Self::SslCipher(v) => Tlv::SslCipher(Cow::Owned(v.into_owned())),
            Self::SslSigAlg(v) => Tlv::SslSigAlg(Cow::Owned(v.into_owned())),
            Self::SslKeyAlg(v) => Tlv::SslKeyAlg(Cow::Owned(v.into_owned())),
            Self::Custom(a, v) => Tlv::Custom(a, Cow::Owned(v.into_owned())),
        }
    }
}

/// SSL information from a PROXY protocol header
#[derive(PartialEq, Eq, Clone)]
pub struct SslInfo<'a>(u8, u32, Cow<'a, [u8]>);

impl SslInfo<'_> {
    /// Client connected over SSL/TLS
    ///
    /// The `PP2_CLIENT_SSL` flag indicates that the client connected over SSL/TLS. When this field
    /// is present, the US-ASCII string representation of the TLS version is appended at the end of
    /// the field in the TLV format using the type `PP2_SUBTYPE_SSL_VERSION`.
    pub fn client_ssl(&self) -> bool {
        self.0 & 0x01 != 0
    }

    /// Client certificate presented in the connection
    ///
    /// `PP2_CLIENT_CERT_CONN` indicates that the client provided a certificate over the current
    /// connection.
    pub fn client_cert_conn(&self) -> bool {
        self.0 & 0x02 != 0
    }

    /// Client certificate presented in the session
    ///
    /// `PP2_CLIENT_CERT_SESS` indicates that the client provided a certificate at least once over
    /// the TLS session this connection belongs to.
    pub fn client_cert_sess(&self) -> bool {
        self.0 & 0x04 != 0
    }

    /// Whether the certificate was verified
    ///
    /// The verify field will be zero if the client presented a certificate and it was successfully
    /// verified, and non-zero otherwise.
    pub fn verify(&self) -> u32 {
        self.1
    }

    /// Iterator over all TLV (type-length-value) fields
    pub fn tlvs(&self) -> Tlvs<'_> {
        Tlvs { buf: &self.2 }
    }

    // Convenience accessors for common TLVs

    /// SSL version
    ///
    /// See [`Tlv::SslVersion`] for more information.
    pub fn version(&self) -> Option<&str> {
        tlv_borrowed!(self, SslVersion)
    }

    /// SSL CN
    ///
    /// See [`Tlv::SslCn`] for more information.
    pub fn cn(&self) -> Option<&str> {
        tlv_borrowed!(self, SslCn)
    }

    /// SSL cipher
    ///
    /// See [`Tlv::SslCipher`] for more information.
    pub fn cipher(&self) -> Option<&str> {
        tlv_borrowed!(self, SslCipher)
    }

    /// SSL signature algorithm
    ///
    /// See [`Tlv::SslSigAlg`] for more information.
    pub fn sig_alg(&self) -> Option<&str> {
        tlv_borrowed!(self, SslSigAlg)
    }

    /// SSL key algorithm
    ///
    /// See [`Tlv::SslKeyAlg`] for more information.
    pub fn key_alg(&self) -> Option<&str> {
        tlv_borrowed!(self, SslKeyAlg)
    }

    /// Returns an owned version of this struct
    pub fn into_owned(self) -> SslInfo<'static> {
        SslInfo(self.0, self.1, Cow::Owned(self.2.into_owned()))
    }
}

impl fmt::Debug for SslInfo<'_> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("SslInfo")
            .field("verify", &self.verify())
            .field("client_ssl", &self.client_ssl())
            .field("client_cert_conn", &self.client_cert_conn())
            .field("client_cert_sess", &self.client_cert_sess())
            .field("fields", &self.tlvs().collect::<Vec<_>>())
            .finish()
    }
}

/// A PROXY protocol header
#[derive(Default, PartialEq, Eq, Clone, Debug)]
pub struct Header<'a>(pub(super) Option<Address>, pub(super) Cow<'a, [u8]>);

impl<'a> Header<'a> {
    /// Attempt to parse a PROXY protocol header from the given buffer
    ///
    /// Returns the parsed header and the number of bytes consumed from the buffer. If the header
    /// is incomplete, returns [`ProxyProtocolError::BufferTooShort`] so more data can be read from
    /// the socket.
    ///
    /// If the header is malformed or unsupported, returns [`ProxyProtocolError::Invalid`].
    ///
    /// This function will borrow the buffer for the lifetime of the returned header. If
    /// you need to keep the header around for longer than the buffer, use
    /// [`Header::into_owned`].
    #[instrument(skip_all)]
    pub fn parse(buf: &'a [u8]) -> Result<(Self, usize), ProxyProtocolError> {
        match buf.first() {
            Some(b'P') => super::v1::decode(buf),
            Some(b'\r') => super::v2::decode(buf),
            None => Err(ProxyProtocolError::BufferTooShort),
            _ => Err(ProxyProtocolError::Invalid),
        }
    }

    /// Proxied address information
    ///
    /// If `None`, this indicates so-called "local" mode, where the connection is not proxied.
    /// This is usually the case when the connection is initiated by the proxy itself, e.g. for
    /// health checks.
    pub fn proxied_address(&self) -> Option<&Address> {
        self.0.as_ref()
    }

    /// Iterator that yields all extension TLV (type-length-value) fields present in the header
    ///
    /// See [`Tlv`] for more information on the different types of TLV fields.
    pub fn tlvs(&self) -> Tlvs<'_> {
        Tlvs { buf: &self.1 }
    }

    // Convenience accessors for common fields

    /// Raw ALPN extension data
    ///
    /// See [`Tlv::Alpn`] for more information.
    pub fn alpn(&self) -> Option<&[u8]> {
        tlv_borrowed!(self, Alpn)
    }

    /// Authority - typically the hostname of the client (SNI)
    ///
    /// See [`Tlv::Authority`] for more information.
    pub fn authority(&self) -> Option<&str> {
        tlv_borrowed!(self, Authority)
    }

    /// `CRC32c` checksum of the address information
    ///
    /// See [`Tlv::Crc32c`] for more information.
    pub fn crc32c(&self) -> Option<u32> {
        tlv!(self, Crc32c)
    }

    /// Unique ID of the connection
    ///
    /// See [`Tlv::UniqueId`] for more information.
    pub fn unique_id(&self) -> Option<&[u8]> {
        tlv_borrowed!(self, UniqueId)
    }

    /// SSL information
    ///
    /// See [`Tlv::Ssl`] for more information.
    pub fn ssl(&self) -> Option<SslInfo<'_>> {
        tlv!(self, Ssl)
    }

    /// Network namespace
    ///
    /// See [`Tlv::Netns`] for more information.
    pub fn netns(&self) -> Option<&str> {
        tlv_borrowed!(self, Netns)
    }

    /// Returns an owned version of this struct
    pub fn into_owned(self) -> Header<'static> {
        Header(self.0, Cow::Owned(self.1.into_owned()))
    }
}

#[derive(Debug, PartialEq, Eq, Error)]
pub enum ProxyProtocolError {
    #[error("The buffer is too short to contain a complete PROXY protocol header")]
    BufferTooShort,
    #[error("The PROXY protocol header is malformed")]
    Invalid,
}

#[cfg(test)]
mod tests {
    use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};

    use super::*;

    const V1_UNKNOWN: &[u8] = b"PROXY UNKNOWN\r\n";

    const V1_TCPV4: &[u8] = b"PROXY TCP4 127.0.0.1 192.168.0.1 12345 443\r\n";
    const V1_TCPV6: &[u8] = b"PROXY TCP6 2001:db8::1 ::1 12345 443\r\n";

    const V2_LOCAL: &[u8] =
        b"\r\n\r\n\0\r\nQUIT\n \0\0\x0f\x03\0\x04\x88\x9d\xa1\xdf \0\x05\0\0\0\0\0";

    const V2_TCPV4: &[u8] = &[
        13, 10, 13, 10, 0, 13, 10, 81, 85, 73, 84, 10, 33, 17, 0, 12, 127, 0, 0, 1, 192, 168, 0, 1,
        48, 57, 1, 187,
    ];
    const V2_TCPV6: &[u8] = &[
        13, 10, 13, 10, 0, 13, 10, 81, 85, 73, 84, 10, 33, 33, 0, 36, 32, 1, 13, 184, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 48, 57, 1, 187,
    ];
    const V2_TCPV4_TLV: &[u8] = &[
        13, 10, 13, 10, 0, 13, 10, 81, 85, 73, 84, 10, 33, 17, 0, 104, 127, 0, 0, 1, 192, 168, 0,
        1, 48, 57, 1, 187, 3, 0, 4, 211, 153, 216, 216, 5, 0, 4, 49, 50, 51, 52, 32, 0, 75, 7, 0,
        0, 0, 0, 33, 0, 7, 84, 76, 83, 118, 49, 46, 51, 34, 0, 9, 108, 111, 99, 97, 108, 104, 111,
        115, 116, 37, 0, 7, 82, 83, 65, 52, 48, 57, 54, 36, 0, 10, 82, 83, 65, 45, 83, 72, 65, 50,
        53, 54, 35, 0, 22, 84, 76, 83, 95, 65, 69, 83, 95, 50, 53, 54, 95, 71, 67, 77, 95, 83, 72,
        65, 51, 56, 52,
    ];

    #[test]
    fn parse_proxy_header_too_short() {
        for case in &[
            V1_TCPV4,
            V1_TCPV6,
            V1_UNKNOWN,
            V2_TCPV4,
            V2_TCPV6,
            V2_TCPV4_TLV,
            V2_LOCAL,
        ] {
            for i in 0..case.len() {
                assert!(matches!(
                    Header::parse(&case[..i]),
                    Err(ProxyProtocolError::BufferTooShort)
                ));
            }

            Header::parse(case).expect("failed to parse header");
        }
    }

    #[test]
    fn parse_proxy_header_v1_unterminated() {
        let line = b"PROXY TCP4 THISISSTORYALLABOUTHOWMYLIFEGOTFLIPPEDTURNEDUPSIDEDOWNANDIDLIKETOTAKEAMINUTEJUSTSITRIGHTTHEREANDILLTELLYOUHOWIGOTTHEPRINCEOFAIR";
        assert!(matches!(
            Header::parse(line),
            Err(ProxyProtocolError::Invalid)
        ));
    }

    #[test]
    fn parse_proxy_header_v1() {
        let (res, consumed) = Header::parse(V1_TCPV4).expect("failed to parse");
        assert_eq!(consumed, V1_TCPV4.len());
        assert_eq!(
            res.0,
            Some(Address {
                protocol: Protocol::Stream,
                source: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 12345),
                destination: SocketAddr::new(IpAddr::V4(Ipv4Addr::new(192, 168, 0, 1)), 443),
            })
        );
        assert_eq!(res.1, vec![0; 0]);

        let (res, consumed) = Header::parse(V1_TCPV6).expect("failed to parse");
        assert_eq!(consumed, V1_TCPV6.len());
        assert_eq!(
            res.0,
            Some(Address {
                protocol: Protocol::Stream,
                source: SocketAddr::new(
                    IpAddr::V6(Ipv6Addr::new(0x2001, 0xdb8, 0, 0, 0, 0, 0, 1)),
                    12345
                ),
                destination: SocketAddr::new(IpAddr::V6(Ipv6Addr::LOCALHOST), 443),
            })
        );
        assert_eq!(res.1, vec![0; 0]);
    }

    #[test]
    fn parse_proxy_header_v2() {
        let (res, consumed) = Header::parse(V2_LOCAL).expect("failed to parse");
        assert_eq!(consumed, V2_LOCAL.len());
        assert_eq!(res.0, None);

        let (res, consumed) = Header::parse(V2_TCPV4).expect("failed to parse");
        assert_eq!(consumed, V2_TCPV4.len());
        assert_eq!(
            res.0,
            Some(Address {
                protocol: Protocol::Stream,
                source: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 12345),
                destination: SocketAddr::new(IpAddr::V4(Ipv4Addr::new(192, 168, 0, 1)), 443),
            })
        );
        assert_eq!(res.1, vec![0; 0]);

        let (res, consumed) = Header::parse(V2_TCPV6).expect("failed to parse");
        assert_eq!(consumed, V2_TCPV6.len());
        assert_eq!(
            res.0,
            Some(Address {
                protocol: Protocol::Stream,
                source: SocketAddr::new(
                    IpAddr::V6(Ipv6Addr::new(0x2001, 0xdb8, 0, 0, 0, 0, 0, 1)),
                    12345
                ),
                destination: SocketAddr::new(IpAddr::V6(Ipv6Addr::LOCALHOST), 443),
            })
        );
        assert_eq!(res.1, vec![0; 0]);
    }

    #[test]
    fn parse_proxy_header_v2_with_tlvs() {
        let (res, _) = Header::parse(V2_TCPV4_TLV).expect("failed to parse");

        let mut fields = res.tlvs();
        assert_eq!(fields.next(), Some(Ok(Tlv::Crc32c(0xd399_d8d8))));
        assert_eq!(fields.next(), Some(Ok(Tlv::UniqueId(b"1234"[..].into()))));

        let ssl = fields
            .next()
            .expect("next tlv missing")
            .expect("tlv parsing failed");
        let Tlv::Ssl(ssl) = ssl else {
            panic!("expected SSL TLV");
        };

        assert_eq!(ssl.verify(), 0);
        assert!(ssl.client_ssl());
        assert!(ssl.client_cert_conn());
        assert!(ssl.client_cert_sess());

        let mut f = ssl.tlvs();

        assert_eq!(f.next(), Some(Ok(Tlv::SslVersion("TLSv1.3".into()))));
        assert_eq!(f.next(), Some(Ok(Tlv::SslCn("localhost".into()))));
        assert_eq!(f.next(), Some(Ok(Tlv::SslKeyAlg("RSA4096".into()))));
        assert_eq!(f.next(), Some(Ok(Tlv::SslSigAlg("RSA-SHA256".into()))));
        assert_eq!(
            f.next(),
            Some(Ok(Tlv::SslCipher("TLS_AES_256_GCM_SHA384".into())))
        );
        assert!(f.next().is_none());

        assert!(fields.next().is_none());
    }
}
