haproxy_protocol/

lib.rs

#![deny(warnings)]
#![warn(unused_extern_crates)]
#![deny(clippy::todo)]
#![deny(clippy::unimplemented)]
#![deny(clippy::unwrap_used)]
#![deny(clippy::expect_used)]
#![deny(clippy::panic)]
#![deny(clippy::unreachable)]
#![deny(clippy::await_holding_lock)]
#![deny(clippy::needless_pass_by_value)]
#![deny(clippy::trivially_copy_pass_by_ref)]

use crate::parse::{parse_proxy_hdr_v1, parse_proxy_hdr_v2};
use std::num::NonZeroUsize;

#[cfg(feature = "tokio")]
use crate::parse::{V1_MAX_LEN, V1_MIN_LEN};

const NZ_ONE: NonZeroUsize = NonZeroUsize::new(1).expect("Invalid compile time constant");

mod parse;

#[derive(Debug, PartialEq, Eq, Clone, Copy)]
#[repr(u8)]
enum Protocol {
    Unspec = 0x00,
    TcpV4 = 0x11,
    UdpV4 = 0x12,
    TcpV6 = 0x21,
    UdpV6 = 0x22,
    // UnixStream = 0x31,
    // UnixDgram = 0x32,
}

#[derive(Debug, PartialEq, Eq, Clone, Copy)]
#[repr(u8)]
enum Command {
    Local = 0x00,
    Proxy = 0x01,
}

#[derive(Debug, PartialEq, Eq, Clone)]
enum Address {
    None,
    V4 {
        src: std::net::SocketAddrV4,
        dst: std::net::SocketAddrV4,
    },
    V6 {
        src: std::net::SocketAddrV6,
        dst: std::net::SocketAddrV6,
    },
    // Unix {
    //     src: PathBuf,
    //     dst: PathBuf,
    // }
}

#[derive(Debug, Clone)]
pub enum RemoteAddress {
    Local,
    Invalid,
    TcpV4 {
        src: std::net::SocketAddrV4,
        dst: std::net::SocketAddrV4,
    },
    UdpV4 {
        src: std::net::SocketAddrV4,
        dst: std::net::SocketAddrV4,
    },
    TcpV6 {
        src: std::net::SocketAddrV6,
        dst: std::net::SocketAddrV6,
    },
    UdpV6 {
        src: std::net::SocketAddrV6,
        dst: std::net::SocketAddrV6,
    },
}

#[derive(Debug)]
pub enum Error {
    Incomplete { need: NonZeroUsize },
    Invalid,
    UnableToComplete,
}

#[derive(Debug, Clone)]
pub struct ProxyHdrV2 {
    command: Command,
    protocol: Protocol,
    // address_family: AddressFamily,
    // length: u16,
    address: Address,
}

impl ProxyHdrV2 {
    pub fn parse(input_data: &[u8]) -> Result<(usize, Self), Error> {
        match parse_proxy_hdr_v2(input_data) {
            Ok((remainder, hdr)) => {
                let took = input_data.len() - remainder.len();
                Ok((took, hdr))
            }
            Err(nom::Err::Incomplete(nom::Needed::Size(need))) => Err(Error::Incomplete { need }),
            // We always know exactly how much is needed for hdr v2
            Err(nom::Err::Incomplete(nom::Needed::Unknown)) => Err(Error::UnableToComplete),

            Err(nom::Err::Error(err)) => {
                tracing::error!(?err);
                Err(Error::Invalid)
            }
            Err(nom::Err::Failure(err)) => {
                tracing::error!(?err);
                Err(Error::Invalid)
            }
        }
    }

    pub fn to_remote_addr(self) -> RemoteAddress {
        match (self.command, self.protocol, self.address) {
            (Command::Local, _, _) => RemoteAddress::Local,
            (Command::Proxy, Protocol::TcpV4, Address::V4 { src, dst }) => {
                RemoteAddress::TcpV4 { src, dst }
            }
            (Command::Proxy, Protocol::UdpV4, Address::V4 { src, dst }) => {
                RemoteAddress::UdpV4 { src, dst }
            }
            (Command::Proxy, Protocol::TcpV6, Address::V6 { src, dst }) => {
                RemoteAddress::TcpV6 { src, dst }
            }
            (Command::Proxy, Protocol::UdpV6, Address::V6 { src, dst }) => {
                RemoteAddress::UdpV6 { src, dst }
            }
            _ => RemoteAddress::Invalid,
        }
    }
}

#[derive(Debug, Clone)]
pub struct ProxyHdrV1 {
    protocol: Protocol,
    address: Address,
}

impl ProxyHdrV1 {
    pub fn parse(input_data: &[u8]) -> Result<(usize, Self), Error> {
        match parse_proxy_hdr_v1(input_data) {
            Ok((remainder, hdr)) => {
                let took = input_data.len() - remainder.len();
                Ok((took, hdr))
            }
            Err(nom::Err::Incomplete(nom::Needed::Size(need))) => Err(Error::Incomplete { need }),
            // We aren't sure how much we need but we need *something*.
            Err(nom::Err::Incomplete(nom::Needed::Unknown)) => {
                Err(Error::Incomplete { need: NZ_ONE })
            }

            Err(nom::Err::Error(err)) => {
                tracing::error!(?err);
                Err(Error::Invalid)
            }
            Err(nom::Err::Failure(err)) => {
                tracing::error!(?err);
                Err(Error::Invalid)
            }
        }
    }

