stream-tungstenite 0.6.1

//! WebSocket connector - establishes WebSocket connections over transports.

use async_trait::async_trait;
use futures_util::{Sink, Stream};
use std::fmt::Debug;
use tokio::net::TcpStream;
use tokio_tungstenite::{
    client_async_tls_with_config, Connector as TlsConnector, MaybeTlsStream, WebSocketStream,
};
use tungstenite::client::IntoClientRequest;
use tungstenite::handshake::client::Response;
use tungstenite::protocol::WebSocketConfig;
use tungstenite::Message;

use crate::error::ConnectError;
use crate::transport::Transport;

/// Type alias for the default WebSocket stream type
pub type DefaultWsStream = WebSocketStream<MaybeTlsStream<TcpStream>>;

/// WebSocket stream trait - abstraction over different stream types
pub trait WsStream:
    Stream<Item = Result<Message, tungstenite::Error>>
    + Sink<Message, Error = tungstenite::Error>
    + Unpin
    + Send
    + 'static
{
}

impl<T> WsStream for T where
    T: Stream<Item = Result<Message, tungstenite::Error>>
        + Sink<Message, Error = tungstenite::Error>
        + Unpin
        + Send
        + 'static
{
}

/// Connector trait - responsible for establishing WebSocket connections
#[async_trait]
pub trait Connector: Send + Sync + 'static {
    /// The WebSocket stream type produced by this connector
    type Stream: WsStream;

    /// Connect to the given URI and return a WebSocket stream
    async fn connect(&self, uri: &str) -> Result<(Self::Stream, Response), ConnectError>;

    /// Get the connector name (for logging/debugging)
    fn name(&self) -> &'static str {
        "connector"
    }
}

/// Default WebSocket connector using TLS transport
#[derive(Clone)]
pub struct DefaultConnector {
    /// WebSocket protocol configuration
    ws_config: Option<WebSocketConfig>,
    /// Whether to disable Nagle's algorithm
    disable_nagle: bool,
    /// TLS connector configuration
    tls_connector: Option<TlsConnector>,
}

// Manual Debug implementation since TlsConnector doesn't implement Debug
impl Debug for DefaultConnector {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("DefaultConnector")
            .field("ws_config", &self.ws_config)
            .field("disable_nagle", &self.disable_nagle)
            .field("tls_connector", &self.tls_connector.is_some())
            .finish()
    }
}

impl DefaultConnector {
    /// Create a new default connector with system default TLS configuration
    #[must_use]
    pub const fn new() -> Self {
        Self {
            ws_config: None,
            disable_nagle: false,
            tls_connector: None,
        }
    }

    /// Set WebSocket configuration
    #[must_use]
    pub const fn with_ws_config(mut self, config: WebSocketConfig) -> Self {
        self.ws_config = Some(config);
        self
    }

    /// Disable Nagle's algorithm (`TCP_NODELAY`)
    #[must_use]
    pub const fn with_nodelay(mut self, nodelay: bool) -> Self {
        self.disable_nagle = nodelay;
        self
    }

    /// Set custom TLS connector
    ///
    /// This allows full control over TLS configuration including:
    /// - Custom CA certificates
    /// - Client certificate authentication
    /// - Certificate verification policies
    /// - Protocol version constraints
    ///
    /// # Example
    /// ```no_run
    /// # use stream_tungstenite::connection::DefaultConnector;
    /// # #[cfg(feature = "native-tls")]
    /// # {
    /// use tokio_tungstenite::Connector;
    /// use native_tls::TlsConnector as NativeTlsConnector;
    ///
    /// let tls = NativeTlsConnector::builder()
    ///     .danger_accept_invalid_certs(true) // For testing only!
    ///     .build()
    ///     .unwrap();
    ///
    /// let connector = DefaultConnector::new()
    ///     .with_tls_connector(Connector::NativeTls(tls));
    /// # }
    /// ```
    #[must_use]
    pub fn with_tls_connector(mut self, connector: TlsConnector) -> Self {
        self.tls_connector = Some(connector);
        self
    }

    /// Create a connector optimized for low latency
    #[must_use]
    pub fn low_latency() -> Self {
        Self {
            ws_config: Some(
                WebSocketConfig::default()
                    .max_message_size(Some(64 << 20)) // 64 MB
                    .max_frame_size(Some(16 << 20)), // 16 MB
            ),
            disable_nagle: true,
            tls_connector: None,
        }
    }

    // === Convenience methods for common TLS configurations ===

