<commit_before>use std::{net::SocketAddr, process::exit};

use anyhow::bail;

use sozu_command_lib::{
    certificate::{calculate_fingerprint, split_certificate_chain},
    config::{Config, FileListenerProtocolConfig, Listener, ProxyProtocolConfig},
    proxy::{
        ActivateListener, AddCertificate, Backend, CertificateAndKey, CertificateFingerprint,
        Cluster, DeactivateListener, HttpFrontend, ListenerType, LoadBalancingParams, PathRule,
        ProxyRequestData, RemoveBackend, RemoveCertificate, RemoveListener, ReplaceCertificate,
        RulePosition, TcpFrontend, TcpListener, TlsVersion,
    },
};

use crate::{
    cli::{
        ApplicationCmd, BackendCmd, HttpFrontendCmd, HttpListenerCmd, HttpsListenerCmd,
        LoggingLevel, TcpFrontendCmd, TcpListenerCmd,
    },
    ctl::CommandManager,
};

impl CommandManager {
    pub fn backend_command(&mut self, cmd: BackendCmd) -> Result<(), anyhow::Error> {
        match cmd {
            BackendCmd::Add {
                id,
                backend_id,
                address,
                sticky_id,
                backup,
            } => self.order_command(ProxyRequestData::AddBackend(Backend {
                cluster_id: id,
                address,
                backend_id,
                load_balancing_parameters: Some(LoadBalancingParams::default()),
                sticky_id,
                backup,
            })),
            BackendCmd::Remove {
                id,
                backend_id,
                address,
            } => self.order_command(ProxyRequestData::RemoveBackend(RemoveBackend {
                cluster_id: id,
                address,
                backend_id,
            })),
        }
    }

    pub fn application_command(&mut self, cmd: ApplicationCmd) -> Result<(), anyhow::Error> {
        match cmd {
            ApplicationCmd::Add {
                id,
                sticky_session,
                https_redirect,
                send_proxy,
                expect_proxy,
                load_balancing_policy,
            } => {
                let proxy_protocol = match (send_proxy, expect_proxy) {
                    (true, true) => Some(ProxyProtocolConfig::RelayHeader),
                    (true, false) => Some(ProxyProtocolConfig::SendHeader),
                    (false, true) => Some(ProxyProtocolConfig::ExpectHeader),
                    _ => None,
                };
                self.order_command(ProxyRequestData::AddCluster(Cluster {
                    cluster_id: id,
                    sticky_session,
                    https_redirect,
                    proxy_protocol,
                    load_balancing: load_balancing_policy,
                    load_metric: None,
                    answer_503: None,
                }))
            }
            ApplicationCmd::Remove { id } => {
                self.order_command(ProxyRequestData::RemoveCluster { cluster_id: id })
            }
        }
    }

    pub fn tcp_frontend_command(&mut self, cmd: TcpFrontendCmd) -> Result<(), anyhow::Error> {
        match cmd {
            TcpFrontendCmd::Add { id, address } => {
                self.order_command(ProxyRequestData::AddTcpFrontend(TcpFrontend {
                    cluster_id: id,
                    address,
                }))
            }
            TcpFrontendCmd::Remove { id, address } => {
                self.order_command(ProxyRequestData::RemoveTcpFrontend(TcpFrontend {
                    cluster_id: id,
                    address,
                }))
            }
        }
    }

    pub fn http_frontend_command(&mut self, cmd: HttpFrontendCmd) -> Result<(), anyhow::Error> {
        match cmd {
            HttpFrontendCmd::Add {
                hostname,
                path_begin,
                address,
                method,
                route,
            } => self.order_command(ProxyRequestData::AddHttpFrontend(HttpFrontend {
                route: route.into(),
                address,
                hostname,
                path: PathRule::Prefix(path_begin.unwrap_or_else(|| "".to_string())),
                method: method.map(String::from),
                position: RulePosition::Tree,
            })),

            HttpFrontendCmd::Remove {
                hostname,
                path_begin,
                address,
                method,
                route,
            } => self.order_command(ProxyRequestData::RemoveHttpFrontend(HttpFrontend {
                route: route.into(),
                address,
                hostname,
                path: PathRule::Prefix(path_begin.unwrap_or_else(|| "".to_string())),
                method: method.map(String::from),
                position: RulePosition::Tree,
            })),
        }
    }

    pub fn https_frontend_command(&mut self, cmd: HttpFrontendCmd) -> Result<(), anyhow::Error> {
        match cmd {
            HttpFrontendCmd::Add {
                hostname,
                path_begin,
                address,
                method,
                route,
            } => self.order_command(ProxyRequestData::AddHttpsFrontend(HttpFrontend {
                route: route.into(),
                address,
                hostname,
                path: PathRule::Prefix(path_begin.unwrap_or_else(|| "".to_string())),
                method: method.map(String::from),
                position: RulePosition::Tree,
            })),
            HttpFrontendCmd::Remove {
                hostname,
                path_begin,
                address,
                method,
                route,
            } => self.order_command(ProxyRequestData::RemoveHttpsFrontend(HttpFrontend {
                route: route.into(),
                address,
                hostname,
                path: PathRule::Prefix(path_begin.unwrap_or_else(|| "".to_string())),
                method: method.map(String::from),
                position: RulePosition::Tree,
            })),
        }
    }

    pub fn https_listener_command(&mut self, cmd: HttpsListenerCmd) -> Result<(), anyhow::Error> {
        match cmd {
            HttpsListenerCmd::Add {
                address,
                public_address,
                answer_404,
                answer_503,
                tls_versions,
                cipher_list,
                rustls_cipher_list,
                expect_proxy,
                sticky_name,
            } => {
                let mut listener = Listener::new(address, FileListenerProtocolConfig::Https);
                listener.public_address = public_address;
                listener.answer_404 = answer_404;
                listener.answer_503 = answer_503;
                listener.expect_proxy = Some(expect_proxy);
                if let Some(sticky_name) = sticky_name {
                    listener.sticky_name = sticky_name;
                }
                listener.cipher_list = cipher_list;
                listener.tls_versions = if tls_versions.is_empty() {
                    None
                } else {
                    Some(tls_versions)
                };
                listener.rustls_cipher_list = if rustls_cipher_list.is_empty() {
                    None
                } else {
                    Some(rustls_cipher_list)
                };
                match listener.to_tls(None, None, None) {
                    Some(conf) => self.order_command(ProxyRequestData::AddHttpsListener(conf)),
                    None => bail!("Error creating HTTPS listener"),
                }
            }
            HttpsListenerCmd::Remove { address } => {
                self.remove_listener(address, ListenerType::HTTPS)
            }
            HttpsListenerCmd::Activate { address } => {
                self.activate_listener(address, ListenerType::HTTPS)
            }
            HttpsListenerCmd::Deactivate { address } => {
                self.deactivate_listener(address, ListenerType::HTTPS)
            }
        }
    }

    pub fn http_listener_command(&mut self, cmd: HttpListenerCmd) -> Result<(), anyhow::Error> {
        match cmd {
            HttpListenerCmd::Add {
                address,
                public_address,
                answer_404,
                answer_503,
                expect_proxy,
                sticky_name,
            } => {
                let mut listener = Listener::new(address, FileListenerProtocolConfig::Http);
                listener.public_address = public_address;
                listener.answer_404 = answer_404;
                listener.answer_503 = answer_503;
                listener.expect_proxy = Some(expect_proxy);
                if let Some(sticky_name) = sticky_name {
                    listener.sticky_name = sticky_name;
                }
                match listener.to_http(None, None, None) {
                    Some(conf) => self.order_command(ProxyRequestData::AddHttpListener(conf)),
                    None => bail!("Error creating HTTP listener"),
                }
            }
            HttpListenerCmd::Remove { address } => {
                self.remove_listener(address, ListenerType::HTTP)
            }
            HttpListenerCmd::Activate { address } => {
                self.activate_listener(address, ListenerType::HTTP)
            }
            HttpListenerCmd::Deactivate { address } => {
                self.deactivate_listener(address, ListenerType::HTTP)
            }
        }
    }

    pub fn tcp_listener_command(&mut self, cmd: TcpListenerCmd) -> Result<(), anyhow::Error> {
        match cmd {
            TcpListenerCmd::Add {
                address,
                public_address,
                expect_proxy,
            } => self.order_command(ProxyRequestData::AddTcpListener(TcpListener {
                address,
                public_address,
                expect_proxy,
                front_timeout: 60,
                back_timeout: 30,
                connect_timeout: 3,
            })),
            TcpListenerCmd::Remove { address } => self.remove_listener(address, ListenerType::TCP),
            TcpListenerCmd::Activate { address } => {
                self.activate_listener(address, ListenerType::TCP)
            }
            TcpListenerCmd::Deactivate { address } => {
                self.deactivate_listener(address, ListenerType::TCP)
            }
        }
    }

    pub fn remove_listener(
        &mut self,
        address: SocketAddr,
        proxy: ListenerType,
    ) -> Result<(), anyhow::Error> {
        self.order_command(ProxyRequestData::RemoveListener(RemoveListener {
            address,
            proxy,
        }))
    }

    pub fn activate_listener(
        &mut self,
        address: SocketAddr,
        proxy: ListenerType,
    ) -> Result<(), anyhow::Error> {
        self.order_command(ProxyRequestData::ActivateListener(ActivateListener {
            address,
            proxy,
            from_scm: false,
        }))
    }

    pub fn deactivate_listener(
        &mut self,
        address: SocketAddr,
        proxy: ListenerType,
    ) -> Result<(), anyhow::Error> {
        self.order_command(ProxyRequestData::DeactivateListener(DeactivateListener {
            address,
            proxy,
            to_scm: false,
        }))
    }

    pub fn logging_filter(&mut self, filter: &LoggingLevel) -> Result<(), anyhow::Error> {
        self.order_command(ProxyRequestData::Logging(filter.to_string().to_lowercase()))
    }