    pub fn to_remote_addr(self) -> RemoteAddress {
        match (self.protocol, self.address) {
            (Protocol::TcpV4, Address::V4 { src, dst }) => RemoteAddress::TcpV4 { src, dst },
            (Protocol::UdpV4, Address::V4 { src, dst }) => RemoteAddress::UdpV4 { src, dst },
            (Protocol::TcpV6, Address::V6 { src, dst }) => RemoteAddress::TcpV6 { src, dst },
            (Protocol::UdpV6, Address::V6 { src, dst }) => RemoteAddress::UdpV6 { src, dst },
            _ => RemoteAddress::Invalid,
        }
    }
}

#[cfg(feature = "tokio")]
#[derive(Debug)]
pub enum AsyncReadError {
    Io(std::io::Error),
    Invalid,
    UnableToComplete,
    RequestTooLarge,
    InconsistentRead,
}

#[cfg(feature = "tokio")]
impl ProxyHdrV2 {
    pub async fn parse_from_read<S>(mut stream: S) -> Result<(S, Self), AsyncReadError>
    where
        S: tokio::io::AsyncReadExt + std::marker::Unpin,
    {
        use tracing::{debug, error};

        const HDR_SIZE_LIMIT: usize = 512;

        let mut buf = vec![0; 16];

        // First we need to read the exact amount to get up to the *length* field. This will
        // let us then proceed to parse the early header and return how much we need to continue
        // to read.
        let mut took = stream
            .read_exact(&mut buf)
            .await
            .map_err(AsyncReadError::Io)?;

        match ProxyHdrV2::parse(&buf) {
            // Okay, we got a valid header - this can occur with proxy for local conditions.
            Ok((_, hdr)) => return Ok((stream, hdr)),
            // We need more bytes, this is the precise amount we need.
            Err(Error::Incomplete { need }) => {
                let resize_to = buf.len() + usize::from(need);
                // Limit the amount so that we don't overflow anything or allocate a buffer that
                // is too large. Nice try hackers.
                if resize_to > HDR_SIZE_LIMIT {
                    error!(
                        "proxy header request was larger than {} bytes, refusing to proceed.",
                        HDR_SIZE_LIMIT
                    );
                    return Err(AsyncReadError::RequestTooLarge);
                }
                buf.resize(resize_to, 0);
            }
            Err(Error::Invalid) => {
                debug!(proxy_binary_dump = %hex::encode(&buf));
                error!("proxy header was invalid");
                return Err(AsyncReadError::Invalid);
            }
            Err(Error::UnableToComplete) => {
                debug!(proxy_binary_dump = %hex::encode(&buf));
                error!("proxy header was incomplete");
                return Err(AsyncReadError::UnableToComplete);
            }
        };

        // Now read any remaining bytes into the buffer.
        took += stream
            .read_exact(&mut buf[16..])
            .await
            .map_err(AsyncReadError::Io)?;

        match ProxyHdrV2::parse(&buf) {
            Ok((hdr_took, _)) if hdr_took != took => {
                // We took inconsistent byte amounts, error.
                error!("proxy header read an inconsistent amount from stream.");
                Err(AsyncReadError::InconsistentRead)
            }
            Ok((_, hdr)) =>
            // HAPPY!!!!!
            {
                Ok((stream, hdr))
            }
            Err(Error::Incomplete { need: _ }) => {
                error!("proxy header could not be read to the end.");
                Err(AsyncReadError::UnableToComplete)
            }
            Err(Error::Invalid) => {
                debug!(proxy_binary_dump = %hex::encode(&buf));
                error!("proxy header was invalid");
                Err(AsyncReadError::Invalid)
            }
            Err(Error::UnableToComplete) => {
                debug!(proxy_binary_dump = %hex::encode(&buf));
                error!("proxy header was incomplete");
                Err(AsyncReadError::UnableToComplete)
            }
        }
    }
}

#[cfg(feature = "tokio")]
impl ProxyHdrV1 {
    pub async fn parse_from_read<S>(mut stream: S) -> Result<(S, Self), AsyncReadError>
    where
        S: tokio::io::AsyncReadExt + std::marker::Unpin,
    {
        use tracing::{debug, error};

        // This is the maximum size of the buffer we could possibly need.
        let mut buf = [0; V1_MAX_LEN + 1];

        // First we need to read the exact amount to get up to the *length* field. This will
        // let us then proceed to parse the early header and return how much we need to continue
        // to read.
        let mut took = stream
            .read_exact(&mut buf[..V1_MIN_LEN])
            .await
            .map_err(AsyncReadError::Io)?;

        loop {
            match ProxyHdrV1::parse(&buf) {
                Ok((hdr_took, _)) if hdr_took != took => {
                    // We took inconsistent byte amounts, error.
                    error!("proxy header read an inconsistent amount from stream.");
                    return Err(AsyncReadError::InconsistentRead);
                }
                Ok((_, hdr)) =>
                // HAPPY!!!!!
                {
                    return Ok((stream, hdr));
                }
                Err(Error::Incomplete { need }) => {
                    // We need more data, read it and then continue the loop.
                    // Now read any remaining bytes into the buffer.
                    took += stream
                        .read_exact(&mut buf[took..took + need.get()])
                        .await
                        .map_err(AsyncReadError::Io)?;

                    continue;
                }
                Err(Error::Invalid) => {
                    debug!(proxy_binary_dump = %hex::encode(buf));
                    error!("proxy header was invalid");
                    return Err(AsyncReadError::Invalid);
                }
                Err(Error::UnableToComplete) => {
                    debug!(proxy_binary_dump = %hex::encode(buf));
                    error!("proxy header was incomplete");
                    return Err(AsyncReadError::UnableToComplete);
                }
            }
        } // end loop
    }
}