    /// Disable certificate verification (DANGEROUS - use only for testing!)
    ///
    /// This creates a TLS connector that accepts any certificate, including:
    /// - Self-signed certificates
    /// - Expired certificates
    /// - Certificates with wrong hostname
    ///
    /// # Security Warning
    /// This makes your connection vulnerable to man-in-the-middle attacks.
    /// **NEVER use this in production!**
    ///
    /// # Example
    /// ```no_run
    /// # use stream_tungstenite::connection::DefaultConnector;
    /// let connector = DefaultConnector::new()
    ///     .danger_accept_invalid_certs()
    ///     .expect("Failed to create insecure connector");
    /// ```
    #[cfg(all(feature = "native-tls", not(feature = "__rustls-tls")))]
    pub fn danger_accept_invalid_certs(mut self) -> Result<Self, ConnectError> {
        use native_tls::TlsConnector as NativeTlsConnector;

        let tls = NativeTlsConnector::builder()
            .danger_accept_invalid_certs(true)
            .danger_accept_invalid_hostnames(true)
            .build()
            .map_err(|e| {
                ConnectError::Tls(format!("Failed to build insecure TLS connector: {e}"))
            })?;

        self.tls_connector = Some(TlsConnector::NativeTls(tls));
        Ok(self)
    }

    /// Disable certificate verification using rustls (DANGEROUS - use only for testing!)
    ///
    /// # Security Warning
    /// This makes your connection vulnerable to man-in-the-middle attacks.
    /// **NEVER use this in production!**
    #[cfg(feature = "__rustls-tls")]
    pub fn danger_accept_invalid_certs(mut self) -> Result<Self, ConnectError> {
        use rustls::client::danger::{HandshakeSignatureValid, ServerCertVerifier};
        use rustls::pki_types::{CertificateDer, ServerName, UnixTime};
        use rustls::{ClientConfig, DigitallySignedStruct, SignatureScheme};
        use std::sync::Arc;

        #[derive(Debug)]
        struct NoVerifier;

        impl ServerCertVerifier for NoVerifier {
            fn verify_server_cert(
                &self,
                _end_entity: &CertificateDer<'_>,
                _intermediates: &[CertificateDer<'_>],
                _server_name: &ServerName<'_>,
                _ocsp_response: &[u8],
                _now: UnixTime,
            ) -> Result<rustls::client::danger::ServerCertVerified, rustls::Error> {
                Ok(rustls::client::danger::ServerCertVerified::assertion())
            }

            fn verify_tls12_signature(
                &self,
                _message: &[u8],
                _cert: &CertificateDer<'_>,
                _dss: &DigitallySignedStruct,
            ) -> Result<HandshakeSignatureValid, rustls::Error> {
                Ok(HandshakeSignatureValid::assertion())
            }

            fn verify_tls13_signature(
                &self,
                _message: &[u8],
                _cert: &CertificateDer<'_>,
                _dss: &DigitallySignedStruct,
            ) -> Result<HandshakeSignatureValid, rustls::Error> {
                Ok(HandshakeSignatureValid::assertion())
            }

            fn supported_verify_schemes(&self) -> Vec<SignatureScheme> {
                vec![
                    SignatureScheme::RSA_PKCS1_SHA256,
                    SignatureScheme::RSA_PKCS1_SHA384,
                    SignatureScheme::RSA_PKCS1_SHA512,
                    SignatureScheme::ECDSA_NISTP256_SHA256,
                    SignatureScheme::ECDSA_NISTP384_SHA384,
                    SignatureScheme::ECDSA_NISTP521_SHA512,
                    SignatureScheme::RSA_PSS_SHA256,
                    SignatureScheme::RSA_PSS_SHA384,
                    SignatureScheme::RSA_PSS_SHA512,
                    SignatureScheme::ED25519,
                ]
            }
        }

        let mut config = ClientConfig::builder()
            .dangerous()
            .with_custom_certificate_verifier(Arc::new(NoVerifier))
            .with_no_client_auth();

        config.alpn_protocols = vec![b"http/1.1".to_vec()];

        self.tls_connector = Some(TlsConnector::Rustls(Arc::new(config)));
        Ok(self)
    }

    /// Add a custom CA certificate to trust
    ///
    /// This is useful when connecting to servers with self-signed certificates
    /// or internal CAs.
    ///
    /// # Example
    /// ```no_run
    /// # use stream_tungstenite::connection::DefaultConnector;
    /// # #[cfg(feature = "native-tls")]
    /// # {
    /// use native_tls::Certificate;
    /// use std::fs;
    ///
    /// let ca_cert = fs::read("ca-cert.pem").unwrap();
    /// let cert = Certificate::from_pem(&ca_cert).unwrap();
    ///
    /// let connector = DefaultConnector::new()
    ///     .with_custom_ca_cert(cert)
    ///     .expect("Failed to add CA cert");
    /// # }
    /// ```
    #[cfg(feature = "native-tls")]
    pub fn with_custom_ca_cert(
        mut self,
        cert: native_tls::Certificate,
    ) -> Result<Self, ConnectError> {
        use native_tls::TlsConnector as NativeTlsConnector;

        let tls = NativeTlsConnector::builder()
            .add_root_certificate(cert)
            .build()
            .map_err(|e| {
                ConnectError::Tls(format!("Failed to build TLS connector with custom CA: {e}"))
            })?;

        self.tls_connector = Some(TlsConnector::NativeTls(tls));
        Ok(self)
    }

    /// Configure client certificate authentication
    ///
    /// This is required when the server uses mutual TLS (mTLS) authentication.
    ///
    /// # Example
    /// ```no_run
    /// # use stream_tungstenite::connection::DefaultConnector;
    /// # #[cfg(feature = "native-tls")]
    /// # {
    /// use native_tls::Identity;
    /// use std::fs;
    ///
    /// let pkcs12 = fs::read("client-cert.p12").unwrap();
    /// let identity = Identity::from_pkcs12(&pkcs12, "password").unwrap();
    ///
    /// let connector = DefaultConnector::new()
    ///     .with_client_identity(identity)
    ///     .expect("Failed to configure client cert");
    /// # }
    /// ```
    #[cfg(feature = "native-tls")]
    pub fn with_client_identity(
        mut self,
        identity: native_tls::Identity,
    ) -> Result<Self, ConnectError> {
        use native_tls::TlsConnector as NativeTlsConnector;

        let tls = NativeTlsConnector::builder()
            .identity(identity)
            .build()
            .map_err(|e| {
                ConnectError::Tls(format!(
                    "Failed to build TLS connector with client identity: {e}"
                ))
            })?;

        self.tls_connector = Some(TlsConnector::NativeTls(tls));
        Ok(self)
    }
}

impl Default for DefaultConnector {
    fn default() -> Self {
        Self::new()
    }
}

#[async_trait]
impl Connector for DefaultConnector {
    type Stream = DefaultWsStream;