    pub fn add_certificate(
        &mut self,
        address: SocketAddr,
        certificate_path: &str,
        certificate_chain_path: &str,
        key_path: &str,
        versions: Vec<TlsVersion>,
    ) -> Result<(), anyhow::Error> {
        if let Some(new_certificate) =
            load_full_certificate(certificate_path, certificate_chain_path, key_path, versions)?
        {
            self.order_command(ProxyRequestData::AddCertificate(AddCertificate {
                address,
                certificate: new_certificate,
                names: vec![],
                expired_at: None,
            }))?;
        }
        Ok(())
    }

    pub fn replace_certificate(
        &mut self,
        address: SocketAddr,
        new_certificate_path: &str,
        new_certificate_chain_path: &str,
        new_key_path: &str,
        old_certificate_path: Option<&str>,
        old_fingerprint: Option<&str>,
        versions: Vec<TlsVersion>,
    ) -> Result<(), anyhow::Error> {
        if old_certificate_path.is_some() && old_fingerprint.is_some() {
            bail!("Error: Either provide the old certificate's path or its fingerprint");
        }

        if old_certificate_path.is_none() && old_fingerprint.is_none() {
            bail!("Error: Either provide the old certificate's path or its fingerprint");
        }

        if let Some(new_certificate) = load_full_certificate(
            new_certificate_path,
            new_certificate_chain_path,
            new_key_path,
            versions,
        )? {
            if let Some(old_fingerprint) = old_fingerprint.and_then(|s| {
        match hex::decode(s) {
            Ok(v) => Some(CertificateFingerprint(v)),
            Err(e) => {
                eprintln!("Error decoding the certificate fingerprint (expected hexadecimal data): {:?}", e);
                None
            }
        }
    }).or_else(|| old_certificate_path.and_then(get_certificate_fingerprint)) {
     self.order_command( ProxyRequestData::ReplaceCertificate(ReplaceCertificate {
        address,
        new_certificate,
        old_fingerprint,
        new_names: vec![],
        new_expired_at: None,
      }))?;
    }
        }
        Ok(())
    }

    pub fn remove_certificate(
        &mut self,
        address: SocketAddr,
        certificate_path: Option<&str>,
        fingerprint: Option<&str>,
    ) -> Result<(), anyhow::Error> {
        if certificate_path.is_some() && fingerprint.is_some() {
            bail!("Error: Either provide the certificate's path or its fingerprint");
        }

        if certificate_path.is_none() && fingerprint.is_none() {
            bail!("Error: Either provide the certificate's path or its fingerprint");
        }

        if let Some(fingerprint) = fingerprint
            .and_then(|s| match hex::decode(s) {
                Ok(v) => Some(CertificateFingerprint(v)),
                Err(e) => {
                    eprintln!(
                    "Error decoding the certificate fingerprint (expected hexadecimal data): {:?}",
                    e
                );
                    None
                }
            })
            .or_else(|| certificate_path.and_then(get_certificate_fingerprint))
        {
            self.order_command(ProxyRequestData::RemoveCertificate(RemoveCertificate {
                address,
                fingerprint,
            }))?
        }
        Ok(())
    }
}

fn get_certificate_fingerprint(certificate_path: &str) -> Option<CertificateFingerprint> {
    match Config::load_file_bytes(certificate_path) {
        Ok(data) => match calculate_fingerprint(&data) {
            Some(fingerprint) => Some(CertificateFingerprint(fingerprint)),
            None => {
                eprintln!("could not calculate finrprint for certificate");
                exit(1);
            }
        },
        Err(e) => {
            eprintln!("could not load file: {:?}", e);
            exit(1);
        }
    }
}

fn load_full_certificate(
    certificate_path: &str,
    certificate_chain_path: &str,
    key_path: &str,
    versions: Vec<TlsVersion>,
) -> Result<Option<CertificateAndKey>, anyhow::Error> {
    match Config::load_file(certificate_path) {
        Err(e) => {
            bail!("could not load certificate: {:?}", e);
        }
        Ok(certificate) => {
            match Config::load_file(certificate_chain_path).map(split_certificate_chain) {
                Err(e) => {
                    bail!("could not load certificate chain: {:?}", e);
                }
                Ok(certificate_chain) => match Config::load_file(key_path) {
                    Err(e) => {
                        bail!("could not load key: {:?}", e);
                    }
                    Ok(key) => Ok(Some(CertificateAndKey {
                        certificate,
                        certificate_chain,
                        key,
                        versions,
                    })),
                },
            }
        }
    }
}
<commit_msg>refactor certificate logic in ctl, with Results instead of Options<commit_after>use std::net::SocketAddr;

use anyhow::{bail, Context};

use sozu_command_lib::{
    certificate::{calculate_fingerprint, split_certificate_chain},
    config::{Config, FileListenerProtocolConfig, Listener, ProxyProtocolConfig},
    proxy::{
        ActivateListener, AddCertificate, Backend, CertificateAndKey, CertificateFingerprint,
        Cluster, DeactivateListener, HttpFrontend, ListenerType, LoadBalancingParams, PathRule,
        ProxyRequestData, RemoveBackend, RemoveCertificate, RemoveListener, ReplaceCertificate,
        RulePosition, TcpFrontend, TcpListener, TlsVersion,
    },
};

use crate::{
    cli::{
        ApplicationCmd, BackendCmd, HttpFrontendCmd, HttpListenerCmd, HttpsListenerCmd,
        LoggingLevel, TcpFrontendCmd, TcpListenerCmd,
    },
    ctl::CommandManager,
};

impl CommandManager {
    pub fn backend_command(&mut self, cmd: BackendCmd) -> Result<(), anyhow::Error> {
        match cmd {
            BackendCmd::Add {
                id,
                backend_id,
                address,
                sticky_id,
                backup,
            } => self.order_command(ProxyRequestData::AddBackend(Backend {
                cluster_id: id,
                address,
                backend_id,
                load_balancing_parameters: Some(LoadBalancingParams::default()),
                sticky_id,
                backup,
            })),
            BackendCmd::Remove {
                id,
                backend_id,
                address,
            } => self.order_command(ProxyRequestData::RemoveBackend(RemoveBackend {
                cluster_id: id,
                address,
                backend_id,
            })),
        }
    }

    pub fn application_command(&mut self, cmd: ApplicationCmd) -> Result<(), anyhow::Error> {
        match cmd {
            ApplicationCmd::Add {
                id,
                sticky_session,
                https_redirect,
                send_proxy,
                expect_proxy,
                load_balancing_policy,
            } => {
                let proxy_protocol = match (send_proxy, expect_proxy) {
                    (true, true) => Some(ProxyProtocolConfig::RelayHeader),
                    (true, false) => Some(ProxyProtocolConfig::SendHeader),
                    (false, true) => Some(ProxyProtocolConfig::ExpectHeader),
                    _ => None,
                };
                self.order_command(ProxyRequestData::AddCluster(Cluster {
                    cluster_id: id,
                    sticky_session,
                    https_redirect,
                    proxy_protocol,
                    load_balancing: load_balancing_policy,
                    load_metric: None,
                    answer_503: None,
                }))
            }
            ApplicationCmd::Remove { id } => {
                self.order_command(ProxyRequestData::RemoveCluster { cluster_id: id })
            }
        }
    }

    pub fn tcp_frontend_command(&mut self, cmd: TcpFrontendCmd) -> Result<(), anyhow::Error> {
        match cmd {
            TcpFrontendCmd::Add { id, address } => {
                self.order_command(ProxyRequestData::AddTcpFrontend(TcpFrontend {
                    cluster_id: id,
                    address,
                }))
            }
            TcpFrontendCmd::Remove { id, address } => {
                self.order_command(ProxyRequestData::RemoveTcpFrontend(TcpFrontend {
                    cluster_id: id,
                    address,
                }))
            }
        }
    }

    pub fn http_frontend_command(&mut self, cmd: HttpFrontendCmd) -> Result<(), anyhow::Error> {
        match cmd {
            HttpFrontendCmd::Add {
                hostname,
                path_begin,
                address,
                method,
                route,
            } => self.order_command(ProxyRequestData::AddHttpFrontend(HttpFrontend {
                route: route.into(),
                address,
                hostname,
                path: PathRule::Prefix(path_begin.unwrap_or_else(|| "".to_string())),
                method: method.map(String::from),
                position: RulePosition::Tree,
            })),

            HttpFrontendCmd::Remove {
                hostname,
                path_begin,
                address,
                method,
                route,
            } => self.order_command(ProxyRequestData::RemoveHttpFrontend(HttpFrontend {
                route: route.into(),
                address,
                hostname,
                path: PathRule::Prefix(path_begin.unwrap_or_else(|| "".to_string())),
                method: method.map(String::from),
                position: RulePosition::Tree,
            })),
        }
    }

    pub fn https_frontend_command(&mut self, cmd: HttpFrontendCmd) -> Result<(), anyhow::Error> {
        match cmd {
            HttpFrontendCmd::Add {
                hostname,
                path_begin,
                address,
                method,
                route,
            } => self.order_command(ProxyRequestData::AddHttpsFrontend(HttpFrontend {
                route: route.into(),
                address,
                hostname,
                path: PathRule::Prefix(path_begin.unwrap_or_else(|| "".to_string())),
                method: method.map(String::from),
                position: RulePosition::Tree,
            })),
            HttpFrontendCmd::Remove {
                hostname,
                path_begin,
                address,
                method,
                route,
            } => self.order_command(ProxyRequestData::RemoveHttpsFrontend(HttpFrontend {
                route: route.into(),
                address,
                hostname,
                path: PathRule::Prefix(path_begin.unwrap_or_else(|| "".to_string())),
                method: method.map(String::from),
                position: RulePosition::Tree,
            })),
        }
    }

