//! Forked from fast-socks5 v1.0.0 (upstream `src/server.rs`;
//! extraction; ownership per ADR-013, see the crate root's `mod.rs`).
//!
//! Extracted region: the explicit
//! typestate (`states`, `Socks5ServerProtocol<T, S>`, auth
//! negotiation with the `AuthMethod<T>` machinery, `read_command`,
//! `reply_success`, `reply_error`, `transfer`, `wait_on_tcp`,
//! `ErrorContext<T>`), `SocksServerError`, and the published reply
//! encoder `new_reply` (carried change #2: upstream's is private,
//! `fn new_reply` at `server.rs:1380`; here it is `pub` so the
//! BIND reply#2 writer and consumer reply readers reuse the
//! 10/22-byte reply form instead of re-encoding it). The cut regions
//! (the legacy `Socks5Server`/`Socks5Socket`/`Config`/`Authentication`
//! layer, the `run_*` kernel-net handlers, `DnsResolveHelper`, the
//! net-only tests) stay in `/workspace/fast-socks5`, the not-forked
//! inventory row 1 in `docs/architecture/rfc1928.md`.
//!
//! Extraction-time adjustments (mechanical, ADR-011 §Amendment):
//! `log` calls dropped; `anyhow::Context` dropped; `reply_error`'s
//! `"0.0.0.0:0".parse().unwrap()` replaced with a `const` parsed at
//! compile time (no `unwrap` outside tests). The upstream
//! `Config<A>`/`Authentication` impls that consumed
//! `SocksServerError::AuthenticationRejected` are cut; the variant
//! stays (the typestate's `accept_password_auth` produces it).
//!
//! Post-fork owned changes (extend, never replace, per convention 17):
//! the keep-or-cut sweep (review 002 N-1) cut the producer-replaced
//! machinery: the `auth_method_enums!` macro (upstream's enum
//! generator) with its generated `StandardAuthentication` enum +
//! `allow_no_auth`, `skip_auth_this_is_not_rfc_compliant`, and
//! `CheckResult`'s `Option`/`Result` impls (`bool` stays); none had
//! an in-tree consumer once the producer drove `negotiate_auth`
//! directly. Kept with note (the hand-port map, ADR-013):
//! `transfer`/`wait_on_tcp` are differential-suite-consumed (their
//! semantics stay pinned there), and the error vocabulary keeps its
//! upstream variants. One variant was added, not absorbed:
//! `SocksServerError::Refused`: the honest refusal vocabulary for
//! the producer's in-stream policy refusals (review 002 N-4; never
//! reaches the wire). Recorded in docs/architecture/rfc1928.md.
//!
//! The BIND note (ADR-007): BIND's driving is producer-side code;
//! the typestate already returns `TCPBind` from `read_command()`,
//! `reply_success(sock_addr)` carries reply#1's BND.ADDR, and
//! reply#2 is written with the published `new_reply`. No BIND code
//! lives here; this file only guarantees the surface is present.

use crate::rfc1928::consts;
use crate::rfc1928::util::stream::read_exact;
use crate::rfc1928::util::stream::ConnectError;
use crate::rfc1928::util::target_addr::read_address;
use crate::rfc1928::util::target_addr::TargetAddr;
use crate::rfc1928::AddrError;
use crate::rfc1928::ReplyError;
use crate::rfc1928::Socks5Command;
use crate::rfc1928::UdpHeaderError;
use std::io;
use std::marker::PhantomData;
use std::net::Ipv4Addr;
use std::net::SocketAddr;
use std::net::SocketAddrV4;
use std::string::FromUtf8Error;
use tokio::io::AsyncRead;
use tokio::io::AsyncReadExt;
use tokio::io::AsyncWrite;
use tokio::io::AsyncWriteExt;

