socker 1.0.0

//! The four SOCKS5 exchange messages: the client greeting and server method
//! choice, and the request/response pair.

use super::{
    VERSION,
    types::{
        Address,
        AuthenticationMethod,
        CommandType,
        Reply,
    },
};
use crate::{
    DecodeError,
    DecodeStatus,
    EncodeError,
    Message,
};

/// The client's greeting (`VER NMETHODS METHODS`), the first message of a
/// SOCKS5 negotiation.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ClientGreeting {
    /// The authentication methods the client is willing to use, in order of
    /// preference. May include unassigned or private-use values; may be empty
    /// only if the peer tolerates it, as the specification expects at least
    /// one method.
    pub methods: Box<[AuthenticationMethod]>,
}

impl Message for ClientGreeting {
    fn encoded_len(&self) -> usize {
        self.methods.len().saturating_add(2)
    }

    fn encode(&self, buffer: &mut Vec<u8>) -> Result<(), EncodeError> {
        let count = u8::try_from(self.methods.len())
            .map_err(|_error| EncodeError::TooLong("more than 255 authentication methods"))?;
        buffer.reserve(self.encoded_len());
        buffer.push(VERSION);
        buffer.push(count);
        buffer.extend(self.methods.iter().map(|&method| u8::from(method)));
        Ok(())
    }

    fn decode(source: &[u8]) -> Result<DecodeStatus<(Self, usize)>, DecodeError> {
        let Some((&version, after_version)) = source.split_first() else {
            return Ok(DecodeStatus::Partial);
        };
        if version != VERSION {
            return Err(DecodeError::InvalidVersion {
                expected: VERSION,
                actual:   version,
            });
        }
        let Some((&count, rest)) = after_version.split_first() else {
            return Ok(DecodeStatus::Partial);
        };
        let count = usize::from(count);
        let Some(method_bytes) = rest.get(.. count) else {
            return Ok(DecodeStatus::Partial);
        };
        let methods = method_bytes
            .iter()
            .map(|&byte| AuthenticationMethod::from(byte))
            .collect::<Vec<_>>()
            .into_boxed_slice();
        Ok(DecodeStatus::Complete((
            Self {
                methods,
            },
            count.saturating_add(2),
        )))
    }
}

/// The server's method selection (`VER METHOD`), the response to
/// [`ClientGreeting`].
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ServerChoice {
    /// The authentication method the server chose; `0xFF`
    /// ([`AuthenticationMethod::NO_ACCEPTABLE_METHODS`]) signals that none of
    /// the offered methods were acceptable.
    pub method: AuthenticationMethod,
}

impl Message for ServerChoice {
    fn encoded_len(&self) -> usize {
        2
    }

    fn encode(&self, buffer: &mut Vec<u8>) -> Result<(), EncodeError> {
        buffer.reserve(2);
        buffer.push(VERSION);
        buffer.push(u8::from(self.method));
        Ok(())
    }

    fn decode(source: &[u8]) -> Result<DecodeStatus<(Self, usize)>, DecodeError> {
        let Some((&version, after_version)) = source.split_first() else {
            return Ok(DecodeStatus::Partial);
        };
        if version != VERSION {
            return Err(DecodeError::InvalidVersion {
                expected: VERSION,
                actual:   version,
            });
        }
        let Some((&method, _rest)) = after_version.split_first() else {
            return Ok(DecodeStatus::Partial);
        };
        Ok(DecodeStatus::Complete((
            Self {
                method: AuthenticationMethod::from(method),
            },
            2,
        )))
    }
}

/// A SOCKS5 request (`VER CMD RSV ATYP DST.ADDR DST.PORT`), sent by the
/// client after the negotiation completes.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Request {
    /// The command to perform (`CONNECT`, `BIND`, `UDP_ASSOCIATE`, or any
    /// unassigned value).
    pub command: CommandType,

    /// The target address; a domain name defers DNS resolution to the server
    /// (the SOCKS5h mode).
    pub address: Address,

    /// The target port, in the usual byte order of the surrounding payload.
    pub port: u16,
}

impl Message for Request {
    fn encoded_len(&self) -> usize {
        self.address.encoded_len().saturating_add(5)
    }

    fn encode(&self, buffer: &mut Vec<u8>) -> Result<(), EncodeError> {
        buffer.reserve(self.encoded_len());
        buffer.push(VERSION);
        buffer.push(u8::from(self.command));
        buffer.push(0x00); // RSV
        self.address.encode_into(buffer)?;
        buffer.extend_from_slice(&self.port.to_be_bytes());
        Ok(())
    }