    pub fn https_listener_command(&mut self, cmd: HttpsListenerCmd) -> Result<(), anyhow::Error> {
        match cmd {
            HttpsListenerCmd::Add {
                address,
                public_address,
                answer_404,
                answer_503,
                tls_versions,
                cipher_list,
                rustls_cipher_list,
                expect_proxy,
                sticky_name,
            } => {
                let mut listener = Listener::new(address, FileListenerProtocolConfig::Https);
                listener.public_address = public_address;
                listener.answer_404 = answer_404;
                listener.answer_503 = answer_503;
                listener.expect_proxy = Some(expect_proxy);
                if let Some(sticky_name) = sticky_name {
                    listener.sticky_name = sticky_name;
                }
                listener.cipher_list = cipher_list;
                listener.tls_versions = if tls_versions.is_empty() {
                    None
                } else {
                    Some(tls_versions)
                };
                listener.rustls_cipher_list = if rustls_cipher_list.is_empty() {
                    None
                } else {
                    Some(rustls_cipher_list)
                };
                match listener.to_tls(None, None, None) {
                    Some(conf) => self.order_command(ProxyRequestData::AddHttpsListener(conf)),
                    None => bail!("Error creating HTTPS listener"),
                }
            }
            HttpsListenerCmd::Remove { address } => {
                self.remove_listener(address, ListenerType::HTTPS)
            }
            HttpsListenerCmd::Activate { address } => {
                self.activate_listener(address, ListenerType::HTTPS)
            }
            HttpsListenerCmd::Deactivate { address } => {
                self.deactivate_listener(address, ListenerType::HTTPS)
            }
        }
    }

    pub fn http_listener_command(&mut self, cmd: HttpListenerCmd) -> Result<(), anyhow::Error> {
        match cmd {
            HttpListenerCmd::Add {
                address,
                public_address,
                answer_404,
                answer_503,
                expect_proxy,
                sticky_name,
            } => {
                let mut listener = Listener::new(address, FileListenerProtocolConfig::Http);
                listener.public_address = public_address;
                listener.answer_404 = answer_404;
                listener.answer_503 = answer_503;
                listener.expect_proxy = Some(expect_proxy);
                if let Some(sticky_name) = sticky_name {
                    listener.sticky_name = sticky_name;
                }
                match listener.to_http(None, None, None) {
                    Some(conf) => self.order_command(ProxyRequestData::AddHttpListener(conf)),
                    None => bail!("Error creating HTTP listener"),
                }
            }
            HttpListenerCmd::Remove { address } => {
                self.remove_listener(address, ListenerType::HTTP)
            }
            HttpListenerCmd::Activate { address } => {
                self.activate_listener(address, ListenerType::HTTP)
            }
            HttpListenerCmd::Deactivate { address } => {
                self.deactivate_listener(address, ListenerType::HTTP)
            }
        }
    }

    pub fn tcp_listener_command(&mut self, cmd: TcpListenerCmd) -> Result<(), anyhow::Error> {
        match cmd {
            TcpListenerCmd::Add {
                address,
                public_address,
                expect_proxy,
            } => self.order_command(ProxyRequestData::AddTcpListener(TcpListener {
                address,
                public_address,
                expect_proxy,
                front_timeout: 60,
                back_timeout: 30,
                connect_timeout: 3,
            })),
            TcpListenerCmd::Remove { address } => self.remove_listener(address, ListenerType::TCP),
            TcpListenerCmd::Activate { address } => {
                self.activate_listener(address, ListenerType::TCP)
            }
            TcpListenerCmd::Deactivate { address } => {
                self.deactivate_listener(address, ListenerType::TCP)
            }
        }
    }

    pub fn remove_listener(
        &mut self,
        address: SocketAddr,
        proxy: ListenerType,
    ) -> Result<(), anyhow::Error> {
        self.order_command(ProxyRequestData::RemoveListener(RemoveListener {
            address,
            proxy,
        }))
    }

    pub fn activate_listener(
        &mut self,
        address: SocketAddr,
        proxy: ListenerType,
    ) -> Result<(), anyhow::Error> {
        self.order_command(ProxyRequestData::ActivateListener(ActivateListener {
            address,
            proxy,
            from_scm: false,
        }))
    }

    pub fn deactivate_listener(
        &mut self,
        address: SocketAddr,
        proxy: ListenerType,
    ) -> Result<(), anyhow::Error> {
        self.order_command(ProxyRequestData::DeactivateListener(DeactivateListener {
            address,
            proxy,
            to_scm: false,
        }))
    }

    pub fn logging_filter(&mut self, filter: &LoggingLevel) -> Result<(), anyhow::Error> {
        self.order_command(ProxyRequestData::Logging(filter.to_string().to_lowercase()))
    }

    pub fn add_certificate(
        &mut self,
        address: SocketAddr,
        certificate_path: &str,
        certificate_chain_path: &str,
        key_path: &str,
        versions: Vec<TlsVersion>,
    ) -> Result<(), anyhow::Error> {
        let new_certificate =
            load_full_certificate(certificate_path, certificate_chain_path, key_path, versions)
                .with_context(|| "Could not load the full certificate")?;

        self.order_command(ProxyRequestData::AddCertificate(AddCertificate {
            address,
            certificate: new_certificate,
            names: vec![],
            expired_at: None,
        }))
    }

    pub fn replace_certificate(
        &mut self,
        address: SocketAddr,
        new_certificate_path: &str,
        new_certificate_chain_path: &str,
        new_key_path: &str,
        old_certificate_path: Option<&str>,
        old_fingerprint: Option<&str>,
        versions: Vec<TlsVersion>,
    ) -> Result<(), anyhow::Error> {
        let old_fingerprint = match (old_certificate_path, old_fingerprint) {
            (None, None) | (Some(_), Some(_)) => {
                bail!("Error: Please provide either one, the old certificate's path OR its fingerprint")
            }
            (Some(old_certificate_path), None) => {
                get_fingerprint_from_certificate_path(old_certificate_path).with_context(|| {
                    "Could not retrieve the fingerprint from the given certificate path"
                })?
            }
            (None, Some(fingerprint)) => decode_fingerprint(fingerprint)
                .with_context(|| "Error decoding the given fingerprint")?,
        };

        let new_certificate = load_full_certificate(
            new_certificate_path,
            new_certificate_chain_path,
            new_key_path,
            versions,
        )
        .with_context(|| "Could not load the full certificate")?;

        self.order_command(ProxyRequestData::ReplaceCertificate(ReplaceCertificate {
            address,
            new_certificate,
            old_fingerprint,
            new_names: vec![],
            new_expired_at: None,
        }))?;

        Ok(())
    }

    pub fn remove_certificate(
        &mut self,
        address: SocketAddr,
        certificate_path: Option<&str>,
        fingerprint: Option<&str>,
    ) -> Result<(), anyhow::Error> {
        let fingerprint = match (certificate_path, fingerprint) {
            (None, None) | (Some(_), Some(_)) => {
                bail!("Error: Please provide either one, the path OR the fingerprint of the certificate")
            }
            (Some(certificate_path), None) => {
                get_fingerprint_from_certificate_path(certificate_path).with_context(|| {
                    "Could not retrieve the finger print from the given certificate path"
                })?
            }
            (None, Some(fingerprint)) => decode_fingerprint(fingerprint)
                .with_context(|| "Error decoding the given fingerprint")?,
        };

        self.order_command(ProxyRequestData::RemoveCertificate(RemoveCertificate {
            address,
            fingerprint,
        }))
    }
}

fn get_fingerprint_from_certificate_path(
    certificate_path: &str,
) -> anyhow::Result<CertificateFingerprint> {
    let bytes = Config::load_file_bytes(certificate_path).with_context(|| {
        format!(
            "could not load certificate file on path {}",
            certificate_path
        )
    })?;

    let parsed_bytes = calculate_fingerprint(&bytes).with_context(|| {
        format!(
            "could not calculate fingerprint for the certificate at {}",
            certificate_path
        )
    })?;

    Ok(CertificateFingerprint(parsed_bytes))
}

fn decode_fingerprint(fingerprint: &str) -> anyhow::Result<CertificateFingerprint> {
    let bytes = hex::decode(fingerprint)
        .with_context(|| "Failed at decoding the string (expected hexadecimal data)")?;
    Ok(CertificateFingerprint(bytes))
}

fn load_full_certificate(
    certificate_path: &str,
    certificate_chain_path: &str,
    key_path: &str,
    versions: Vec<TlsVersion>,
) -> Result<CertificateAndKey, anyhow::Error> {
    let certificate = Config::load_file(certificate_path).with_context(|| {
        format!(
            "Could not load certificate file on path {}",
            certificate_path
        )
    })?;

    let certificate_chain = Config::load_file(certificate_chain_path)
        .map(split_certificate_chain)
        .with_context(|| {
            format!(
                "could not load certificate chain on path: {}",
                certificate_chain_path
            )
        })?;

    let key = Config::load_file(key_path)
        .with_context(|| format!("Could not load key file on path {}", key_path))?;

    Ok(CertificateAndKey {
        certificate,
        certificate_chain,
        key,
        versions,
    })
}
<|endoftext|>


<commit_before>// Copyright 2012 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

/*!
 * A simple map based on a vector for small integer keys. Space requirements
 * are O(highest integer key).
 */

use core::container::{Container, Mutable, Map, Set};
use core::iter::{BaseIter, ReverseIter};
use core::option::{Some, None};
use core::prelude::*;

pub struct SmallIntMap<T> {
    priv v: ~[Option<T>],
}