#[derive(thiserror::Error, Debug)]
pub enum SocksServerError {
    #[error("i/o error when {context}: {source}")]
    Io {
        source: io::Error,
        context: &'static str,
    },
    #[error("string error when {context}: {source}")]
    FromUtf8 {
        source: FromUtf8Error,
        context: &'static str,
    },
    #[error(transparent)]
    ConnectError(#[from] ConnectError),
    #[error(transparent)]
    UdpHeaderError(#[from] UdpHeaderError),
    #[error(transparent)]
    AddrError(#[from] AddrError),
    #[error("BUG: {0}")] // should be unreachable
    Bug(&'static str),
    #[error("Refused: {0}")]
    Refused(&'static str),
    #[error("Auth method unacceptable `{0:?}`.")]
    AuthMethodUnacceptable(Vec<u8>),
    #[error("Unsupported SOCKS version `{0}`.")]
    UnsupportedSocksVersion(u8),
    #[error("Unsupported SOCKS command `{0}`.")]
    UnknownCommand(u8),
    #[error("Unexpected garbage received on TCP stream used for UDP proxy keep-alive: `{0}`")]
    UnexpectedUdpControlGarbage(u8),
    #[error("Empty username received")]
    EmptyUsername,
    #[error("Empty password received")]
    EmptyPassword,
    #[error("Authentication rejected")]
    AuthenticationRejected,
    #[error("End of stream")]
    EOF,
}

impl SocksServerError {
    pub fn to_reply_error(&self) -> ReplyError {
        match self {
            SocksServerError::UnknownCommand(_) => ReplyError::CommandNotSupported,
            SocksServerError::AddrError(err) => err.to_reply_error(),
            _ => ReplyError::GeneralFailure,
        }
    }
}

pub trait ErrorContext<T> {
    fn err_when(self, context: &'static str) -> Result<T, SocksServerError>;
}

impl<T> ErrorContext<T> for Result<T, io::Error> {
    fn err_when(self, context: &'static str) -> Result<T, SocksServerError> {
        self.map_err(|source| SocksServerError::Io { source, context })
    }
}

impl<T> ErrorContext<T> for Result<T, FromUtf8Error> {
    fn err_when(self, context: &'static str) -> Result<T, SocksServerError> {
        self.map_err(|source| SocksServerError::FromUtf8 { source, context })
    }
}
pub mod states {
    pub struct Opened;
    pub struct Authenticated;
    pub struct CommandRead;
}

pub struct Socks5ServerProtocol<T, S> {
    inner: T,
    _state: PhantomData<S>,
}

impl<T, S> Socks5ServerProtocol<T, S> {
    fn new(inner: T) -> Self {
        Socks5ServerProtocol {
            inner,
            _state: PhantomData,
        }
    }
}

impl<T> Socks5ServerProtocol<T, states::Opened> {
    /// Start handling the SOCKS5 protocol flow, wrapping a client socket.
    pub fn start(inner: T) -> Self {
        Self::new(inner)
    }
}

/// Closure-result adapter for [`Socks5ServerProtocol
/// ::accept_password_auth`]'s credential checker: `bool` (`is_good`).
/// The `Option`/`Result` impls upstream-adjacent shape were cut in
/// the keep-or-cut sweep (review 002 N-1); every in-tree checker
/// returns `bool`; re-add with a typed-rejection design if a door
/// ever wants it.
pub trait CheckResult {
    fn is_good(&self) -> bool;
}

impl CheckResult for bool {
    fn is_good(&self) -> bool {
        *self
    }
}

impl<T> Socks5ServerProtocol<T, states::Authenticated> {
    /// Finish handling the authentication method-specific part of the protocol,
    /// returning back to the overall SOCKS5 flow.
    pub fn finish_auth<A: AuthMethodSuccessState<T>>(auth: A) -> Self {
        Self::new(auth.into_inner())
    }

    /// Handle the SOCKS5 auth negotiation supporting only the `NoAuthentication` method.
    pub async fn accept_no_auth(inner: T) -> Result<Self, SocksServerError>
    where
        T: AsyncWrite + AsyncRead + Unpin,
    {
        Ok(Socks5ServerProtocol::start(inner)
            .negotiate_auth(&[NoAuthentication])
            .await?
            .finish_auth())
    }

    /// Handle the SOCKS5 auth negotiation supporting only the `PasswordAuthentication` method,
    /// and verify the provided username and password using the provided closure.
    ///
    /// The closure can mutate state variables and/or return a result as `Option`/`Result`.
    pub async fn accept_password_auth<F, R>(
        inner: T,
        mut check: F,
    ) -> Result<(Self, R), SocksServerError>
    where
        T: AsyncWrite + AsyncRead + Unpin,
        F: FnMut(String, String) -> R,
        R: CheckResult,
    {
        let (user, pass, auth) = Socks5ServerProtocol::start(inner)
            .negotiate_auth(&[PasswordAuthentication])
            .await?
            .read_username_password()
            .await?;
        let check_result = check(user, pass);
        if check_result.is_good() {
            Ok((auth.accept().await?.finish_auth(), check_result))
        } else {
            auth.reject().await?;
            Err(SocksServerError::AuthenticationRejected)
        }
    }
}

/// A trait for the final successful state of an authentication method's implementation.
///
/// This allows `Socks5ServerProtocol<T, states::Authenticated>::finish_authentication` to
/// let the user continue with the protocol after the socket has been handed off to the
/// authentication method.
pub trait AuthMethodSuccessState<T> {
    fn into_inner(self) -> T;

    fn finish_auth(self) -> Socks5ServerProtocol<T, states::Authenticated>
    where
        Self: Sized,
    {
        Socks5ServerProtocol::finish_auth(self)
    }
}

/// A metadata trait for authentication methods, essentially binding an ID value
/// (as used in the method negotiation) to an actual implementation of the method.
///
/// Use blank structs for individual protocol implementations;
/// `negotiate_auth` takes any slice of `AuthMethod`s, so a custom
/// selection is a slice of the method structs (upstream's
/// enum-generator macro was cut (review 002 N-1); nothing in-tree
/// drove a multi-method slice).
pub trait AuthMethod<T>: Copy {
    type StartingState;
    fn method_id(self) -> u8;
    fn into_starting_state(self, inner: T) -> Self::StartingState;
}

pub struct NoAuthenticationImpl<T>(T);

impl<T> AuthMethodSuccessState<T> for NoAuthenticationImpl<T> {
    fn into_inner(self) -> T {
        self.0
    }
}

/// The "NO AUTHENTICATION REQUIRED" auth method, ID 00h as specifed by RFC 1928.
///
/// As the dummy no-auth method, it only has one state. Once it's been negotiated,
/// you can immediately continue with `finish_authentication`.
///
/// Or not so immediately: if you want to use no-authentication with e.g. IP address
/// allowlisting or TLS client certificate auth for TLS-wrapped SOCKS5, this is your
/// opportunity to reject the no-authentication by dropping the connection!
#[derive(Debug, Clone, Copy)]
pub struct NoAuthentication;

impl<T> AuthMethod<T> for NoAuthentication {
    type StartingState = NoAuthenticationImpl<T>;

    fn method_id(self) -> u8 {
        0x00
    }

    fn into_starting_state(self, inner: T) -> Self::StartingState {
        NoAuthenticationImpl(inner)
    }
}

mod password_states {
    pub struct Started;
    pub struct Received;
    pub struct Finished;
}

pub struct PasswordAuthenticationImpl<T, S> {
    inner: T,
    _state: PhantomData<S>,
}

pub type PasswordAuthenticationStarted<T> = PasswordAuthenticationImpl<T, password_states::Started>;

impl<T, S> PasswordAuthenticationImpl<T, S> {
    fn new(inner: T) -> Self {
        PasswordAuthenticationImpl {
            inner,
            _state: PhantomData,
        }
    }
}

impl<T: AsyncRead + Unpin> PasswordAuthenticationImpl<T, password_states::Started> {
    /// Handle the username and password sent by the client.
    pub async fn read_username_password(
        self,
    ) -> Result<
        (
            String,
            String,
            PasswordAuthenticationImpl<T, password_states::Received>,
        ),
        SocksServerError,
    > {
        let mut socket = self.inner;
        let [_version, user_len] = read_exact!(socket, [0u8; 2]).err_when("reading user len")?;

        if user_len < 1 {
            return Err(SocksServerError::EmptyUsername);
        }

        let username =
            read_exact!(socket, vec![0u8; user_len as usize]).err_when("reading username")?;

        let [pass_len] = read_exact!(socket, [0u8; 1]).err_when("reading password len")?;

        if pass_len < 1 {
            return Err(SocksServerError::EmptyPassword);
        }

        let password =
            read_exact!(socket, vec![0u8; pass_len as usize]).err_when("reading password")?;

        let username = String::from_utf8(username).err_when("converting username")?;
        let password = String::from_utf8(password).err_when("converting password")?;

        Ok((username, password, PasswordAuthenticationImpl::new(socket)))
    }
}

impl<T: AsyncWrite + Unpin> PasswordAuthenticationImpl<T, password_states::Received> {
    /// Notify the client with a "SUCCEEDED" reply and proceed to finish the authentication.
    pub async fn accept(
        mut self,
    ) -> Result<PasswordAuthenticationImpl<T, password_states::Finished>, SocksServerError> {
        self.inner
            .write_all(&[1, consts::SOCKS5_REPLY_SUCCEEDED])
            .await
            .err_when("replying auth success")?;

        Ok(PasswordAuthenticationImpl::new(self.inner))
    }

    /// Notify the client with a "NOT_ACCEPTABLE" reply and drop the socket.
    ///
    /// The status byte is `0xff`: RFC 1929 requires only that a
    /// rejection be any non-`0x00` byte; `0xff` is the convention
    /// upstream picked (borrowed from the method-negotiation phase's
    /// `SOCKS5_AUTH_METHOD_NOT_ACCEPTABLE` constant, which is a
    /// different protocol phase's vocabulary). The bytes interoperate;
    /// the pairing works only because both in-tree sides write/check
    /// the same arbitrary constant.
    pub async fn reject(mut self) -> Result<(), SocksServerError> {
        self.inner
            .write_all(&[1, consts::SOCKS5_AUTH_METHOD_NOT_ACCEPTABLE])
            .await
            .err_when("replying with auth method not acceptable")?;

        Ok(())
    }
}

impl<T> AuthMethodSuccessState<T> for PasswordAuthenticationImpl<T, password_states::Finished> {
    fn into_inner(self) -> T {
        self.inner
    }
}

/// The "USERNAME/PASSWORD" auth method, ID 02h as specified by RFC 1928.
#[derive(Debug, Clone, Copy)]
pub struct PasswordAuthentication;

impl<T> AuthMethod<T> for PasswordAuthentication {
    type StartingState = PasswordAuthenticationImpl<T, password_states::Started>;

    fn method_id(self) -> u8 {
        0x02
    }

    fn into_starting_state(self, inner: T) -> Self::StartingState {
        PasswordAuthenticationImpl::new(inner)
    }
}

macro_rules! try_notify {
    ($proto:expr, $e:expr) => {
        match $e {
            Ok(res) => res,
            Err(err) => {
                let _ = $proto.reply_error(&err.to_reply_error()).await;
                return Err(err.into());
            }
        }
    };
}

impl<T: AsyncRead + AsyncWrite + Unpin> Socks5ServerProtocol<T, states::Opened> {
    /// Negotiate an authentication method from a list of supported ones and initialize it.
    ///
    /// Internally, this reads the list of authentication methods provided by the client, and
    /// picks the first one for which there exists an implementation in `server_methods`.
    ///
    /// If none of the auth methods requested by the client are in `server_methods`,
    /// returns a `SocksServerError::AuthMethodUnacceptable`.
    pub async fn negotiate_auth<M: AuthMethod<T>>(
        mut self,
        server_methods: &[M],
    ) -> Result<M::StartingState, SocksServerError> {
        let [version, methods_len] =
            read_exact!(self.inner, [0u8; 2]).err_when("reading methods")?;

        if version != consts::SOCKS5_VERSION {
            return Err(SocksServerError::UnsupportedSocksVersion(version));
        }

        // {METHODS available from the client}
        // eg. (non-auth) {0, 1}
        // eg. (auth)     {0, 1, 2}
        let methods =
            read_exact!(self.inner, vec![0u8; methods_len as usize]).err_when("reading methods")?;

        // server_methods order matter!
        // the server could choose to prioritize methods
        for server_method in server_methods {
            for client_method_id in methods.iter() {
                if server_method.method_id() == *client_method_id {
                    self.inner
                        .write_all(&[consts::SOCKS5_VERSION, *client_method_id])
                        .await
                        .err_when("replying with auth method")?;
                    return Ok(server_method.into_starting_state(self.inner));
                }
            }
        }

        self.inner
            .write_all(&[
                consts::SOCKS5_VERSION,
                consts::SOCKS5_AUTH_METHOD_NOT_ACCEPTABLE,
            ])
            .await
            .err_when("replying with method not acceptable")?;
        Err(SocksServerError::AuthMethodUnacceptable(methods))
    }
}

impl<T: AsyncRead + AsyncWrite + Unpin> Socks5ServerProtocol<T, states::CommandRead> {
    /// Reply success to the client according to the RFC.
    /// This consumes the wrapper as after this message actual proxying should begin.
    pub async fn reply_success(mut self, sock_addr: SocketAddr) -> Result<T, SocksServerError> {
        self.inner
            .write_all(&new_reply(&ReplyError::Succeeded, sock_addr))
            .await
            .err_when("writing successful reply")?;

        self.inner.flush().await.err_when("flushing auth reply")?;

        Ok(self.inner)
    }

    /// Reply error to the client with the reply code according to the RFC.
    pub async fn reply_error(mut self, error: &ReplyError) -> Result<(), SocksServerError> {
        let reply = new_reply(error, V4_UNSPECIFIED_ZERO);
        self.inner
            .write_all(&reply)
            .await
            .err_when("writing unsuccessful reply")?;

        self.inner.flush().await.err_when("flushing auth reply")?;

        Ok(())
    }
}

impl<T: AsyncRead + AsyncWrite + Unpin> Socks5ServerProtocol<T, states::Authenticated> {
    /// Decide to whether or not, accept the authentication method.
    /// Don't forget that the methods list sent by the client, contains one or more methods.
    ///
    /// # Request
    /// ```text
    ///          +----+-----+-------+------+----------+----------+
    ///          |VER | CMD |  RSV  | ATYP | DST.ADDR | DST.PORT |
    ///          +----+-----+-------+------+----------+----------+
    ///          | 1  |  1  |   1   |  1   | Variable |    2     |
    ///          +----+-----+-------+------+----------+----------+
    /// ```
    ///
    /// It the request is correct, it should returns a ['SocketAddr'].
    ///
    pub async fn read_command(
        mut self,
    ) -> Result<
        (
            Socks5ServerProtocol<T, states::CommandRead>,
            Socks5Command,
            TargetAddr,
        ),
        SocksServerError,
    > {
        let [version, cmd, _rsv, address_type] =
            read_exact!(self.inner, [0u8; 4]).err_when("reading command")?;

        if version != consts::SOCKS5_VERSION {
            return Err(SocksServerError::UnsupportedSocksVersion(version));
        }

        let mut proto = Socks5ServerProtocol::new(self.inner);

        // Guess address type
        let target_addr = try_notify!(proto, read_address(&mut proto.inner, address_type).await);

        let cmd = try_notify!(
            proto,
            Socks5Command::from_u8(cmd).ok_or(SocksServerError::UnknownCommand(cmd))
        );

        Ok((proto, cmd, target_addr))
    }
}

/// The RFC reply form's placeholder BND.ADDR for error replies
/// (statically constructed; upstream parsed `"0.0.0.0:0"` at runtime).
const V4_UNSPECIFIED_ZERO: SocketAddr = SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::UNSPECIFIED, 0));

/// Generate reply code according to the RFC.
///
/// Carried change #2 (ADR-011): upstream keeps this helper private
/// (fast-socks5 `server.rs:1380`); in-tree it is `pub` so the BIND
/// reply#2 writer and the consumer's reply readers reuse the
/// 10-byte (v4) / 22-byte (v6) reply form instead of re-encoding it.
pub fn new_reply(error: &ReplyError, sock_addr: SocketAddr) -> Vec<u8> {
    let (addr_type, mut ip_oct, mut port) = match sock_addr {
        SocketAddr::V4(sock) => (
            consts::SOCKS5_ADDR_TYPE_IPV4,
            sock.ip().octets().to_vec(),
            sock.port().to_be_bytes().to_vec(),
        ),
        SocketAddr::V6(sock) => (
            consts::SOCKS5_ADDR_TYPE_IPV6,
            sock.ip().octets().to_vec(),
            sock.port().to_be_bytes().to_vec(),
        ),
    };

    let mut reply = vec![
        consts::SOCKS5_VERSION,
        error.as_u8(), // transform the error into byte code
        0x00,          // reserved
        addr_type,     // address type (ipv4, v6, domain)
    ];
    reply.append(&mut ip_oct);
    reply.append(&mut port);

    reply
}

/// Wait until a TCP stream (that's not supposed to receive anything) closes.
///
/// This is intended for cancelling the `transfer_udp` task.
pub async fn wait_on_tcp<I>(stream: &mut I) -> Result<(), SocksServerError>
where
    I: AsyncRead + Unpin,
{
    let mut buf = [0; 1];
    match stream.read(&mut buf).await {
        Ok(0) => Err(SocksServerError::EOF),
        Ok(_) => Err(SocksServerError::UnexpectedUdpControlGarbage(buf[0])),
        Err(err) => Err(err).err_when("waiting on UDP control stream"),
    }
}

/// Run a bidirectional proxy between two streams.
/// Using 2 different generators, because they could be different structs with same traits.
pub async fn transfer<I, O>(mut inbound: I, mut outbound: O)
where
    I: AsyncRead + AsyncWrite + Unpin,
    O: AsyncRead + AsyncWrite + Unpin,
{
    let _ = tokio::io::copy_bidirectional(&mut inbound, &mut outbound).await;
}

#[cfg(test)]
mod tests {
    use super::new_reply;
    use super::states;
    use super::Socks5ServerProtocol;
    use super::SocksServerError;
    use crate::rfc1928::util::target_addr::TargetAddr;
    use crate::rfc1928::ReplyError;
    use crate::rfc1928::Socks5Command;
    use std::net::SocketAddr;
    use tokio::io::AsyncReadExt;
    use tokio::io::AsyncWriteExt;

    /// `SocksServerError::to_reply_error`'s full mapping table (review
    /// 001 L-3): `UnknownCommand` → `CommandNotSupported`,
    /// `AddrError` delegates to its own table, and every remaining
    /// variant falls through to `GeneralFailure`, the mapping the
    /// producer's error surface relies on when mapping session errors
    /// to reply codes (`try_notify!`'s secondary-error write). The
    /// fall-through arms (`Io`/`FromUtf8`/`EOF`/auth/refusal variants)
    /// are exercised here directly, since driving each from the wire
    /// is not one-test-cheap. Upstream `src/server.rs:62-70`. The
    /// table also pins the vocabulary honesty (review 002 N-4):
    /// `Bug` and `Refused` are log-vocabulary variants; both fall
    /// through to `GeneralFailure` (they never reach the wire from
    /// the typestate; the producer's refusal sites write the RFC
    /// reply before the session error).
    #[test]
    fn to_reply_error_mapping_table() {
        assert!(matches!(
            SocksServerError::UnknownCommand(0x7f).to_reply_error(),
            ReplyError::CommandNotSupported
        ));
        assert!(matches!(
            SocksServerError::AddrError(crate::rfc1928::AddrError::IncorrectAddressType)
                .to_reply_error(),
            ReplyError::AddressTypeNotSupported
        ));
        for fall_through in [
            SocksServerError::EOF,
            SocksServerError::EmptyUsername,
            SocksServerError::EmptyPassword,
            SocksServerError::AuthenticationRejected,
            SocksServerError::UnsupportedSocksVersion(0x04),
            SocksServerError::Bug("test"),
            SocksServerError::Refused("test"),
            SocksServerError::AuthMethodUnacceptable(vec![0x02]),
            SocksServerError::UnexpectedUdpControlGarbage(0xde),
            SocksServerError::Io {
                source: std::io::Error::other("x"),
                context: "reading",
            },
            SocksServerError::FromUtf8 {
                source: String::from_utf8(vec![0xff]).unwrap_err(),
                context: "converting",
            },
        ] {
            assert!(matches!(
                fall_through.to_reply_error(),
                ReplyError::GeneralFailure
            ));
        }
    }

    /// The refusal vocabulary (review 002 N-4): `Refused` is the
    /// honest display for the producer's in-stream policy refusals:
    /// not `BUG:` (a working-as-designed refusal is not a bug) and
    /// not "Authentication rejected" (an ACL refusal is not an auth
    /// failure). It is log/observability vocabulary: never encoded
    /// on the wire by the typestate (only `to_reply_error` maps it,
    /// and the producer's refusal sites write the RFC reply before
    /// the session error). Upstream has no such variant; that is the
    /// fork's owned addition, recorded in the module doc.
    #[test]
    fn refused_variant_display_is_honest_and_wire_invisible() {
        assert_eq!(
            SocksServerError::Refused("policy says no").to_string(),
            "Refused: policy says no"
        );
        assert!(matches!(
            SocksServerError::Refused("policy says no").to_reply_error(),
            ReplyError::GeneralFailure
        ));
    }

    #[test]
    fn new_reply_v4_form() {
        let addr: SocketAddr = "127.0.0.1:1080".parse().unwrap();
        let reply = new_reply(&ReplyError::Succeeded, addr);
        assert_eq!(reply, vec![0x05, 0x00, 0x00, 0x01, 127, 0, 0, 1, 4, 56]);
        assert_eq!(reply.len(), 10);
    }

    #[test]
    fn new_reply_v6_form() {
        let addr: SocketAddr = "[2001:db8::1]:80".parse().unwrap();
        let reply = new_reply(&ReplyError::CommandNotSupported, addr);
        assert_eq!(reply.len(), 22);
        assert_eq!(reply[0], 0x05);
        assert_eq!(reply[1], 0x07);
        assert_eq!(reply[2], 0x00);
        assert_eq!(reply[3], 0x04);
        assert_eq!(
            &reply[4..20],
            &0x2001_0db8_0000_0000_0000_0000_0000_0001u128.to_be_bytes()
        );
        assert_eq!(&reply[20..22], &80u16.to_be_bytes());
    }

    #[test]
    fn new_reply_error_codes_are_carried() {
        let addr: SocketAddr = "10.0.0.1:1".parse().unwrap();
        for (code, byte) in [
            (ReplyError::Succeeded, 0x00),
            (ReplyError::GeneralFailure, 0x01),
            (ReplyError::ConnectionNotAllowed, 0x02),
            (ReplyError::NetworkUnreachable, 0x03),
            (ReplyError::HostUnreachable, 0x04),
            (ReplyError::ConnectionRefused, 0x05),
            (ReplyError::ConnectionTimeout, 0x06),
            (ReplyError::TtlExpired, 0x06),
            (ReplyError::CommandNotSupported, 0x07),
            (ReplyError::AddressTypeNotSupported, 0x08),
        ] {
            assert_eq!(new_reply(&code, addr)[1], byte);
        }
    }

    #[tokio::test]
    async fn version_mismatch_is_refused() {
        let (mut client, server) = tokio::io::duplex(64);
        let handle: tokio::task::JoinHandle<Result<(), SocksServerError>> =
            tokio::spawn(async move {
                Socks5ServerProtocol::<_, states::Authenticated>::accept_no_auth(server)
                    .await
                    .map(|_| ())
            });
        client.write_all(&[0x04, 0x01, 0x00]).await.unwrap();
        let err = handle.await.unwrap().unwrap_err();
        assert!(matches!(
            err,
            SocksServerError::UnsupportedSocksVersion(0x04)
        ));
    }

    #[tokio::test]
    async fn no_acceptable_method_replies_0xff() {
        let (mut client, server) = tokio::io::duplex(64);
        let handle: tokio::task::JoinHandle<Result<(), SocksServerError>> =
            tokio::spawn(async move {
                Socks5ServerProtocol::<_, states::Authenticated>::accept_no_auth(server)
                    .await
                    .map(|_| ())
            });
        client.write_all(&[0x05, 0x01, 0x02]).await.unwrap();
        let mut method_reply = [0u8; 2];
        client.read_exact(&mut method_reply).await.unwrap();
        assert_eq!(method_reply, [0x05, 0xff]);
        let err = handle.await.unwrap().unwrap_err();
        assert!(matches!(err, SocksServerError::AuthMethodUnacceptable(ref m) if m == &vec![0x02]));
    }

    #[tokio::test]
    async fn unknown_command_replies_0x07() {
        let (mut client, server) = tokio::io::duplex(64);
        let handle: tokio::task::JoinHandle<Result<(), SocksServerError>> =
            tokio::spawn(async move {
                let proto =
                    Socks5ServerProtocol::<_, states::Authenticated>::accept_no_auth(server)
                        .await
                        .unwrap();
                proto.read_command().await.map(|_| ())
            });
        client.write_all(&[0x05, 0x01, 0x00]).await.unwrap();
        let mut method_reply = [0u8; 2];
        client.read_exact(&mut method_reply).await.unwrap();
        assert_eq!(method_reply, [0x05, 0x00]);
        client
            .write_all(&[0x05, 0x7f, 0x00, 0x01, 192, 168, 0, 1, 0, 80])
            .await
            .unwrap();
        let mut reply = [0u8; 10];
        client.read_exact(&mut reply).await.unwrap();
        assert_eq!(reply, [0x05, 0x07, 0x00, 0x01, 0, 0, 0, 0, 0, 0]);
        let err = handle.await.unwrap().unwrap_err();
        assert!(matches!(err, SocksServerError::UnknownCommand(0x7f)));
    }

    #[tokio::test]
    async fn success_reply_form_round_trip() {
        let (mut client, server) = tokio::io::duplex(64);
        let handle: tokio::task::JoinHandle<Result<(Socks5Command, TargetAddr), SocksServerError>> =
            tokio::spawn(async move {
                let proto =
                    Socks5ServerProtocol::<_, states::Authenticated>::accept_no_auth(server)
                        .await
                        .unwrap();
                let (proto, cmd, target) = proto.read_command().await.unwrap();
                let bound: SocketAddr = "198.51.100.7:443".parse().unwrap();
                let mut inner = proto.reply_success(bound).await.unwrap();
                inner.write_all(b"ok").await.unwrap();
                Ok((cmd, target))
            });
        client.write_all(&[0x05, 0x01, 0x00]).await.unwrap();
        let mut method_reply = [0u8; 2];
        client.read_exact(&mut method_reply).await.unwrap();
        assert_eq!(method_reply, [0x05, 0x00]);
        client
            .write_all(&[0x05, 0x01, 0x00, 0x01, 192, 0, 2, 1, 0, 80])
            .await
            .unwrap();
        let mut reply = [0u8; 10];
        client.read_exact(&mut reply).await.unwrap();
        assert_eq!(reply, [0x05, 0x00, 0x00, 0x01, 198, 51, 100, 7, 1, 187]);
        let mut data = [0u8; 2];
        client.read_exact(&mut data).await.unwrap();
        assert_eq!(&data, b"ok");
        let (cmd, target) = handle.await.unwrap().unwrap();
        assert_eq!(cmd, Socks5Command::TCPConnect);
        assert_eq!(target, TargetAddr::Ip("192.0.2.1:80".parse().unwrap()));
    }

    #[tokio::test]
    async fn error_reply_form_round_trip() {
        let (mut client, server) = tokio::io::duplex(64);
        let handle: tokio::task::JoinHandle<Result<(), SocksServerError>> =
            tokio::spawn(async move {
                let proto =
                    Socks5ServerProtocol::<_, states::Authenticated>::accept_no_auth(server)
                        .await
                        .unwrap();
                let (proto, _cmd, _target) = proto.read_command().await.unwrap();
                proto.reply_error(&ReplyError::ConnectionNotAllowed).await
            });
        client.write_all(&[0x05, 0x01, 0x00]).await.unwrap();
        let mut method_reply = [0u8; 2];
        client.read_exact(&mut method_reply).await.unwrap();
        assert_eq!(method_reply, [0x05, 0x00]);
        client
            .write_all(&[0x05, 0x01, 0x00, 0x03, 3, b'i', b'p', b'v', b'6', 0, 80])
            .await
            .unwrap();
        let mut reply = [0u8; 10];
        client.read_exact(&mut reply).await.unwrap();
        assert_eq!(reply, [0x05, 0x02, 0x00, 0x01, 0, 0, 0, 0, 0, 0]);
        handle.await.unwrap().unwrap();
    }

    #[tokio::test]
    async fn password_auth_accepts_and_rejects() {
        let (mut client, server) = tokio::io::duplex(64);
        let handle: tokio::task::JoinHandle<Result<(), SocksServerError>> =
            tokio::spawn(async move {
                let (proto, _ok) =
                    Socks5ServerProtocol::<_, states::Authenticated>::accept_password_auth(
                        server,
                        |user, pass| user == "alice" && pass == "secret",
                    )
                    .await?;
                drop(proto);
                Ok(())
            });
        client.write_all(&[0x05, 0x01, 0x02]).await.unwrap();
        let mut method_reply = [0u8; 2];
        client.read_exact(&mut method_reply).await.unwrap();
        assert_eq!(method_reply, [0x05, 0x02]);
        client
            .write_all(&[
                0x01, 0x05, b'a', b'l', b'i', b'c', b'e', 0x06, b's', b'e', b'c', b'r', b'e', b't',
            ])
            .await
            .unwrap();
        let mut status = [0u8; 2];
        client.read_exact(&mut status).await.unwrap();
        assert_eq!(status, [0x01, 0x00]);
        handle.await.unwrap().unwrap();

        let (mut client, server) = tokio::io::duplex(64);
        let handle: tokio::task::JoinHandle<Result<(), SocksServerError>> =
            tokio::spawn(async move {
                Socks5ServerProtocol::<_, states::Authenticated>::accept_password_auth(
                    server,
                    |_user, _pass| false,
                )
                .await
                .map(|_| ())
            });
        client.write_all(&[0x05, 0x01, 0x02]).await.unwrap();
        let mut method_reply = [0u8; 2];
        client.read_exact(&mut method_reply).await.unwrap();
        assert_eq!(method_reply, [0x05, 0x02]);
        client
            .write_all(&[0x01, 0x03, b'b', b'o', b'b', 0x01, b'x'])
            .await
            .unwrap();
        let mut status = [0u8; 2];
        client.read_exact(&mut status).await.unwrap();
        assert_eq!(status, [0x01, 0xff]);
        let err = handle.await.unwrap().unwrap_err();
        assert!(matches!(err, SocksServerError::AuthenticationRejected));
    }
}