    fn decode(source: &[u8]) -> Result<DecodeStatus<(Self, usize)>, DecodeError> {
        let Some((&version, after_version)) = source.split_first() else {
            return Ok(DecodeStatus::Partial);
        };
        if version != VERSION {
            return Err(DecodeError::InvalidVersion {
                expected: VERSION,
                actual:   version,
            });
        }
        let Some((&command, after_command)) = after_version.split_first() else {
            return Ok(DecodeStatus::Partial);
        };
        let Some((&reserved, rest)) = after_command.split_first() else {
            return Ok(DecodeStatus::Partial);
        };
        if reserved != 0x00 {
            return Err(DecodeError::Malformed("reserved byte is not zero"));
        }
        let (address, consumed) = match Address::decode_from(rest)? {
            | DecodeStatus::Complete(complete) => complete,
            | DecodeStatus::Partial => return Ok(DecodeStatus::Partial),
        };
        let Some(after_address) = rest.get(consumed ..) else {
            return Ok(DecodeStatus::Partial);
        };
        let Some(&[port_high, port_low]) = after_address.first_chunk::<2>() else {
            return Ok(DecodeStatus::Partial);
        };
        Ok(DecodeStatus::Complete((
            Self {
                command: CommandType::from(command),
                address,
                port: u16::from_be_bytes([port_high, port_low]),
            },
            consumed.saturating_add(5),
        )))
    }
}

/// A SOCKS5 response (`VER REP RSV ATYP BND.ADDR BND.PORT`), sent by the
/// server in answer to a [`Request`].
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Response {
    /// The reply code; anything but [`Reply::SUCCESS`] describes a failure.
    pub reply: Reply,

    /// The server-bound address reported back to the client.
    pub address: Address,

    /// The server-bound port reported back to the client.
    pub port: u16,
}

impl Response {
    /// A `COMMAND_NOT_SUPPORTED` failure response with an unspecified bound
    /// address.
    pub const COMMAND_NOT_SUPPORTED: Self = Self::error_reply(Reply::COMMAND_NOT_SUPPORTED);
    /// A `HOST_UNREACHABLE` failure response with an unspecified bound
    /// address.
    pub const HOST_UNREACHABLE: Self = Self::error_reply(Reply::HOST_UNREACHABLE);

    /// A response with the given reply code and an unspecified IPv4 address
    /// and port, the conventional shape of pure failure replies.
    #[must_use]
    pub const fn error_reply(reply: Reply) -> Self {
        Self {
            reply,
            address: Address::Ipv4(std::net::Ipv4Addr::UNSPECIFIED),
            port: 0,
        }
    }
}

impl Message for Response {
    fn encoded_len(&self) -> usize {
        self.address.encoded_len().saturating_add(5)
    }

    fn encode(&self, buffer: &mut Vec<u8>) -> Result<(), EncodeError> {
        buffer.reserve(self.encoded_len());
        buffer.push(VERSION);
        buffer.push(u8::from(self.reply));
        buffer.push(0x00); // RSV
        self.address.encode_into(buffer)?;
        buffer.extend_from_slice(&self.port.to_be_bytes());
        Ok(())
    }

    fn decode(source: &[u8]) -> Result<DecodeStatus<(Self, usize)>, DecodeError> {
        let Some((&version, after_version)) = source.split_first() else {
            return Ok(DecodeStatus::Partial);
        };
        if version != VERSION {
            return Err(DecodeError::InvalidVersion {
                expected: VERSION,
                actual:   version,
            });
        }
        let Some((&reply, after_reply)) = after_version.split_first() else {
            return Ok(DecodeStatus::Partial);
        };
        let Some((&reserved, rest)) = after_reply.split_first() else {
            return Ok(DecodeStatus::Partial);
        };
        if reserved != 0x00 {
            return Err(DecodeError::Malformed("reserved byte is not zero"));
        }
        let (address, consumed) = match Address::decode_from(rest)? {
            | DecodeStatus::Complete(complete) => complete,
            | DecodeStatus::Partial => return Ok(DecodeStatus::Partial),
        };
        let Some(after_address) = rest.get(consumed ..) else {
            return Ok(DecodeStatus::Partial);
        };
        let Some(&[port_high, port_low]) = after_address.first_chunk::<2>() else {
            return Ok(DecodeStatus::Partial);
        };
        Ok(DecodeStatus::Complete((
            Self {
                reply: Reply::from(reply),
                address,
                port: u16::from_be_bytes([port_high, port_low]),
            },
            consumed.saturating_add(5),
        )))
    }
}