impl<V> BaseIter<(uint, &'self V)> for SmallIntMap<V> {
    /// Visit all key-value pairs in order
    fn each(&self, it: &fn(&(uint, &'self V)) -> bool) {
        for uint::range(0, self.v.len()) |i| {
            match self.v[i] {
              Some(ref elt) => if !it(&(i, elt)) { break },
              None => ()
            }
        }
    }

    fn size_hint(&self) -> Option<uint> { Some(self.len()) }
}

impl<V> ReverseIter<(uint, &'self V)> for SmallIntMap<V> {
    /// Visit all key-value pairs in reverse order
    fn each_reverse(&self, it: &fn(&(uint, &'self V)) -> bool) {
        for uint::range_rev(self.v.len(), 0) |i| {
            match self.v[i - 1] {
              Some(ref elt) => if !it(&(i - 1, elt)) { break },
              None => ()
            }
        }
    }
}

impl<V> Container for SmallIntMap<V> {
    /// Return the number of elements in the map
    fn len(&const self) -> uint {
        let mut sz = 0;
        for uint::range(0, vec::uniq_len(&const self.v)) |i| {
            match self.v[i] {
                Some(_) => sz += 1,
                None => {}
            }
        }
        sz
    }

    /// Return true if the map contains no elements
    fn is_empty(&const self) -> bool { self.len() == 0 }
}

impl<V> Mutable for SmallIntMap<V> {
    /// Clear the map, removing all key-value pairs.
    fn clear(&mut self) { self.v.clear() }
}

impl<V> Map<uint, V> for SmallIntMap<V> {
    /// Return true if the map contains a value for the specified key
    fn contains_key(&self, key: &uint) -> bool {
        self.find(key).is_some()
    }

    /// Visit all keys in order
    fn each_key(&self, blk: &fn(key: &uint) -> bool) {
        self.each(|&(k, _)| blk(&k))
    }

    /// Visit all values in order
    fn each_value(&self, blk: &fn(value: &V) -> bool) {
        self.each(|&(_, v)| blk(v))
    }

    /// Visit all key-value pairs in order
    fn mutate_values(&mut self, it: &fn(&uint, &'self mut V) -> bool) {
        for uint::range(0, self.v.len()) |i| {
            match self.v[i] {
              Some(ref mut elt) => if !it(&i, elt) { break },
              None => ()
            }
        }
    }

    /// Iterate over the map and mutate the contained values
    fn find(&self, key: &uint) -> Option<&'self V> {
        if *key < self.v.len() {
            match self.v[*key] {
              Some(ref value) => Some(value),
              None => None
            }
        } else {
            None
        }
    }

    /// Insert a key-value pair into the map. An existing value for a
    /// key is replaced by the new value. Return true if the key did
    /// not already exist in the map.
    fn insert(&mut self, key: uint, value: V) -> bool {
        let exists = self.contains_key(&key);
        let len = self.v.len();
        if len <= key {
            vec::grow_fn(&mut self.v, key - len + 1, |_| None);
        }
        self.v[key] = Some(value);
        !exists
    }

    /// Remove a key-value pair from the map. Return true if the key
    /// was present in the map, otherwise false.
    fn remove(&mut self, key: &uint) -> bool {
        if *key >= self.v.len() {
            return false;
        }
        let removed = self.v[*key].is_some();
        self.v[*key] = None;
        removed
    }
}

pub impl<V> SmallIntMap<V> {
    /// Create an empty SmallIntMap
    fn new() -> SmallIntMap<V> { SmallIntMap{v: ~[]} }

    fn get(&self, key: &uint) -> &'self V {
        self.find(key).expect("key not present")
    }
}

pub impl<V:Copy> SmallIntMap<V> {
    fn update_with_key(&mut self, key: uint, val: V,
                       ff: &fn(uint, V, V) -> V) -> bool {
        let new_val = match self.find(&key) {
            None => val,
            Some(orig) => ff(key, *orig, val)
        };
        self.insert(key, new_val)
    }

    fn update(&mut self, key: uint, newval: V, ff: &fn(V, V) -> V) -> bool {
        self.update_with_key(key, newval, |_k, v, v1| ff(v,v1))
    }
}

#[cfg(test)]
mod tests {
    use super::SmallIntMap;

    #[test]
    fn test_len() {
        let mut map = SmallIntMap::new();
        fail_unless!(map.len() == 0);
        fail_unless!(map.is_empty());
        fail_unless!(map.insert(5, 20));
        fail_unless!(map.len() == 1);
        fail_unless!(!map.is_empty());
        fail_unless!(map.insert(11, 12));
        fail_unless!(map.len() == 2);
        fail_unless!(!map.is_empty());
        fail_unless!(map.insert(14, 22));
        fail_unless!(map.len() == 3);
        fail_unless!(!map.is_empty());
    }

    #[test]
    fn test_clear() {
        let mut map = SmallIntMap::new();
        fail_unless!(map.insert(5, 20));
        fail_unless!(map.insert(11, 12));
        fail_unless!(map.insert(14, 22));
        map.clear();
        fail_unless!(map.is_empty());
        fail_unless!(map.find(&5).is_none());
        fail_unless!(map.find(&11).is_none());
        fail_unless!(map.find(&14).is_none());
    }

    #[test]
    fn test_insert_with_key() {
        let mut map = SmallIntMap::new();

        // given a new key, initialize it with this new count, given
        // given an existing key, add more to its count
        fn addMoreToCount(_k: uint, v0: uint, v1: uint) -> uint {
            v0 + v1
        }

        fn addMoreToCount_simple(v0: uint, v1: uint) -> uint {
            v0 + v1
        }

        // count integers
        map.update(3, 1, addMoreToCount_simple);
        map.update_with_key(9, 1, addMoreToCount);
        map.update(3, 7, addMoreToCount_simple);
        map.update_with_key(5, 3, addMoreToCount);
        map.update_with_key(3, 2, addMoreToCount);

        // check the total counts
        fail_unless!(map.find(&3).get() == &10);
        fail_unless!(map.find(&5).get() == &3);
        fail_unless!(map.find(&9).get() == &1);

        // sadly, no sevens were counted
        fail_unless!(map.find(&7).is_none());
    }
}
<commit_msg>smallintmap: add find_mut method<commit_after>// Copyright 2012 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

/*!
 * A simple map based on a vector for small integer keys. Space requirements
 * are O(highest integer key).
 */

use core::container::{Container, Mutable, Map, Set};
use core::iter::{BaseIter, ReverseIter};
use core::option::{Some, None};
use core::prelude::*;

pub struct SmallIntMap<T> {
    priv v: ~[Option<T>],
}

impl<V> BaseIter<(uint, &'self V)> for SmallIntMap<V> {
    /// Visit all key-value pairs in order
    fn each(&self, it: &fn(&(uint, &'self V)) -> bool) {
        for uint::range(0, self.v.len()) |i| {
            match self.v[i] {
              Some(ref elt) => if !it(&(i, elt)) { break },
              None => ()
            }
        }
    }

    fn size_hint(&self) -> Option<uint> { Some(self.len()) }
}

impl<V> ReverseIter<(uint, &'self V)> for SmallIntMap<V> {
    /// Visit all key-value pairs in reverse order
    fn each_reverse(&self, it: &fn(&(uint, &'self V)) -> bool) {
        for uint::range_rev(self.v.len(), 0) |i| {
            match self.v[i - 1] {
              Some(ref elt) => if !it(&(i - 1, elt)) { break },
              None => ()
            }
        }
    }
}

impl<V> Container for SmallIntMap<V> {
    /// Return the number of elements in the map
    fn len(&const self) -> uint {
        let mut sz = 0;
        for uint::range(0, vec::uniq_len(&const self.v)) |i| {
            match self.v[i] {
                Some(_) => sz += 1,
                None => {}
            }
        }
        sz
    }

    /// Return true if the map contains no elements
    fn is_empty(&const self) -> bool { self.len() == 0 }
}

impl<V> Mutable for SmallIntMap<V> {
    /// Clear the map, removing all key-value pairs.
    fn clear(&mut self) { self.v.clear() }
}

impl<V> Map<uint, V> for SmallIntMap<V> {
    /// Return true if the map contains a value for the specified key
    fn contains_key(&self, key: &uint) -> bool {
        self.find(key).is_some()
    }

    /// Visit all keys in order
    fn each_key(&self, blk: &fn(key: &uint) -> bool) {
        self.each(|&(k, _)| blk(&k))
    }

    /// Visit all values in order
    fn each_value(&self, blk: &fn(value: &V) -> bool) {
        self.each(|&(_, v)| blk(v))
    }

    /// Iterate over the map and mutate the contained values
    fn mutate_values(&mut self, it: &fn(&uint, &'self mut V) -> bool) {
        for uint::range(0, self.v.len()) |i| {
            match self.v[i] {
              Some(ref mut elt) => if !it(&i, elt) { break },
              None => ()
            }
        }
    }

    /// Return a reference to the value corresponding to the key
    fn find(&self, key: &uint) -> Option<&'self V> {
        if *key < self.v.len() {
            match self.v[*key] {
              Some(ref value) => Some(value),
              None => None
            }
        } else {
            None
        }
    }

    /// Insert a key-value pair into the map. An existing value for a
    /// key is replaced by the new value. Return true if the key did
    /// not already exist in the map.
    fn insert(&mut self, key: uint, value: V) -> bool {
        let exists = self.contains_key(&key);
        let len = self.v.len();
        if len <= key {
            vec::grow_fn(&mut self.v, key - len + 1, |_| None);
        }
        self.v[key] = Some(value);
        !exists
    }

    /// Remove a key-value pair from the map. Return true if the key
    /// was present in the map, otherwise false.
    fn remove(&mut self, key: &uint) -> bool {
        if *key >= self.v.len() {
            return false;
        }
        let removed = self.v[*key].is_some();
        self.v[*key] = None;
        removed
    }
}

pub impl<V> SmallIntMap<V> {
    /// Create an empty SmallIntMap
    fn new() -> SmallIntMap<V> { SmallIntMap{v: ~[]} }

    fn get(&self, key: &uint) -> &'self V {
        self.find(key).expect("key not present")
    }

    /// Return a mutable reference to the value corresponding to the key
    fn find_mut(&mut self, key: &uint) -> Option<&'self mut V> {
        if *key < self.v.len() {
            match self.v[*key] {
              Some(ref mut value) => Some(value),
              None => None
            }
        } else {
            None
        }
    }
}

pub impl<V:Copy> SmallIntMap<V> {
    fn update_with_key(&mut self, key: uint, val: V,
                       ff: &fn(uint, V, V) -> V) -> bool {
        let new_val = match self.find(&key) {
            None => val,
            Some(orig) => ff(key, *orig, val)
        };
        self.insert(key, new_val)
    }

    fn update(&mut self, key: uint, newval: V, ff: &fn(V, V) -> V) -> bool {
        self.update_with_key(key, newval, |_k, v, v1| ff(v,v1))
    }
}

#[cfg(test)]
mod tests {
    use super::SmallIntMap;
    use core::prelude::*;

    #[test]
    fn test_find_mut() {
        let mut m = SmallIntMap::new();
        fail_unless!(m.insert(1, 12));
        fail_unless!(m.insert(2, 8));
        fail_unless!(m.insert(5, 14));
        let new = 100;
        match m.find_mut(&5) {
            None => fail!(), Some(x) => *x = new
        }
        assert_eq!(m.find(&5), Some(&new));
    }

    #[test]
    fn test_len() {
        let mut map = SmallIntMap::new();
        fail_unless!(map.len() == 0);
        fail_unless!(map.is_empty());
        fail_unless!(map.insert(5, 20));
        fail_unless!(map.len() == 1);
        fail_unless!(!map.is_empty());
        fail_unless!(map.insert(11, 12));
        fail_unless!(map.len() == 2);
        fail_unless!(!map.is_empty());
        fail_unless!(map.insert(14, 22));
        fail_unless!(map.len() == 3);
        fail_unless!(!map.is_empty());
    }

    #[test]
    fn test_clear() {
        let mut map = SmallIntMap::new();
        fail_unless!(map.insert(5, 20));
        fail_unless!(map.insert(11, 12));
        fail_unless!(map.insert(14, 22));
        map.clear();
        fail_unless!(map.is_empty());
        fail_unless!(map.find(&5).is_none());
        fail_unless!(map.find(&11).is_none());
        fail_unless!(map.find(&14).is_none());
    }

    #[test]
    fn test_insert_with_key() {
        let mut map = SmallIntMap::new();

        // given a new key, initialize it with this new count, given
        // given an existing key, add more to its count
        fn addMoreToCount(_k: uint, v0: uint, v1: uint) -> uint {
            v0 + v1
        }

        fn addMoreToCount_simple(v0: uint, v1: uint) -> uint {
            v0 + v1
        }

        // count integers
        map.update(3, 1, addMoreToCount_simple);
        map.update_with_key(9, 1, addMoreToCount);
        map.update(3, 7, addMoreToCount_simple);
        map.update_with_key(5, 3, addMoreToCount);
        map.update_with_key(3, 2, addMoreToCount);

        // check the total counts
        fail_unless!(map.find(&3).get() == &10);
        fail_unless!(map.find(&5).get() == &3);
        fail_unless!(map.find(&9).get() == &1);

        // sadly, no sevens were counted
        fail_unless!(map.find(&7).is_none());
    }
}
<|endoftext|>


<commit_before>use std::{
    collections::BTreeMap,
    io::{self, ErrorKind},
    mem,
    net::{IpAddr, SocketAddr},
    pin::Pin,
    sync::Arc,
    task::{Context, Poll},
    time::Duration as StdDuration,
};

use log::{error, trace};
use parking_lot::Mutex as ParkingMutex;
use shadowsocks::relay::socks5::Address;
use smoltcp::{
    iface::{Interface, InterfaceBuilder, Routes, SocketHandle},
    phy::{DeviceCapabilities, Medium},
    socket::{TcpSocket, TcpSocketBuffer, TcpState},
    time::{Duration, Instant},
    wire::{IpAddress, IpCidr, Ipv4Address, Ipv6Address, TcpPacket},
};
use tokio::{
    io::{AsyncRead, AsyncWrite, ReadBuf},
    sync::{mpsc, Notify},
    task::JoinHandle,
    time,
};

use crate::local::{
    context::ServiceContext,
    loadbalancing::PingBalancer,
    net::AutoProxyClientStream,
    utils::{establish_tcp_tunnel, to_ipv4_mapped},
};

use super::virt_device::VirtTunDevice;

struct TcpSocketManager {
    iface: Interface<'static, VirtTunDevice>,
    manager_notify: Arc<Notify>,
}

impl TcpSocketManager {
    fn notify(&self) {
        self.manager_notify.notify_waiters();
    }
}

type SharedTcpSocketManager = Arc<ParkingMutex<TcpSocketManager>>;

struct TcpConnection {
    socket_handle: SocketHandle,
    manager: SharedTcpSocketManager,
}

impl Drop for TcpConnection {
    fn drop(&mut self) {
        let mut manager = self.manager.lock();
        let socket = manager.iface.get_socket::<TcpSocket>(self.socket_handle);
        socket.close();
    }
}

impl TcpConnection {
    fn new(socket: TcpSocket<'static>, manager: SharedTcpSocketManager) -> TcpConnection {
        let socket_handle = {
            let mut manager = manager.lock();
            manager.iface.add_socket(socket)
        };

        TcpConnection { socket_handle, manager }
    }
}

impl AsyncRead for TcpConnection {
    fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
        let mut manager = self.manager.lock();
        {
            let socket = manager.iface.get_socket::<TcpSocket>(self.socket_handle);
            if !socket.is_open() {
                return Ok(()).into();
            }

            if socket.can_recv() {
                let recv_buf = unsafe { mem::transmute::<_, &mut [u8]>(buf.unfilled_mut()) };
                let n = match socket.recv_slice(recv_buf) {
                    Ok(n) => n,
                    Err(err) => return Err(io::Error::new(ErrorKind::Other, err)).into(),
                };
                buf.advance(n);
            } else {
                socket.register_recv_waker(cx.waker());
                return Poll::Pending;
            }
        }

        manager.notify();
        Ok(()).into()
    }
}

impl AsyncWrite for TcpConnection {
    fn poll_write(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8]) -> Poll<io::Result<usize>> {
        let mut manager = self.manager.lock();
        let n = {
            let socket = manager.iface.get_socket::<TcpSocket>(self.socket_handle);
            if !socket.is_open() {
                return Err(ErrorKind::BrokenPipe.into()).into();
            }
            if socket.can_send() {
                match socket.send_slice(buf) {
                    Ok(n) => n,
                    Err(err) => return Err(io::Error::new(ErrorKind::Other, err)).into(),
                }
            } else {
                socket.register_send_waker(cx.waker());
                return Poll::Pending;
            }
        };

        manager.notify();
        Ok(n).into()
    }

    fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        Ok(()).into()
    }

    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        let mut manager = self.manager.lock();
        {
            let socket = manager.iface.get_socket::<TcpSocket>(self.socket_handle);
            // close the transmission half.
            if socket.is_open() {
                socket.close();
            }

            if socket.state() != TcpState::Closed {
                socket.register_send_waker(cx.waker());
                return Poll::Pending;
            }
        }
        manager.notify();
        Ok(()).into()
    }
}

pub struct TcpTun {
    context: Arc<ServiceContext>,
    manager: SharedTcpSocketManager,
    manager_handle: JoinHandle<()>,
    manager_notify: Arc<Notify>,
    balancer: PingBalancer,
    iface_rx: mpsc::UnboundedReceiver<Vec<u8>>,
    iface_tx: mpsc::Sender<Vec<u8>>,
}

impl Drop for TcpTun {
    fn drop(&mut self) {
        self.manager_handle.abort();
    }
}

impl TcpTun {
    pub fn new(context: Arc<ServiceContext>, balancer: PingBalancer, mtu: u32) -> TcpTun {
        let mut capabilities = DeviceCapabilities::default();
        capabilities.medium = Medium::Ip;
        capabilities.max_transmission_unit = mtu as usize;

        let (virt, iface_rx, iface_tx) = VirtTunDevice::new(capabilities);

        let iface_builder = InterfaceBuilder::new(virt, vec![]);
        let iface_ipaddrs = [
            IpCidr::new(IpAddress::v4(0, 0, 0, 1), 0),
            IpCidr::new(IpAddress::v6(0, 0, 0, 0, 0, 0, 0, 1), 0),
        ];
        let mut iface_routes = Routes::new(BTreeMap::new());
        iface_routes
            .add_default_ipv4_route(Ipv4Address::new(0, 0, 0, 1))
            .expect("IPv4 route");
        iface_routes
            .add_default_ipv6_route(Ipv6Address::new(0, 0, 0, 0, 0, 0, 0, 1))
            .expect("IPv6 route");
        let iface = iface_builder
            .any_ip(true)
            .ip_addrs(iface_ipaddrs)
            .routes(iface_routes)
            .finalize();

        let manager_notify = Arc::new(Notify::new());
        let manager = Arc::new(ParkingMutex::new(TcpSocketManager {
            iface,
            manager_notify: manager_notify.clone(),
        }));

        let manager_handle = {
            let manager = manager.clone();
            let manager_notify = manager_notify.clone();
            tokio::spawn(async move {
                loop {
                    let next_duration = {
                        let mut manager = manager.lock();

                        let before_poll = Instant::now();
                        let updated_sockets = match manager.iface.poll(before_poll) {
                            Ok(u) => u,
                            Err(err) => {
                                error!("VirtDevice::poll error: {}", err);
                                false
                            }
                        };

                        let after_poll = Instant::now();

                        if updated_sockets {
                            trace!("VirtDevice::poll costed {}", after_poll - before_poll);
                        }

                        let next_duration = manager
                            .iface
                            .poll_delay(after_poll)
                            .unwrap_or(Duration::from_millis(50));

                        next_duration
                    };

                    tokio::task::yield_now().await;

                    tokio::select! {
                        _ = time::sleep(StdDuration::from(next_duration)) => {}
                        _ = manager_notify.notified() => {}
                    }
                }
            })
        };

        TcpTun {
            context,
            manager,
            manager_handle,
            manager_notify,
            balancer,
            iface_rx,
            iface_tx,
        }
    }

    pub async fn handle_packet(
        &mut self,
        src_addr: SocketAddr,
        dst_addr: SocketAddr,
        tcp_packet: &TcpPacket<&[u8]>,
    ) -> io::Result<()> {
        // TCP first handshake packet, create a new Connection
        if tcp_packet.syn() && !tcp_packet.ack() {
            let accept_opts = self.context.accept_opts();
            let send_buffer_size = accept_opts.tcp.send_buffer_size.unwrap_or(4096);
            let recv_buffer_size = accept_opts.tcp.recv_buffer_size.unwrap_or(4096);

            let mut socket = TcpSocket::new(
                TcpSocketBuffer::new(vec![0u8; recv_buffer_size as usize]),
                TcpSocketBuffer::new(vec![0u8; send_buffer_size as usize]),
            );
            socket.set_ack_delay(None);
            if let Err(err) = socket.listen(dst_addr) {
                return Err(io::Error::new(ErrorKind::Other, err));
            }

            trace!("created TCP connection for {} <-> {}", src_addr, dst_addr);

            let connection = TcpConnection::new(socket, self.manager.clone());

            // establish a tunnel
            let context = self.context.clone();
            let balancer = self.balancer.clone();
            tokio::spawn(async move {
                if let Err(err) = handle_redir_client(context, balancer, connection, src_addr, dst_addr).await {
                    error!("TCP tunnel failure, {} <-> {}, error: {}", src_addr, dst_addr, err);
                }
            });
        }

        Ok(())
    }

    pub async fn drive_interface_state(&mut self, frame: &[u8]) {
        if let Err(..) = self.iface_tx.send(frame.to_vec()).await {
            panic!("interface send channel closed unexpectly");
        }

        // Wake up and poll the interface.
        self.manager_notify.notify_waiters();
    }

    pub async fn recv_packet(&mut self) -> Vec<u8> {
        match self.iface_rx.recv().await {
            Some(v) => v,
            None => unreachable!("channel closed unexpectedly"),
        }
    }
}

/// Established Client Transparent Proxy
///
/// This method must be called after handshaking with client (for example, socks5 handshaking)
async fn establish_client_tcp_redir<'a>(
    context: Arc<ServiceContext>,
    balancer: PingBalancer,
    mut stream: TcpConnection,
    peer_addr: SocketAddr,
    addr: &Address,
) -> io::Result<()> {
    let server = balancer.best_tcp_server();
    let svr_cfg = server.server_config();

    let mut remote = AutoProxyClientStream::connect(context, &server, addr).await?;

    establish_tcp_tunnel(svr_cfg, &mut stream, &mut remote, peer_addr, addr).await
}