    async fn connect(&self, uri: &str) -> Result<(Self::Stream, Response), ConnectError> {
        // Parse the URI and create request
        let request = uri
            .into_client_request()
            .map_err(|e| ConnectError::InvalidUri(format!("Failed to parse URI '{uri}': {e}")))?;

        // Extract host and port
        let host = request
            .uri()
            .host()
            .ok_or_else(|| ConnectError::InvalidUri("No host in URI".into()))?;

        let port = request
            .uri()
            .port_u16()
            .unwrap_or_else(|| match request.uri().scheme_str() {
                Some("wss") => 443,
                _ => 80,
            });

        let addr = format!("{host}:{port}");

        // Establish TCP connection
        let socket = TcpStream::connect(&addr).await.map_err(|e| {
            tracing::debug!(addr = %addr, error = ?e, "TCP connection failed");
            ConnectError::TcpConnect(e.to_string())
        })?;

        // Apply socket options
        if self.disable_nagle {
            socket
                .set_nodelay(true)
                .map_err(|e| ConnectError::Io(format!("Failed to set TCP_NODELAY: {e}")))?;
        }

        // Perform WebSocket upgrade (with TLS if wss://)
        let (ws_stream, response) = client_async_tls_with_config(
            request,
            socket,
            self.ws_config,
            self.tls_connector.clone(),
        )
        .await
        .map_err(|e| {
            tracing::debug!(uri = %uri, error = ?e, "WebSocket connection failed");
            ConnectError::WebSocketUpgrade(e.to_string())
        })?;

        tracing::debug!(uri = %uri, "WebSocket connection established");
        Ok((ws_stream, response))
    }

    fn name(&self) -> &'static str {
        "default"
    }
}

/// Connector that wraps a Transport implementation
#[derive(Debug, Clone)]
pub struct TransportConnector<T: Transport> {
    #[allow(dead_code)]
    transport: T,
    ws_config: Option<WebSocketConfig>,
}

impl<T: Transport> TransportConnector<T> {
    /// Create a new connector with the given transport
    #[must_use]
    pub const fn new(transport: T) -> Self {
        Self {
            transport,
            ws_config: None,
        }
    }

    /// Set WebSocket configuration
    #[must_use]
    pub const fn with_ws_config(mut self, config: WebSocketConfig) -> Self {
        self.ws_config = Some(config);
        self
    }
}

/// Connection info returned after successful connection
#[derive(Debug, Clone)]
pub struct ConnectionInfo {
    /// The URI connected to
    pub uri: String,
    /// HTTP response from the WebSocket upgrade
    pub response_status: u16,
    /// Response headers (selected)
    pub response_headers: Vec<(String, String)>,
}

impl ConnectionInfo {
    /// Create connection info from response
    #[must_use]
    pub fn from_response(uri: &str, response: &Response) -> Self {
        let headers = response
            .headers()
            .iter()
            .map(|(k, v)| (k.to_string(), v.to_str().unwrap_or("").to_string()))
            .collect();

        Self {
            uri: uri.to_string(),
            response_status: response.status().as_u16(),
            response_headers: headers,
        }
    }
}

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

    #[test]
    fn test_default_connector_creation() {
        let connector = DefaultConnector::new();
        assert_eq!(connector.name(), "default");
    }

    #[test]
    fn test_low_latency_connector() {
        let connector = DefaultConnector::low_latency();
        assert!(connector.disable_nagle);
        assert!(connector.ws_config.is_some());
    }
}