async fn handle_redir_client(
    context: Arc<ServiceContext>,
    balancer: PingBalancer,
    s: TcpConnection,
    peer_addr: SocketAddr,
    mut daddr: SocketAddr,
) -> io::Result<()> {
    // Get forward address from socket
    //
    // Try to convert IPv4 mapped IPv6 address for dual-stack mode.
    if let SocketAddr::V6(ref a) = daddr {
        if let Some(v4) = to_ipv4_mapped(a.ip()) {
            daddr = SocketAddr::new(IpAddr::from(v4), a.port());
        }
    }
    let target_addr = Address::from(daddr);
    establish_client_tcp_redir(context, balancer, s, peer_addr, &target_addr).await
}
<commit_msg>SocketHandle should be removed from iface when connection closed<commit_after>use std::{
    collections::BTreeMap,
    io::{self, ErrorKind},
    mem,
    net::{IpAddr, SocketAddr},
    pin::Pin,
    sync::Arc,
    task::{Context, Poll},
    time::Duration as StdDuration,
};

use log::{error, trace};
use parking_lot::Mutex as ParkingMutex;
use shadowsocks::relay::socks5::Address;
use smoltcp::{
    iface::{Interface, InterfaceBuilder, Routes, SocketHandle},
    phy::{DeviceCapabilities, Medium},
    socket::{TcpSocket, TcpSocketBuffer, TcpState},
    time::{Duration, Instant},
    wire::{IpAddress, IpCidr, Ipv4Address, Ipv6Address, TcpPacket},
};
use tokio::{
    io::{AsyncRead, AsyncWrite, ReadBuf},
    sync::{mpsc, Notify},
    task::JoinHandle,
    time,
};

use crate::local::{
    context::ServiceContext,
    loadbalancing::PingBalancer,
    net::AutoProxyClientStream,
    utils::{establish_tcp_tunnel, to_ipv4_mapped},
};

use super::virt_device::VirtTunDevice;

struct TcpSocketManager {
    iface: Interface<'static, VirtTunDevice>,
    manager_notify: Arc<Notify>,
}

impl TcpSocketManager {
    fn notify(&self) {
        self.manager_notify.notify_waiters();
    }
}

type SharedTcpSocketManager = Arc<ParkingMutex<TcpSocketManager>>;

struct TcpConnection {
    socket_handle: SocketHandle,
    manager: SharedTcpSocketManager,
}

impl Drop for TcpConnection {
    fn drop(&mut self) {
        let mut manager = self.manager.lock();
        manager.iface.remove_socket(self.socket_handle);
    }
}

impl TcpConnection {
    fn new(socket: TcpSocket<'static>, manager: SharedTcpSocketManager) -> TcpConnection {
        let socket_handle = {
            let mut manager = manager.lock();
            manager.iface.add_socket(socket)
        };

        TcpConnection { socket_handle, manager }
    }
}

impl AsyncRead for TcpConnection {
    fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
        let mut manager = self.manager.lock();
        {
            let socket = manager.iface.get_socket::<TcpSocket>(self.socket_handle);
            if !socket.is_open() {
                return Ok(()).into();
            }

            if socket.can_recv() {
                let recv_buf = unsafe { mem::transmute::<_, &mut [u8]>(buf.unfilled_mut()) };
                let n = match socket.recv_slice(recv_buf) {
                    Ok(n) => n,
                    Err(err) => return Err(io::Error::new(ErrorKind::Other, err)).into(),
                };
                buf.advance(n);
            } else {
                socket.register_recv_waker(cx.waker());
                return Poll::Pending;
            }
        }

        manager.notify();
        Ok(()).into()
    }
}

impl AsyncWrite for TcpConnection {
    fn poll_write(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8]) -> Poll<io::Result<usize>> {
        let mut manager = self.manager.lock();
        let n = {
            let socket = manager.iface.get_socket::<TcpSocket>(self.socket_handle);
            if !socket.is_open() {
                return Err(ErrorKind::BrokenPipe.into()).into();
            }
            if socket.can_send() {
                match socket.send_slice(buf) {
                    Ok(n) => n,
                    Err(err) => return Err(io::Error::new(ErrorKind::Other, err)).into(),
                }
            } else {
                socket.register_send_waker(cx.waker());
                return Poll::Pending;
            }
        };

        manager.notify();
        Ok(n).into()
    }

    fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        Ok(()).into()
    }

    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        let mut manager = self.manager.lock();
        {
            let socket = manager.iface.get_socket::<TcpSocket>(self.socket_handle);
            // close the transmission half.
            if socket.is_open() {
                socket.close();
            }

            if socket.state() != TcpState::Closed {
                socket.register_send_waker(cx.waker());
                return Poll::Pending;
            }
        }
        manager.notify();
        Ok(()).into()
    }
}

pub struct TcpTun {
    context: Arc<ServiceContext>,
    manager: SharedTcpSocketManager,
    manager_handle: JoinHandle<()>,
    manager_notify: Arc<Notify>,
    balancer: PingBalancer,
    iface_rx: mpsc::UnboundedReceiver<Vec<u8>>,
    iface_tx: mpsc::Sender<Vec<u8>>,
}

impl Drop for TcpTun {
    fn drop(&mut self) {
        self.manager_handle.abort();
    }
}

impl TcpTun {
    pub fn new(context: Arc<ServiceContext>, balancer: PingBalancer, mtu: u32) -> TcpTun {
        let mut capabilities = DeviceCapabilities::default();
        capabilities.medium = Medium::Ip;
        capabilities.max_transmission_unit = mtu as usize;

        let (virt, iface_rx, iface_tx) = VirtTunDevice::new(capabilities);

        let iface_builder = InterfaceBuilder::new(virt, vec![]);
        let iface_ipaddrs = [
            IpCidr::new(IpAddress::v4(0, 0, 0, 1), 0),
            IpCidr::new(IpAddress::v6(0, 0, 0, 0, 0, 0, 0, 1), 0),
        ];
        let mut iface_routes = Routes::new(BTreeMap::new());
        iface_routes
            .add_default_ipv4_route(Ipv4Address::new(0, 0, 0, 1))
            .expect("IPv4 route");
        iface_routes
            .add_default_ipv6_route(Ipv6Address::new(0, 0, 0, 0, 0, 0, 0, 1))
            .expect("IPv6 route");
        let iface = iface_builder
            .any_ip(true)
            .ip_addrs(iface_ipaddrs)
            .routes(iface_routes)
            .finalize();

        let manager_notify = Arc::new(Notify::new());
        let manager = Arc::new(ParkingMutex::new(TcpSocketManager {
            iface,
            manager_notify: manager_notify.clone(),
        }));

        let manager_handle = {
            let manager = manager.clone();
            let manager_notify = manager_notify.clone();
            tokio::spawn(async move {
                loop {
                    let next_duration = {
                        let mut manager = manager.lock();

                        let before_poll = Instant::now();
                        let updated_sockets = match manager.iface.poll(before_poll) {
                            Ok(u) => u,
                            Err(err) => {
                                error!("VirtDevice::poll error: {}", err);
                                false
                            }
                        };

                        let after_poll = Instant::now();

                        if updated_sockets {
                            trace!("VirtDevice::poll costed {}", after_poll - before_poll);
                        }

                        let next_duration = manager
                            .iface
                            .poll_delay(after_poll)
                            .unwrap_or(Duration::from_millis(50));

                        next_duration
                    };

                    tokio::task::yield_now().await;

                    tokio::select! {
                        _ = time::sleep(StdDuration::from(next_duration)) => {}
                        _ = manager_notify.notified() => {}
                    }
                }
            })
        };

        TcpTun {
            context,
            manager,
            manager_handle,
            manager_notify,
            balancer,
            iface_rx,
            iface_tx,
        }
    }

    pub async fn handle_packet(
        &mut self,
        src_addr: SocketAddr,
        dst_addr: SocketAddr,
        tcp_packet: &TcpPacket<&[u8]>,
    ) -> io::Result<()> {
        // TCP first handshake packet, create a new Connection
        if tcp_packet.syn() && !tcp_packet.ack() {
            let accept_opts = self.context.accept_opts();
            let send_buffer_size = accept_opts.tcp.send_buffer_size.unwrap_or(4096);
            let recv_buffer_size = accept_opts.tcp.recv_buffer_size.unwrap_or(4096);

            let mut socket = TcpSocket::new(
                TcpSocketBuffer::new(vec![0u8; recv_buffer_size as usize]),
                TcpSocketBuffer::new(vec![0u8; send_buffer_size as usize]),
            );
            socket.set_ack_delay(None);
            if let Err(err) = socket.listen(dst_addr) {
                return Err(io::Error::new(ErrorKind::Other, err));
            }

            trace!("created TCP connection for {} <-> {}", src_addr, dst_addr);

            let connection = TcpConnection::new(socket, self.manager.clone());

            // establish a tunnel
            let context = self.context.clone();
            let balancer = self.balancer.clone();
            tokio::spawn(async move {
                if let Err(err) = handle_redir_client(context, balancer, connection, src_addr, dst_addr).await {
                    error!("TCP tunnel failure, {} <-> {}, error: {}", src_addr, dst_addr, err);
                }
            });
        }

        Ok(())
    }

    pub async fn drive_interface_state(&mut self, frame: &[u8]) {
        if let Err(..) = self.iface_tx.send(frame.to_vec()).await {
            panic!("interface send channel closed unexpectly");
        }

        // Wake up and poll the interface.
        self.manager_notify.notify_waiters();
    }

    pub async fn recv_packet(&mut self) -> Vec<u8> {
        match self.iface_rx.recv().await {
            Some(v) => v,
            None => unreachable!("channel closed unexpectedly"),
        }
    }
}

/// Established Client Transparent Proxy
///
/// This method must be called after handshaking with client (for example, socks5 handshaking)
async fn establish_client_tcp_redir<'a>(
    context: Arc<ServiceContext>,
    balancer: PingBalancer,
    mut stream: TcpConnection,
    peer_addr: SocketAddr,
    addr: &Address,
) -> io::Result<()> {
    let server = balancer.best_tcp_server();
    let svr_cfg = server.server_config();

    let mut remote = AutoProxyClientStream::connect(context, &server, addr).await?;

    establish_tcp_tunnel(svr_cfg, &mut stream, &mut remote, peer_addr, addr).await
}

async fn handle_redir_client(
    context: Arc<ServiceContext>,
    balancer: PingBalancer,
    s: TcpConnection,
    peer_addr: SocketAddr,
    mut daddr: SocketAddr,
) -> io::Result<()> {
    // Get forward address from socket
    //
    // Try to convert IPv4 mapped IPv6 address for dual-stack mode.
    if let SocketAddr::V6(ref a) = daddr {
        if let Some(v4) = to_ipv4_mapped(a.ip()) {
            daddr = SocketAddr::new(IpAddr::from(v4), a.port());
        }
    }
    let target_addr = Address::from(daddr);
    establish_client_tcp_redir(context, balancer, s, peer_addr, &target_addr).await
}
<|endoftext|>


<commit_before>//! Mappings of objects received from the API, with optional helper methods for
//! ease of use.
//!
//! Models can optionally have additional helper methods compiled, by enabling
//! the `model` feature.

#[macro_use]
mod utils;

pub mod event;
pub mod permissions;

mod channel;
mod error;
mod gateway;
mod guild;
mod invite;
mod misc;
mod user;
mod voice;
mod webhook;

pub use self::channel::*;
pub use self::error::Error as ModelError;
pub use self::gateway::*;
pub use self::guild::*;
pub use self::invite::*;
pub use self::misc::*;
pub use self::permissions::Permissions;
pub use self::user::*;
pub use self::voice::*;
pub use self::webhook::*;

use chrono::NaiveDateTime;
use internal::prelude::*;
use parking_lot::RwLock;
use self::utils::*;
use serde::de::Visitor;
use std::collections::HashMap;
use std::fmt::{Display, Formatter, Result as FmtResult};
use std::sync::Arc;

#[cfg(feature = "utils")]
use utils::Colour;

fn default_true() -> bool { true }

macro_rules! id_u64 {
    ($(#[$attr:meta] $name:ident;)*) => {
        $(
            #[$attr]
            #[derive(Copy, Clone, Default, Debug, Eq, Hash, PartialOrd, Ord, Serialize)]
            #[allow(derive_hash_xor_eq)]
            pub struct $name(pub u64);

            impl $name {
                /// Retrieves the time that the Id was created at.
                pub fn created_at(&self) -> NaiveDateTime {
                    let offset = (self.0 >> 22) / 1000;

                    NaiveDateTime::from_timestamp(1_420_070_400 + offset as i64, 0)
                }
            }

            // This is a hack so that functions can accept `IntoIterator<Item=IdType>`, and or `IntoIterator<Item=&IdType>`
            impl AsRef<$name> for $name {
                fn as_ref(&self) -> &Self {
                    self
                }
            }

            impl From<u64> for $name {
                fn from(id_as_u64: u64) -> $name {
                    $name(id_as_u64)
                }
            }

            impl PartialEq for $name {
                fn eq(&self, other: &Self) -> bool {
                    self.0 == other.0
                }
            }

            impl PartialEq<u64> for $name {
                fn eq(&self, u: &u64) -> bool {
                    self.0 == *u
                }
            }

            impl Display for $name {
                fn fmt(&self, f: &mut Formatter) -> FmtResult {
                    Display::fmt(&self.0, f)
                }
            }

            impl<'de> Deserialize<'de> for $name {
                fn deserialize<D: Deserializer<'de>>(deserializer: D) -> StdResult<Self, D::Error> {
                    deserializer.deserialize_u64(U64Visitor).map($name)
                }
            }
        )*
    }
}

id_u64! {
    /// An identifier for a Channel
    ChannelId;
    /// An identifier for an Emoji
    EmojiId;
    /// An identifier for a Guild
    GuildId;
    /// An identifier for an Integration
    IntegrationId;
    /// An identifier for a Message
    MessageId;
    /// An identifier for a Role
    RoleId;
    /// An identifier for a User
    UserId;
    /// An identifier for a [`Webhook`](struct.Webhook.html).
    WebhookId;
    /// An identifier for an audit log entry.
    AuditLogEntryId;
}

/// A container for guilds.
///
/// This is used to differentiate whether a guild itself can be used or whether
/// a guild needs to be retrieved from the cache.
#[allow(large_enum_variant)]
#[derive(Clone, Debug)]
pub enum GuildContainer {
    /// A guild which can have its contents directly searched.
    Guild(PartialGuild),
    /// A guild's id, which can be used to search the cache for a guild.
    Id(GuildId),
}

/// Information about a user's application. An application does not necessarily
/// have an associated bot user.
#[derive(Clone, Debug, Deserialize)]
pub struct ApplicationInfo {
    /// The bot user associated with the application. See [`BotApplication`] for
    /// more information.
    ///
    /// [`BotApplication`]: struct.BotApplication.html
    pub bot: Option<BotApplication>,
    /// Indicator of whether the bot is public.
    ///
    /// If a bot is public, anyone may invite it to their [`Guild`]. While a bot
    /// is private, only the owner may add it to a guild.
    ///
    /// [`Guild`]: struct.Guild.html
    #[serde(default = "default_true")]
    pub bot_public: bool,
    /// Indicator of whether the bot requires an OAuth2 code grant.
    pub bot_require_code_grant: bool,
    /// A description of the application, assigned by the application owner.
    pub description: String,
    /// A set of bitflags assigned to the application, which represent gated
    /// feature flags that have been enabled for the application.
    pub flags: Option<u64>,
    /// A hash pointing to the application's icon.
    ///
    /// This is not necessarily equivalent to the bot user's avatar.
    pub icon: Option<String>,
    /// The unique numeric Id of the application.
    pub id: UserId,
    /// The name assigned to the application by the application owner.
    pub name: String,
    /// A list of redirect URIs assigned to the application.
    pub redirect_uris: Vec<String>,
    /// A list of RPC Origins assigned to the application.
    pub rpc_origins: Vec<String>,
    /// The given secret to the application.
    ///
    /// This is not equivalent to the application's bot user's token.
    pub secret: String,
}

/// Information about an application with an application's bot user.
#[derive(Clone, Debug, Deserialize)]
pub struct BotApplication {
    /// The unique Id of the bot user.
    pub id: UserId,
    /// A hash of the avatar, if one is assigned.
    ///
    /// Can be used to generate a full URL to the avatar.
    pub avatar: Option<String>,
    /// Indicator of whether it is a bot.
    #[serde(default)]
    pub bot: bool,
    /// The discriminator assigned to the bot user.
    ///
    /// While discriminators are not unique, the `username#discriminator` pair
    /// is.
    pub discriminator: u16,
    /// The bot user's username.
    pub name: String,
    /// The token used to authenticate as the bot user.
    ///
    /// **Note**: Keep this information private, as untrusted sources can use it
    /// to perform any action with a bot user.
    pub token: String,
}

/// Information about the current application and its owner.
#[derive(Clone, Debug, Deserialize)]
pub struct CurrentApplicationInfo {
    pub description: String,
    pub icon: Option<String>,
    pub id: UserId,
    pub name: String,
    pub owner: User,
    #[serde(default)] pub rpc_origins: Vec<String>,
}

/// The name of a region that a voice server can be located in.
#[derive(Copy, Clone, Debug, Deserialize, Eq, Hash, PartialEq, PartialOrd, Ord, Serialize)]
pub enum Region {
    #[serde(rename = "amsterdam")] Amsterdam,
    #[serde(rename = "brazil")] Brazil,
    #[serde(rename = "eu-central")] EuCentral,
    #[serde(rename = "eu-west")] EuWest,
    #[serde(rename = "frankfurt")] Frankfurt,
    #[serde(rename = "london")] London,
    #[serde(rename = "sydney")] Sydney,
    #[serde(rename = "us-central")] UsCentral,
    #[serde(rename = "us-east")] UsEast,
    #[serde(rename = "us-south")] UsSouth,
    #[serde(rename = "us-west")] UsWest,
    #[serde(rename = "vip-amsterdam")] VipAmsterdam,
    #[serde(rename = "vip-us-east")] VipUsEast,
    #[serde(rename = "vip-us-west")] VipUsWest,
}

impl Region {
    pub fn name(&self) -> &str {
        match *self {
            Region::Amsterdam => "amsterdam",
            Region::Brazil => "brazil",
            Region::EuCentral => "eu-central",
            Region::EuWest => "eu-west",
            Region::Frankfurt => "frankfurt",
            Region::London => "london",
            Region::Sydney => "sydney",
            Region::UsCentral => "us-central",
            Region::UsEast => "us-east",
            Region::UsSouth => "us-south",
            Region::UsWest => "us-west",
            Region::VipAmsterdam => "vip-amsterdam",
            Region::VipUsEast => "vip-us-east",
            Region::VipUsWest => "vip-us-west",
        }
    }
}

use serde::{Deserialize, Deserializer};
use std::result::Result as StdResult;

fn deserialize_sync_user<'de, D: Deserializer<'de>>(deserializer: D)
                                                    -> StdResult<Arc<RwLock<User>>, D::Error> {
    Ok(Arc::new(RwLock::new(User::deserialize(deserializer)?)))
}
<commit_msg>Make the comment for the `AsRef` impl more clear<commit_after>//! Mappings of objects received from the API, with optional helper methods for
//! ease of use.
//!
//! Models can optionally have additional helper methods compiled, by enabling
//! the `model` feature.

#[macro_use]
mod utils;

pub mod event;
pub mod permissions;

mod channel;
mod error;
mod gateway;
mod guild;
mod invite;
mod misc;
mod user;
mod voice;
mod webhook;

pub use self::channel::*;
pub use self::error::Error as ModelError;
pub use self::gateway::*;
pub use self::guild::*;
pub use self::invite::*;
pub use self::misc::*;
pub use self::permissions::Permissions;
pub use self::user::*;
pub use self::voice::*;
pub use self::webhook::*;

use chrono::NaiveDateTime;
use internal::prelude::*;
use parking_lot::RwLock;
use self::utils::*;
use serde::de::Visitor;
use std::collections::HashMap;
use std::fmt::{Display, Formatter, Result as FmtResult};
use std::sync::Arc;

#[cfg(feature = "utils")]
use utils::Colour;

fn default_true() -> bool { true }

macro_rules! id_u64 {
    ($(#[$attr:meta] $name:ident;)*) => {
        $(
            #[$attr]
            #[derive(Copy, Clone, Default, Debug, Eq, Hash, PartialOrd, Ord, Serialize)]
            #[allow(derive_hash_xor_eq)]
            pub struct $name(pub u64);

            impl $name {
                /// Retrieves the time that the Id was created at.
                pub fn created_at(&self) -> NaiveDateTime {
                    let offset = (self.0 >> 22) / 1000;

                    NaiveDateTime::from_timestamp(1_420_070_400 + offset as i64, 0)
                }
            }

            // This is a hack so functions can accept iterators that either:
            // 1. return the id itself (e.g: `MessageId`)
            // 2. return a reference to it (`&MessageId`).
            impl AsRef<$name> for $name {
                fn as_ref(&self) -> &Self {
                    self
                }
            }

            impl From<u64> for $name {
                fn from(id_as_u64: u64) -> $name {
                    $name(id_as_u64)
                }
            }

            impl PartialEq for $name {
                fn eq(&self, other: &Self) -> bool {
                    self.0 == other.0
                }
            }

            impl PartialEq<u64> for $name {
                fn eq(&self, u: &u64) -> bool {
                    self.0 == *u
                }
            }

            impl Display for $name {
                fn fmt(&self, f: &mut Formatter) -> FmtResult {
                    Display::fmt(&self.0, f)
                }
            }

            impl<'de> Deserialize<'de> for $name {
                fn deserialize<D: Deserializer<'de>>(deserializer: D) -> StdResult<Self, D::Error> {
                    deserializer.deserialize_u64(U64Visitor).map($name)
                }
            }
        )*
    }
}

id_u64! {
    /// An identifier for a Channel
    ChannelId;
    /// An identifier for an Emoji
    EmojiId;
    /// An identifier for a Guild
    GuildId;
    /// An identifier for an Integration
    IntegrationId;
    /// An identifier for a Message
    MessageId;
    /// An identifier for a Role
    RoleId;
    /// An identifier for a User
    UserId;
    /// An identifier for a [`Webhook`](struct.Webhook.html).
    WebhookId;
    /// An identifier for an audit log entry.
    AuditLogEntryId;
}

/// A container for guilds.
///
/// This is used to differentiate whether a guild itself can be used or whether
/// a guild needs to be retrieved from the cache.
#[allow(large_enum_variant)]
#[derive(Clone, Debug)]
pub enum GuildContainer {
    /// A guild which can have its contents directly searched.
    Guild(PartialGuild),
    /// A guild's id, which can be used to search the cache for a guild.
    Id(GuildId),
}

/// Information about a user's application. An application does not necessarily
/// have an associated bot user.
#[derive(Clone, Debug, Deserialize)]
pub struct ApplicationInfo {
    /// The bot user associated with the application. See [`BotApplication`] for
    /// more information.
    ///
    /// [`BotApplication`]: struct.BotApplication.html
    pub bot: Option<BotApplication>,
    /// Indicator of whether the bot is public.
    ///
    /// If a bot is public, anyone may invite it to their [`Guild`]. While a bot
    /// is private, only the owner may add it to a guild.
    ///
    /// [`Guild`]: struct.Guild.html
    #[serde(default = "default_true")]
    pub bot_public: bool,
    /// Indicator of whether the bot requires an OAuth2 code grant.
    pub bot_require_code_grant: bool,
    /// A description of the application, assigned by the application owner.
    pub description: String,
    /// A set of bitflags assigned to the application, which represent gated
    /// feature flags that have been enabled for the application.
    pub flags: Option<u64>,
    /// A hash pointing to the application's icon.
    ///
    /// This is not necessarily equivalent to the bot user's avatar.
    pub icon: Option<String>,
    /// The unique numeric Id of the application.
    pub id: UserId,
    /// The name assigned to the application by the application owner.
    pub name: String,
    /// A list of redirect URIs assigned to the application.
    pub redirect_uris: Vec<String>,
    /// A list of RPC Origins assigned to the application.
    pub rpc_origins: Vec<String>,
    /// The given secret to the application.
    ///
    /// This is not equivalent to the application's bot user's token.
    pub secret: String,
}

/// Information about an application with an application's bot user.
#[derive(Clone, Debug, Deserialize)]
pub struct BotApplication {
    /// The unique Id of the bot user.
    pub id: UserId,
    /// A hash of the avatar, if one is assigned.
    ///
    /// Can be used to generate a full URL to the avatar.
    pub avatar: Option<String>,
    /// Indicator of whether it is a bot.
    #[serde(default)]
    pub bot: bool,
    /// The discriminator assigned to the bot user.
    ///
    /// While discriminators are not unique, the `username#discriminator` pair
    /// is.
    pub discriminator: u16,
    /// The bot user's username.
    pub name: String,
    /// The token used to authenticate as the bot user.
    ///
    /// **Note**: Keep this information private, as untrusted sources can use it
    /// to perform any action with a bot user.
    pub token: String,
}

/// Information about the current application and its owner.
#[derive(Clone, Debug, Deserialize)]
pub struct CurrentApplicationInfo {
    pub description: String,
    pub icon: Option<String>,
    pub id: UserId,
    pub name: String,
    pub owner: User,
    #[serde(default)] pub rpc_origins: Vec<String>,
}

/// The name of a region that a voice server can be located in.
#[derive(Copy, Clone, Debug, Deserialize, Eq, Hash, PartialEq, PartialOrd, Ord, Serialize)]
pub enum Region {
    #[serde(rename = "amsterdam")] Amsterdam,
    #[serde(rename = "brazil")] Brazil,
    #[serde(rename = "eu-central")] EuCentral,
    #[serde(rename = "eu-west")] EuWest,
    #[serde(rename = "frankfurt")] Frankfurt,
    #[serde(rename = "london")] London,
    #[serde(rename = "sydney")] Sydney,
    #[serde(rename = "us-central")] UsCentral,
    #[serde(rename = "us-east")] UsEast,
    #[serde(rename = "us-south")] UsSouth,
    #[serde(rename = "us-west")] UsWest,
    #[serde(rename = "vip-amsterdam")] VipAmsterdam,
    #[serde(rename = "vip-us-east")] VipUsEast,
    #[serde(rename = "vip-us-west")] VipUsWest,
}

impl Region {
    pub fn name(&self) -> &str {
        match *self {
            Region::Amsterdam => "amsterdam",
            Region::Brazil => "brazil",
            Region::EuCentral => "eu-central",
            Region::EuWest => "eu-west",
            Region::Frankfurt => "frankfurt",
            Region::London => "london",
            Region::Sydney => "sydney",
            Region::UsCentral => "us-central",
            Region::UsEast => "us-east",
            Region::UsSouth => "us-south",
            Region::UsWest => "us-west",
            Region::VipAmsterdam => "vip-amsterdam",
            Region::VipUsEast => "vip-us-east",
            Region::VipUsWest => "vip-us-west",
        }
    }
}

use serde::{Deserialize, Deserializer};
use std::result::Result as StdResult;

fn deserialize_sync_user<'de, D: Deserializer<'de>>(deserializer: D)
                                                    -> StdResult<Arc<RwLock<User>>, D::Error> {
    Ok(Arc::new(RwLock::new(User::deserialize(deserializer)?)))
}
<|endoftext|>