mssql_client/client/

connect.rs

//! Connection establishment for SQL Server.
//!
//! This module contains the `impl Client<Disconnected>` block, handling
//! TCP connection, TLS negotiation, PreLogin exchange, and Login7 authentication.

use std::marker::PhantomData;
use std::net::SocketAddr;

use bytes::BytesMut;
use mssql_codec::connection::Connection;
#[cfg(feature = "tls")]
use mssql_tls::{TlsConfig, TlsConnector, TlsNegotiationMode};
use tds_protocol::login7::Login7;
use tds_protocol::packet::MAX_PACKET_SIZE;
use tds_protocol::packet::PacketType;
use tds_protocol::prelogin::{EncryptionLevel, PreLogin};
use tds_protocol::token::{EnvChange, EnvChangeType, Token, TokenParser};
use tokio::net::TcpStream;
use tokio::time::timeout;

use crate::config::Config;
use crate::error::{Error, Result};
#[cfg(feature = "otel")]
use crate::instrumentation::InstrumentationContext;
use crate::state::{Disconnected, Ready};
use crate::statement_cache::StatementCache;

use super::{Client, ConnectionHandle};

impl Client<Disconnected> {
    /// Connect to SQL Server.
    ///
    /// This establishes a connection, performs TLS negotiation (if required),
    /// and authenticates with the server.
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// # use mssql_client::Client;
    /// # async fn ex(config: mssql_client::Config) -> Result<(), mssql_client::Error> {
    /// let client = Client::connect(config).await?;
    /// # let _ = client;
    /// # Ok(())
    /// # }
    /// ```
    pub async fn connect(config: Config) -> Result<Client<Ready>> {
        // FEDAUTH-based credentials (Azure AD / Entra and client certificate)
        // are not yet wired into the login sequence: providers can acquire
        // tokens, but LOGIN7 carries no FEDAUTH feature extension, so the
        // server would receive an empty-credential login and reject it with
        // an opaque error 18456. Fail fast with an actionable error instead.
        // Full FEDAUTH support is tracked in issue #155.
        match &config.credentials {
            mssql_auth::Credentials::SqlServer { .. } => {}
            #[cfg(any(feature = "integrated-auth", feature = "sspi-auth"))]
            mssql_auth::Credentials::Integrated => {}
            // Every other credential type (Azure access token, managed
            // identity, service principal, client certificate) is
            // FEDAUTH-based and rejected here.
            _ => {
                return Err(Error::Config(
                    "Azure AD / Entra and client certificate (FEDAUTH) authentication \
                     are not yet supported: the LOGIN7 FEDAUTH feature extension is not \
                     implemented (tracked in \
                     https://github.com/praxiomlabs/rust-mssql-driver/issues/155). \
                     Use SQL Server or integrated authentication."
                        .into(),
                ));
            }
        }

        let retry = config.retry.clone();
        let max_redirects = config.redirect.max_redirects;
        let follow_redirects = config.redirect.follow_redirects;
        // Overall timeout accounts for retries + redirects per attempt, capped at 5 min.
        let per_attempt = config.timeouts.connect_timeout
            + config.timeouts.tls_timeout
            + config.timeouts.login_timeout;
        let total_attempts = (retry.max_retries + 1) * (max_redirects as u32 + 1);
        let overall = (per_attempt * total_attempts).min(std::time::Duration::from_secs(300));
        let initial_host = config.host.clone();
        let initial_port = config.port;

        let result = timeout(overall, async {
            let mut last_error: Option<Error> = None;

            for retry_attempt in 0..=retry.max_retries {
                if retry_attempt > 0 {
                    let backoff = retry.backoff_for_attempt(retry_attempt);
                    tracing::info!(
                        retry_attempt,
                        backoff_ms = backoff.as_millis() as u64,
                        "retrying connection after transient error"
                    );
                    tokio::time::sleep(backoff).await;
                }

                // Each retry starts fresh with original host/port
                let mut current_config = config.clone();
                let mut redirect_count: u8 = 0;

                let attempt_result = loop {
                    redirect_count += 1;
                    if redirect_count > max_redirects + 1 {
                        break Err(Error::TooManyRedirects { max: max_redirects });
                    }

                    match Self::try_connect(&current_config).await {
                        Ok(client) => break Ok(client),
                        Err(Error::Routing { host, port }) => {
                            if !follow_redirects {
                                break Err(Error::Routing { host, port });
                            }
                            tracing::info!(
                                host = %host,
                                port = port,
                                redirect = redirect_count,
                                max_redirects = max_redirects,
                                "following Azure SQL routing redirect"
                            );
                            current_config = current_config.with_host(&host).with_port(port);
                            continue;
                        }
                        Err(e) => break Err(e),
                    }
                };

                match attempt_result {
                    Ok(client) => return Ok(client),
                    Err(ref e) if e.is_transient() && retry.should_retry(retry_attempt) => {
                        tracing::warn!(
                            retry_attempt,
                            max_retries = retry.max_retries,
                            error = %e,
                            "transient connection error, will retry"
                        );
                        last_error = Some(attempt_result.unwrap_err());
                    }
                    Err(e) => return Err(e),
                }
            }

            // All retries exhausted — return last error
            Err(last_error.expect("at least one attempt was made"))
        })
        .await;

        match result {
            Ok(inner) => inner,
            Err(_elapsed) => Err(Error::ConnectTimeout {
                host: initial_host,
                port: initial_port,
            }),
        }
    }

    async fn try_connect(config: &Config) -> Result<Client<Ready>> {
        // If a named instance is specified, resolve the TCP port via SQL Browser
        let port = if let Some(ref instance) = config.instance {
            let resolved = crate::browser::resolve_instance(
                &config.host,
                instance,
                Some(config.timeouts.connect_timeout),
            )
            .await?;
            tracing::info!(
                host = %config.host,
                instance = %instance,
                resolved_port = resolved,
                database = ?config.database,
                "connecting to named SQL Server instance"
            );
            resolved
        } else {
            tracing::info!(
                host = %config.host,
                port = config.port,
                database = ?config.database,
                "connecting to SQL Server"
            );
            config.port
        };

        // Normalize "." and "(local)" to localhost for TCP.
        // These are standard ADO.NET aliases for the local machine.
        let host = if config.host == "." || config.host.eq_ignore_ascii_case("(local)") {
            "127.0.0.1"
        } else {
            &config.host
        };

        // Step 1: Establish TCP connection
        let tcp_stream = if config.multi_subnet_failover {
            Self::connect_parallel(host, port, config.timeouts.connect_timeout).await?
        } else {
            let addr = format!("{host}:{port}");
            tracing::debug!("establishing TCP connection to {}", addr);
            let stream = timeout(config.timeouts.connect_timeout, TcpStream::connect(&addr))
                .await
                .map_err(|_| Error::ConnectTimeout {
                    host: config.host.clone(),
                    port: config.port,
                })?
                .map_err(Error::from)?;
            stream.set_nodelay(true).map_err(Error::from)?;
            stream
        };

        #[cfg(feature = "tls")]
        {
            // Determine TLS negotiation mode
            let tls_mode = TlsNegotiationMode::from_encrypt_mode(config.strict_mode);

            // Step 2: Handle TDS 8.0 strict mode (TLS before any TDS traffic)
            if tls_mode.is_tls_first() {
                return Self::connect_tds_8(config, tcp_stream).await;
            }

            // Step 3: TDS 7.x flow - PreLogin first, then TLS, then Login7
            Self::connect_tds_7x(config, tcp_stream).await
        }

        #[cfg(not(feature = "tls"))]
        {
            // When TLS feature is disabled, only no_tls connections are supported
            if config.strict_mode {
                return Err(Error::Config(
                    "TDS 8.0 strict mode requires TLS. Enable the 'tls' feature or use Encrypt=no_tls".into()
                ));
            }

            if !config.no_tls {
                return Err(Error::Config(
                    "TLS encryption requires the 'tls' feature. Either enable the 'tls' feature \
                     or use Encrypt=no_tls in your connection string for unencrypted connections."
                        .into(),
                ));
            }

            // Proceed with no-TLS connection
            Self::connect_no_tls(config, tcp_stream).await
        }
    }

    /// Resolve hostname to all IPs and race parallel TCP connections.
    ///
    /// Used when `MultiSubnetFailover=True` for AlwaysOn AG listeners that
    /// span multiple subnets. First successful TCP connection wins.
    async fn connect_parallel(
        host: &str,
        port: u16,
        connect_timeout: std::time::Duration,
    ) -> Result<TcpStream> {
        let addr_str = format!("{host}:{port}");
        let addrs: Vec<SocketAddr> = tokio::net::lookup_host(&addr_str)
            .await
            .map_err(Error::from)?
            .collect();

        if addrs.is_empty() {
            return Err(Error::from(std::io::Error::new(
                std::io::ErrorKind::AddrNotAvailable,
                format!("no addresses resolved for {host}:{port}"),
            )));
        }

        // Single address — no need to spawn tasks
        if addrs.len() == 1 {
            tracing::debug!(addr = %addrs[0], "MultiSubnetFailover: single address resolved");
            let stream = timeout(connect_timeout, TcpStream::connect(addrs[0]))
                .await
                .map_err(|_| Error::ConnectTimeout {
                    host: host.to_string(),
                    port,
                })?
                .map_err(Error::from)?;
            stream.set_nodelay(true).map_err(Error::from)?;
            return Ok(stream);
        }

        let addr_count = addrs.len();
        tracing::debug!(
            host = host,
            port = port,
            resolved_count = addr_count,
            "MultiSubnetFailover: racing parallel connections",
        );

        let mut join_set = tokio::task::JoinSet::new();

        for addr in addrs {
            let dur = connect_timeout;
            join_set.spawn(async move {
                let tcp = timeout(dur, TcpStream::connect(addr)).await.map_err(|_| {
                    std::io::Error::new(
                        std::io::ErrorKind::TimedOut,
                        format!("connection to {addr} timed out"),
                    )
                })??;
                tcp.set_nodelay(true)?;
                Ok::<(TcpStream, SocketAddr), std::io::Error>((tcp, addr))
            });
        }

        let mut last_error: Option<std::io::Error> = None;

        while let Some(result) = join_set.join_next().await {
            match result {
                Ok(Ok((stream, addr))) => {
                    tracing::debug!(addr = %addr, "MultiSubnetFailover: connected");
                    join_set.abort_all();
                    return Ok(stream);
                }
                Ok(Err(e)) => {
                    tracing::debug!(error = %e, "MultiSubnetFailover: attempt failed");
                    last_error = Some(e);
                }
                Err(join_err) => {
                    tracing::debug!(error = %join_err, "MultiSubnetFailover: task failed");
                    last_error = Some(std::io::Error::other(join_err.to_string()));
                }
            }
        }

        // All connections failed
        Err(Error::from(last_error.unwrap_or_else(|| {
            std::io::Error::new(
                std::io::ErrorKind::ConnectionRefused,
                format!("all {addr_count} parallel connection attempts failed for {host}:{port}"),
            )
        })))
    }

    /// Connect using TDS 8.0 strict mode.
    ///
    /// Flow: TCP -> TLS -> PreLogin (encrypted) -> Login7 (encrypted)
    #[cfg(feature = "tls")]
    async fn connect_tds_8(config: &Config, tcp_stream: TcpStream) -> Result<Client<Ready>> {
        tracing::debug!("using TDS 8.0 strict mode (TLS first)");

        // Build TLS configuration from the user's `config.tls` plus the
        // TDS 8.0 strict-mode requirements (see `connection_tls_config`).
        let tls_config = connection_tls_config(config, true);

        let tls_connector = TlsConnector::new(tls_config)?;

        // Perform TLS handshake before any TDS traffic
        let tls_stream = timeout(
            config.timeouts.tls_timeout,
            tls_connector.connect(tcp_stream, &config.host),
        )
        .await
        .map_err(|_| Error::TlsTimeout {
            host: config.host.clone(),
            port: config.port,
        })??;

        tracing::debug!("TLS handshake completed (strict mode)");

        // Create connection wrapper
        let mut connection = Connection::new(tls_stream);
        connection.set_max_message_size(config.max_response_size);

        // Send PreLogin (encrypted in strict mode)
        let prelogin = Self::build_prelogin(config, EncryptionLevel::Required);
        Self::send_prelogin(&mut connection, &prelogin).await?;
        let _prelogin_response = Self::receive_prelogin(&mut connection).await?;

        // Create SSPI negotiator if integrated auth
        #[cfg(any(feature = "integrated-auth", feature = "sspi-auth"))]
        let negotiator = Self::create_negotiator(config)?;
        #[cfg(any(feature = "integrated-auth", feature = "sspi-auth"))]
        let sspi_token = match negotiator {
            Some(ref neg) => Some(neg.initialize()?),
            None => None,
        };
        #[cfg(not(any(feature = "integrated-auth", feature = "sspi-auth")))]
        let sspi_token: Option<Vec<u8>> = None;

        // Send Login7
        let login = Self::build_login7(config, sspi_token);
        Self::send_login7(&mut connection, &login).await?;

        // Process login response (with timeout to prevent hangs during redirect)
        let (server_version, current_database, routing, server_collation) = timeout(
            config.timeouts.login_timeout,
            Self::process_login_response(
                &mut connection,
                #[cfg(any(feature = "integrated-auth", feature = "sspi-auth"))]
                negotiator.as_deref(),
            ),
        )
        .await
        .map_err(|_| Error::LoginTimeout {
            host: config.host.clone(),
            port: config.port,
        })??;

        // Handle routing redirect
        if let Some((host, port)) = routing {
            return Err(Error::Routing { host, port });
        }

        Ok(Client {
            config: config.clone(),
            _state: PhantomData,
            connection: Some(ConnectionHandle::Tls(connection)),
            server_version,
            current_database: current_database.clone(),
            server_collation,
            statement_cache: StatementCache::with_default_size(),
            transaction_descriptor: 0, // Auto-commit mode initially
            needs_reset: false,        // Fresh connection, no reset needed
            in_flight: false,          // No request pending
            #[cfg(feature = "otel")]
            instrumentation: InstrumentationContext::new(config.host.clone(), config.port)
                .with_database(current_database.clone().unwrap_or_default()),
            #[cfg(feature = "always-encrypted")]
            encryption_context: config.column_encryption.clone().map(|cfg| {
                std::sync::Arc::new(crate::encryption::EncryptionContext::from_arc(cfg))
            }),
        })
    }

    /// Connect using TDS 7.x flow.
    ///
    /// Flow: TCP -> PreLogin (clear) -> TLS -> Login7 (encrypted)
    ///
    /// Note: For TDS 7.x, the PreLogin exchange happens over raw TCP before
    /// upgrading to TLS. We use low-level I/O for this initial exchange
    /// since the Connection struct splits the stream immediately.
    #[cfg(feature = "tls")]
    async fn connect_tds_7x(config: &Config, mut tcp_stream: TcpStream) -> Result<Client<Ready>> {
        use bytes::BufMut;
        use tds_protocol::packet::{PACKET_HEADER_SIZE, PacketHeader, PacketStatus};
        use tokio::io::{AsyncReadExt, AsyncWriteExt};

        tracing::debug!("using TDS 7.x flow (PreLogin first)");

        // Build PreLogin packet
        // Determine client encryption level based on configuration
        let client_encryption = if config.no_tls {
            // no_tls: Completely disable TLS
            tracing::warn!(
                "⚠️  no_tls mode enabled. Connection will be UNENCRYPTED. \
                 Credentials and data will be transmitted in plaintext. \
                 This should only be used for development/testing with legacy SQL Server."
            );
            EncryptionLevel::NotSupported
        } else if config.encrypt {
            EncryptionLevel::On
        } else {
            EncryptionLevel::Off
        };
        let prelogin = Self::build_prelogin(config, client_encryption);
        tracing::debug!(encryption = ?client_encryption, "sending PreLogin");
        let prelogin_bytes = prelogin.encode();

        // Manually create and send the PreLogin packet over raw TCP
        let header = PacketHeader::new(
            PacketType::PreLogin,
            PacketStatus::END_OF_MESSAGE,
            (PACKET_HEADER_SIZE + prelogin_bytes.len()) as u16,
        );

        let mut packet_buf = BytesMut::with_capacity(PACKET_HEADER_SIZE + prelogin_bytes.len());
        header.encode(&mut packet_buf);
        packet_buf.put_slice(&prelogin_bytes);

        tcp_stream
            .write_all(&packet_buf)
            .await
            .map_err(Error::from)?;

        // Read PreLogin response
        let mut header_buf = [0u8; PACKET_HEADER_SIZE];
        tcp_stream
            .read_exact(&mut header_buf)
            .await
            .map_err(Error::from)?;

        let response_length = u16::from_be_bytes([header_buf[2], header_buf[3]]) as usize;
        let payload_length = response_length.saturating_sub(PACKET_HEADER_SIZE);

        let mut response_buf = vec![0u8; payload_length];
        tcp_stream
            .read_exact(&mut response_buf)
            .await
            .map_err(Error::from)?;

        let prelogin_response = PreLogin::decode(&response_buf[..])?;

        // Log PreLogin response
        // Note: The server sends its SQL Server product version in PreLogin,
        // NOT the TDS protocol version. The actual TDS version is negotiated
        // in the LOGINACK token after login.
        let client_tds_version = config.tds_version;
        if let Some(ref server_version) = prelogin_response.server_version {
            tracing::debug!(
                requested_tds_version = %client_tds_version,
                server_product_version = %server_version,
                server_product = server_version.product_name(),
                max_tds_version = %server_version.max_tds_version(),
                "PreLogin response received"
            );

            // Warn if the server's max TDS version is lower than requested
            let server_max_tds = server_version.max_tds_version();
            if server_max_tds < client_tds_version && !client_tds_version.is_tds_8() {
                tracing::warn!(
                    requested_tds_version = %client_tds_version,
                    server_max_tds_version = %server_max_tds,
                    server_product = server_version.product_name(),
                    "Server supports lower TDS version than requested. \
                     Connection will use server's maximum: {}",
                    server_max_tds
                );
            }

            // Warn about legacy SQL Server versions (2005 and earlier)
            if server_max_tds.is_legacy() {
                tracing::warn!(
                    server_product = server_version.product_name(),
                    server_max_tds_version = %server_max_tds,
                    "Server uses legacy TDS version. Some features may not be available."
                );
            }
        } else {
            tracing::debug!(
                requested_tds_version = %client_tds_version,
                "PreLogin response received (no version info)"
            );
        }

        // Check server encryption response
        let server_encryption = prelogin_response.encryption;
        tracing::debug!(encryption = ?server_encryption, "server encryption level");

        // Determine negotiated encryption level (follows TDS 7.x rules)
        // - NotSupported + NotSupported = NotSupported (no TLS at all)
        // - Off + Off = Off (TLS for login only, then plain)
        // - On + anything supported = On (full TLS)
        // - Required = On with failure if not possible
        let negotiated_encryption = match (client_encryption, server_encryption) {
            (EncryptionLevel::NotSupported, EncryptionLevel::NotSupported) => {
                EncryptionLevel::NotSupported
            }
            (EncryptionLevel::Off, EncryptionLevel::Off) => EncryptionLevel::Off,
            (EncryptionLevel::On, EncryptionLevel::Off)
            | (EncryptionLevel::On, EncryptionLevel::NotSupported) => {
                return Err(Error::Protocol(
                    "Server does not support requested encryption level".to_string(),
                ));
            }
            _ => EncryptionLevel::On,
        };

        // TLS is required unless negotiated encryption is NotSupported
        // Even with "Off", TLS is used to protect login credentials (per TDS 7.x spec)
        let use_tls = negotiated_encryption != EncryptionLevel::NotSupported;

        if use_tls {
            // Upgrade to TLS with PreLogin wrapping (TDS 7.x style).
            // In TDS 7.x, the TLS handshake is wrapped inside TDS PreLogin
            // packets. Honor the user's `config.tls` (custom root certs,
            // client auth) without the TDS 8.0 strict ALPN.
            let tls_config = connection_tls_config(config, false);

            let tls_connector = TlsConnector::new(tls_config)?;

            // Use PreLogin-wrapped TLS connection for TDS 7.x
            let mut tls_stream = timeout(
                config.timeouts.tls_timeout,
                tls_connector.connect_with_prelogin(tcp_stream, &config.host),
            )
            .await
            .map_err(|_| Error::TlsTimeout {
                host: config.host.clone(),
                port: config.port,
            })??;

            tracing::debug!("TLS handshake completed (PreLogin wrapped)");

            // Check if we need full encryption or login-only encryption
            let login_only_encryption = negotiated_encryption == EncryptionLevel::Off;

            if login_only_encryption {
                // Login-Only Encryption (ENCRYPT_OFF + ENCRYPT_OFF per MS-TDS spec):
                // - Login7 is sent through TLS to protect credentials
                // - Server responds in PLAINTEXT after receiving Login7
                // - All subsequent communication is plaintext
                //
                // We must NOT use Connection with TLS stream because Connection splits
                // the stream and we need to extract the underlying TCP afterward.
                use tokio::io::AsyncWriteExt;

                // Create SSPI negotiator if integrated auth
                // Note: SSPI handshake over login-only encryption is limited —
                // the server response comes in plaintext, so multi-step SSPI
                // may not work. We include the initial token but don't loop.
                #[cfg(any(feature = "integrated-auth", feature = "sspi-auth"))]
                let negotiator = Self::create_negotiator(config)?;
                #[cfg(any(feature = "integrated-auth", feature = "sspi-auth"))]
                let sspi_token = match negotiator {
                    Some(ref neg) => Some(neg.initialize()?),
                    None => None,
                };
                #[cfg(not(any(feature = "integrated-auth", feature = "sspi-auth")))]
                let sspi_token: Option<Vec<u8>> = None;

                // Build and send Login7 directly through TLS
                let login = Self::build_login7(config, sspi_token);
                let login_payload = login.encode();

                // Create TDS packet manually for Login7
                let max_packet = MAX_PACKET_SIZE;
                let max_payload = max_packet - PACKET_HEADER_SIZE;
                let chunks: Vec<_> = login_payload.chunks(max_payload).collect();
                let total_chunks = chunks.len();

                for (i, chunk) in chunks.into_iter().enumerate() {
                    let is_last = i == total_chunks - 1;
                    let status = if is_last {
                        PacketStatus::END_OF_MESSAGE
                    } else {
                        PacketStatus::NORMAL
                    };

                    let header = PacketHeader::new(
                        PacketType::Tds7Login,
                        status,
                        (PACKET_HEADER_SIZE + chunk.len()) as u16,
                    );

                    let mut packet_buf = BytesMut::with_capacity(PACKET_HEADER_SIZE + chunk.len());
                    header.encode(&mut packet_buf);
                    packet_buf.put_slice(chunk);

                    tls_stream
                        .write_all(&packet_buf)
                        .await
                        .map_err(Error::from)?;
                }

                // Flush TLS to ensure all data is sent
                tls_stream.flush().await.map_err(Error::from)?;

                tracing::debug!("Login7 sent through TLS, switching to plaintext for response");

                // Extract the underlying TCP stream from the TLS layer
                // TlsStream::into_inner() returns (IO, ClientConnection)
                // where IO is our TlsPreloginWrapper<TcpStream>
                let (wrapper, _client_conn) = tls_stream.into_inner();
                let tcp_stream = wrapper.into_inner();

                // Create Connection from plain TCP for reading response
                let mut connection = Connection::new(tcp_stream);
                connection.set_max_message_size(config.max_response_size);

                // Process login response (comes in plaintext, with timeout)
                let (server_version, current_database, routing, server_collation) = timeout(
                    config.timeouts.login_timeout,
                    Self::process_login_response(
                        &mut connection,
                        #[cfg(any(feature = "integrated-auth", feature = "sspi-auth"))]
                        negotiator.as_deref(),
                    ),
                )
                .await
                .map_err(|_| Error::LoginTimeout {
                    host: config.host.clone(),
                    port: config.port,
                })??;

                // Handle routing redirect
                if let Some((host, port)) = routing {
                    return Err(Error::Routing { host, port });
                }

                // Store plain TCP connection for subsequent operations
                Ok(Client {
                    config: config.clone(),
                    _state: PhantomData,
                    connection: Some(ConnectionHandle::Plain(connection)),
                    server_version,
                    current_database: current_database.clone(),
                    server_collation,
                    statement_cache: StatementCache::with_default_size(),
                    transaction_descriptor: 0, // Auto-commit mode initially
                    needs_reset: false,        // Fresh connection, no reset needed
                    in_flight: false,          // No request pending
                    #[cfg(feature = "otel")]
                    instrumentation: InstrumentationContext::new(config.host.clone(), config.port)
                        .with_database(current_database.clone().unwrap_or_default()),
                    #[cfg(feature = "always-encrypted")]
                    encryption_context: config.column_encryption.clone().map(|cfg| {
                        std::sync::Arc::new(crate::encryption::EncryptionContext::from_arc(cfg))
                    }),
                })
            } else {
                // Full Encryption (ENCRYPT_ON per MS-TDS spec):
                // - All communication after TLS handshake goes through TLS
                let mut connection = Connection::new(tls_stream);
                connection.set_max_message_size(config.max_response_size);

                // Create SSPI negotiator if integrated auth
                #[cfg(any(feature = "integrated-auth", feature = "sspi-auth"))]
                let negotiator = Self::create_negotiator(config)?;
                #[cfg(any(feature = "integrated-auth", feature = "sspi-auth"))]
                let sspi_token = match negotiator {
                    Some(ref neg) => Some(neg.initialize()?),
                    None => None,
                };
                #[cfg(not(any(feature = "integrated-auth", feature = "sspi-auth")))]
                let sspi_token: Option<Vec<u8>> = None;

                // Send Login7
                let login = Self::build_login7(config, sspi_token);
                Self::send_login7(&mut connection, &login).await?;

                // Process login response (with timeout)
                let (server_version, current_database, routing, server_collation) = timeout(
                    config.timeouts.login_timeout,
                    Self::process_login_response(
                        &mut connection,
                        #[cfg(any(feature = "integrated-auth", feature = "sspi-auth"))]
                        negotiator.as_deref(),
                    ),
                )
                .await
                .map_err(|_| Error::LoginTimeout {
                    host: config.host.clone(),
                    port: config.port,
                })??;

                // Handle routing redirect
                if let Some((host, port)) = routing {
                    return Err(Error::Routing { host, port });
                }

                Ok(Client {
                    config: config.clone(),
                    _state: PhantomData,
                    connection: Some(ConnectionHandle::TlsPrelogin(connection)),
                    server_version,
                    current_database: current_database.clone(),
                    server_collation,
                    statement_cache: StatementCache::with_default_size(),
                    transaction_descriptor: 0, // Auto-commit mode initially
                    needs_reset: false,        // Fresh connection, no reset needed
                    in_flight: false,          // No request pending
                    #[cfg(feature = "otel")]
                    instrumentation: InstrumentationContext::new(config.host.clone(), config.port)
                        .with_database(current_database.clone().unwrap_or_default()),
                    #[cfg(feature = "always-encrypted")]
                    encryption_context: config.column_encryption.clone().map(|cfg| {
                        std::sync::Arc::new(crate::encryption::EncryptionContext::from_arc(cfg))
                    }),
                })
            }
        } else {
            // Server does not require encryption and client doesn't either
            tracing::warn!(
                "Connecting without TLS encryption. This is insecure and should only be \
                 used for development/testing on trusted networks."
            );

            let mut connection = Connection::new(tcp_stream);

            connection.set_max_message_size(config.max_response_size);

            // Create SSPI negotiator if integrated auth
            #[cfg(any(feature = "integrated-auth", feature = "sspi-auth"))]
            let negotiator = Self::create_negotiator(config)?;
            #[cfg(any(feature = "integrated-auth", feature = "sspi-auth"))]
            let sspi_token = match negotiator {
                Some(ref neg) => Some(neg.initialize()?),
                None => None,
            };
            #[cfg(not(any(feature = "integrated-auth", feature = "sspi-auth")))]
            let sspi_token: Option<Vec<u8>> = None;

            // Build and send Login7
            let login = Self::build_login7(config, sspi_token);
            Self::send_login7(&mut connection, &login).await?;

            // Process login response (with timeout)
            let (server_version, current_database, routing, server_collation) = timeout(
                config.timeouts.login_timeout,
                Self::process_login_response(
                    &mut connection,
                    #[cfg(any(feature = "integrated-auth", feature = "sspi-auth"))]
                    negotiator.as_deref(),
                ),
            )
            .await
            .map_err(|_| Error::LoginTimeout {
                host: config.host.clone(),
                port: config.port,
            })??;

            // Handle routing redirect
            if let Some((host, port)) = routing {
                return Err(Error::Routing { host, port });
            }

            Ok(Client {
                config: config.clone(),
                _state: PhantomData,
                connection: Some(ConnectionHandle::Plain(connection)),
                server_version,
                current_database: current_database.clone(),
                server_collation,
                statement_cache: StatementCache::with_default_size(),
                transaction_descriptor: 0, // Auto-commit mode initially
                needs_reset: false,        // Fresh connection, no reset needed
                in_flight: false,          // No request pending
                #[cfg(feature = "otel")]
                instrumentation: InstrumentationContext::new(config.host.clone(), config.port)
                    .with_database(current_database.clone().unwrap_or_default()),
                #[cfg(feature = "always-encrypted")]
                encryption_context: config.column_encryption.clone().map(|cfg| {
                    std::sync::Arc::new(crate::encryption::EncryptionContext::from_arc(cfg))
                }),
            })
        }
    }

    /// Connect without TLS encryption (no_tls mode).
    ///
    /// This method is used when the `tls` feature is disabled and only supports
    /// unencrypted connections via `Encrypt=no_tls`.
    ///
    /// # Security Warning
    ///
    /// This transmits all data including credentials in plaintext. Only use this
    /// for development, testing, or on trusted internal networks where TLS is not
    /// required.
    #[cfg(not(feature = "tls"))]
    async fn connect_no_tls(config: &Config, mut tcp_stream: TcpStream) -> Result<Client<Ready>> {
        use bytes::BufMut;
        use tds_protocol::packet::{PACKET_HEADER_SIZE, PacketHeader, PacketStatus};
        use tokio::io::{AsyncReadExt, AsyncWriteExt};

        tracing::warn!(
            "⚠️  Connecting without TLS (tls feature disabled). \
             Credentials and data will be transmitted in plaintext."
        );

        // Build PreLogin packet with NotSupported encryption
        let prelogin = Self::build_prelogin(config, EncryptionLevel::NotSupported);
        let prelogin_bytes = prelogin.encode();

        // Manually create and send the PreLogin packet over raw TCP
        let header = PacketHeader::new(
            PacketType::PreLogin,
            PacketStatus::END_OF_MESSAGE,
            (PACKET_HEADER_SIZE + prelogin_bytes.len()) as u16,
        );

        let mut packet_buf = BytesMut::with_capacity(PACKET_HEADER_SIZE + prelogin_bytes.len());
        header.encode(&mut packet_buf);
        packet_buf.put_slice(&prelogin_bytes);

        tcp_stream
            .write_all(&packet_buf)
            .await
            .map_err(Error::from)?;

        // Read PreLogin response
        let mut header_buf = [0u8; PACKET_HEADER_SIZE];
        tcp_stream
            .read_exact(&mut header_buf)
            .await
            .map_err(Error::from)?;

        let response_length = u16::from_be_bytes([header_buf[2], header_buf[3]]) as usize;
        let payload_length = response_length.saturating_sub(PACKET_HEADER_SIZE);

        let mut response_buf = vec![0u8; payload_length];
        tcp_stream
            .read_exact(&mut response_buf)
            .await
            .map_err(Error::from)?;

        let prelogin_response = PreLogin::decode(&response_buf[..])?;

        // Check server encryption response - must accept NotSupported
        let server_encryption = prelogin_response.encryption;
        if server_encryption != EncryptionLevel::NotSupported {
            return Err(Error::Config(format!(
                "Server requires encryption (level: {:?}) but TLS feature is disabled. \
                     Either enable the 'tls' feature or configure the server to allow unencrypted connections.",
                server_encryption
            )));
        }

        tracing::debug!("Server accepted unencrypted connection");

        let mut connection = Connection::new(tcp_stream);

        connection.set_max_message_size(config.max_response_size);

        // Create SSPI negotiator if integrated auth
        #[cfg(any(feature = "integrated-auth", feature = "sspi-auth"))]
        let negotiator = Self::create_negotiator(config)?;
        #[cfg(any(feature = "integrated-auth", feature = "sspi-auth"))]
        let sspi_token = match negotiator {
            Some(ref neg) => Some(neg.initialize()?),
            None => None,
        };
        #[cfg(not(any(feature = "integrated-auth", feature = "sspi-auth")))]
        let sspi_token: Option<Vec<u8>> = None;

        // Build and send Login7
        let login = Self::build_login7(config, sspi_token);
        Self::send_login7(&mut connection, &login).await?;

        // Process login response (with timeout)
        let (server_version, current_database, routing, server_collation) = timeout(
            config.timeouts.login_timeout,
            Self::process_login_response(
                &mut connection,
                #[cfg(any(feature = "integrated-auth", feature = "sspi-auth"))]
                negotiator.as_deref(),
            ),
        )
        .await
        .map_err(|_| Error::LoginTimeout {
            host: config.host.clone(),
            port: config.port,
        })??;

        // Handle routing redirect
        if let Some((host, port)) = routing {
            return Err(Error::Routing { host, port });
        }

        Ok(Client {
            config: config.clone(),
            _state: PhantomData,
            connection: Some(ConnectionHandle::Plain(connection)),
            server_version,
            current_database: current_database.clone(),
            server_collation,
            statement_cache: StatementCache::with_default_size(),
            transaction_descriptor: 0,
            needs_reset: false,
            in_flight: false,
            #[cfg(feature = "otel")]
            instrumentation: InstrumentationContext::new(config.host.clone(), config.port)
                .with_database(current_database.clone().unwrap_or_default()),
            #[cfg(feature = "always-encrypted")]
            encryption_context: config.column_encryption.clone().map(|cfg| {
                std::sync::Arc::new(crate::encryption::EncryptionContext::from_arc(cfg))
            }),
        })
    }

    /// Build a PreLogin packet.
    fn build_prelogin(config: &Config, encryption: EncryptionLevel) -> PreLogin {
        // Use the configured TDS version (strict_mode overrides to V8_0)
        let version = if config.strict_mode {
            tds_protocol::version::TdsVersion::V8_0
        } else {
            config.tds_version
        };

        let mut prelogin = PreLogin::new()
            .with_version(version)
            .with_encryption(encryption);

        if config.mars {
            prelogin = prelogin.with_mars(true);
        }

        if let Some(ref instance) = config.instance {
            prelogin = prelogin.with_instance(instance);
        }

        prelogin
    }

    /// Resolve the workstation ID for the LOGIN7 HostName field.
    ///
    /// Per MS-TDS, the LOGIN7 HostName field contains the client machine name
    /// (not the server name). Priority:
    /// 1. `Config::workstation_id` (explicit override)
    /// 2. Machine hostname from environment (`COMPUTERNAME` on Windows, `HOSTNAME` on Linux)
    /// 3. Empty string (fallback)
    fn resolve_workstation_id(config: &Config) -> String {
        if let Some(ref id) = config.workstation_id {
            return id.clone();
        }
        // COMPUTERNAME is set on Windows; HOSTNAME is set on most Linux systems.
        // This avoids adding a dependency for a simple lookup.
        std::env::var("COMPUTERNAME")
            .or_else(|_| std::env::var("HOSTNAME"))
            .unwrap_or_default()
    }

    /// Build a Login7 packet.
    ///
    /// When `sspi_token` is provided (integrated auth), the Login7 packet is
    /// built with the integrated security flag and the initial SSPI blob.
    fn build_login7(config: &Config, sspi_token: Option<Vec<u8>>) -> Login7 {
        // Use the configured TDS version (strict_mode overrides to V8_0)
        let version = if config.strict_mode {
            tds_protocol::version::TdsVersion::V8_0
        } else {
            config.tds_version
        };

        let mut login = Login7::new()
            .with_tds_version(version)
            .with_packet_size(config.packet_size as u32)
            .with_app_name(&config.application_name)
            .with_server_name(&config.host)
            .with_hostname(Self::resolve_workstation_id(config));

        if let Some(ref database) = config.database {
            login = login.with_database(database);
        }

        // ApplicationIntent → LOGIN7 TypeFlags READONLY_INTENT bit
        if config.application_intent == crate::config::ApplicationIntent::ReadOnly {
            login = login.with_read_only_intent(true);
        }

        // Session language → LOGIN7 Language field
        if let Some(ref lang) = config.language {
            login = login.with_language(lang);
        }

        // Set credentials
        if let Some(token) = sspi_token {
            // Integrated auth: set SSPI data and integrated security flag
            login = login.with_integrated_auth(token);
        } else if let mssql_auth::Credentials::SqlServer { username, password } =
            &config.credentials
        {
            login = login.with_sql_auth(username.as_ref(), password.as_ref());
        }

        // When Always Encrypted is configured, add the ColumnEncryption feature extension.
        // Version 1 = client supports column encryption without enclave computations.
        #[cfg(feature = "always-encrypted")]
        if config.column_encryption.is_some() {
            login = login.with_feature(tds_protocol::login7::FeatureExtension {
                feature_id: tds_protocol::login7::FeatureId::ColumnEncryption,
                data: bytes::Bytes::from_static(&[0x01]), // Version 1
            });
            tracing::debug!("Login7: adding ColumnEncryption feature extension (version 1)");
        }

        login
    }

    /// Create an SSPI/GSSAPI negotiator if integrated auth is configured.
    ///
    /// Returns `None` for non-integrated credential types.
    ///
    /// On Windows with `sspi-auth`, uses native Windows SSPI (`secur32.dll`) which
    /// supports all account types including Microsoft Accounts. Falls back to sspi-rs
    /// on non-Windows platforms.
    ///
    /// With `integrated-auth` (Linux/macOS), uses GSSAPI/Kerberos.
    #[cfg(any(feature = "integrated-auth", feature = "sspi-auth"))]
    fn create_negotiator(config: &Config) -> Result<Option<Box<dyn mssql_auth::SspiNegotiator>>> {
        #[allow(clippy::match_like_matches_macro)]
        let is_integrated = match &config.credentials {
            mssql_auth::Credentials::Integrated => true,
            _ => false,
        };

        if !is_integrated {
            return Ok(None);
        }

        // On Windows: prefer native SSPI (secur32.dll) for integrated auth.
        // This handles all Windows account types including Microsoft Accounts,
        // domain accounts, and local accounts — unlike sspi-rs which requires
        // explicit credentials.
        #[cfg(all(windows, feature = "sspi-auth"))]
        let negotiator: Box<dyn mssql_auth::SspiNegotiator> =
            Box::new(mssql_auth::NativeSspiAuth::new(&config.host, config.port)?);

        // On non-Windows: use sspi-rs (pure Rust SSPI implementation)
        #[cfg(all(not(windows), feature = "sspi-auth"))]
        let negotiator: Box<dyn mssql_auth::SspiNegotiator> =
            Box::new(mssql_auth::SspiAuth::new(&config.host, config.port)?);

        #[cfg(all(feature = "integrated-auth", not(feature = "sspi-auth")))]
        let negotiator: Box<dyn mssql_auth::SspiNegotiator> =
            Box::new(mssql_auth::IntegratedAuth::new(&config.host, config.port));

        Ok(Some(negotiator))
    }

    /// Send a PreLogin packet (for use with Connection).
    #[cfg(feature = "tls")]
    async fn send_prelogin<T>(connection: &mut Connection<T>, prelogin: &PreLogin) -> Result<()>
    where
        T: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin,
    {
        let payload = prelogin.encode();
        let max_packet = MAX_PACKET_SIZE;

        connection
            .send_message(PacketType::PreLogin, payload, max_packet)
            .await?;
        Ok(())
    }

    /// Receive a PreLogin response (for use with Connection).
    #[cfg(feature = "tls")]
    async fn receive_prelogin<T>(connection: &mut Connection<T>) -> Result<PreLogin>
    where
        T: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin,
    {
        let message = connection
            .read_message()
            .await?
            .ok_or(Error::ConnectionClosed)?;

        Ok(PreLogin::decode(&message.payload[..])?)
    }

    /// Send a Login7 packet.
    async fn send_login7<T>(connection: &mut Connection<T>, login: &Login7) -> Result<()>
    where
        T: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin,
    {
        let payload = login.encode();
        let max_packet = MAX_PACKET_SIZE;

        connection
            .send_message(PacketType::Tds7Login, payload, max_packet)
            .await?;
        Ok(())
    }

    /// Process the login response tokens, handling SSPI challenge/response if needed.
    ///
    /// When a `negotiator` is provided and the server sends an SSPI challenge token,
    /// this method will automatically perform the multi-step SSPI handshake by:
    /// 1. Calling `negotiator.step(challenge)` to generate a response
    /// 2. Sending the response via an SSPI packet
    /// 3. Reading the next server message and continuing
    ///
    /// Returns: (server_version, database, routing_info)
    #[allow(clippy::never_loop)] // Loop is used when integrated-auth/sspi-auth features are enabled
    async fn process_login_response<T>(
        connection: &mut Connection<T>,
        #[cfg(any(feature = "integrated-auth", feature = "sspi-auth"))] negotiator: Option<
            &dyn mssql_auth::SspiNegotiator,
        >,
    ) -> Result<(
        Option<u32>,
        Option<String>,
        Option<(String, u16)>,
        Option<tds_protocol::token::Collation>,
    )>
    where
        T: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin,
    {
        let mut server_version = None;
        let mut database = None;
        let mut routing = None;
        let mut collation = None;

        'outer: loop {
            let message = connection
                .read_message()
                .await?
                .ok_or(Error::ConnectionClosed)?;

            let response_bytes = message.payload;
            let mut parser = TokenParser::new(response_bytes);

            while let Some(token) = parser.next_token()? {
                match token {
                    Token::LoginAck(ack) => {
                        tracing::info!(
                            version = ack.tds_version,
                            interface = ack.interface,
                            prog_name = %ack.prog_name,
                            "login acknowledged"
                        );
                        server_version = Some(ack.tds_version);
                    }
                    Token::EnvChange(env) => {
                        Self::process_env_change(&env, &mut database, &mut routing, &mut collation);
                    }
                    #[cfg(any(feature = "integrated-auth", feature = "sspi-auth"))]
                    Token::Sspi(sspi_token) => {
                        let neg = negotiator.ok_or_else(|| {
                            Error::Protocol(
                                "server sent SSPI challenge but no negotiator is configured"
                                    .to_string(),
                            )
                        })?;

                        tracing::debug!(
                            challenge_len = sspi_token.data.len(),
                            "received SSPI challenge from server"
                        );

                        if let Some(response) = neg.step(&sspi_token.data)? {
                            tracing::debug!(response_len = response.len(), "sending SSPI response");
                            connection
                                .send_message(
                                    PacketType::Sspi,
                                    bytes::Bytes::from(response),
                                    tds_protocol::packet::MAX_PACKET_SIZE,
                                )
                                .await?;
                        }

                        // After sending the SSPI response, read the next server message
                        continue 'outer;
                    }
                    Token::Error(err) => {
                        return Err(Error::Server {
                            number: err.number,
                            state: err.state,
                            class: err.class,
                            message: err.message.clone(),
                            server: if err.server.is_empty() {
                                None
                            } else {
                                Some(err.server.clone())
                            },
                            procedure: if err.procedure.is_empty() {
                                None
                            } else {
                                Some(err.procedure.clone())
                            },
                            line: err.line as u32,
                        });
                    }
                    Token::Info(info) => {
                        tracing::info!(
                            number = info.number,
                            message = %info.message,
                            "server info message"
                        );
                    }
                    Token::Done(done) => {
                        if done.status.error {
                            return Err(Error::Protocol("login failed".to_string()));
                        }
                        break 'outer;
                    }
                    _ => {}
                }
            }

            // If we consumed all tokens without a Done or SSPI, break
            break;
        }

        Ok((server_version, database, routing, collation))
    }

    /// Process an EnvChange token.
    fn process_env_change(
        env: &EnvChange,
        database: &mut Option<String>,
        routing: &mut Option<(String, u16)>,
        collation: &mut Option<tds_protocol::token::Collation>,
    ) {
        use tds_protocol::token::EnvChangeValue;

        match env.env_type {
            EnvChangeType::Database => {
                if let EnvChangeValue::String(ref new_value) = env.new_value {
                    tracing::debug!(database = %new_value, "database changed");
                    *database = Some(new_value.clone());
                }
            }
            EnvChangeType::Routing => {
                if let EnvChangeValue::Routing { ref host, port } = env.new_value {
                    tracing::info!(host = %host, port = port, "routing redirect received");
                    *routing = Some((host.clone(), port));
                }
            }
            EnvChangeType::SqlCollation => {
                if let EnvChangeValue::Binary(ref data) = env.new_value {
                    if data.len() >= 5 {
                        let c = tds_protocol::token::Collation::from_bytes(
                            data[..5].try_into().unwrap(),
                        );
                        tracing::debug!(
                            lcid = c.lcid,
                            sort_id = c.sort_id,
                            "server collation received"
                        );
                        *collation = Some(c);
                    }
                }
            }
            _ => {
                if let EnvChangeValue::String(ref new_value) = env.new_value {
                    tracing::debug!(
                        env_type = ?env.env_type,
                        new_value = %new_value,
                        "environment change"
                    );
                }
            }
        }
    }
}

/// Build the TLS configuration for an outbound connection.
///
/// Starts from the user's [`Config::tls`] so custom root certificates, client
/// auth, and protocol-version bounds are honored, then layers the
/// connection-specific requirements. `trust_server_certificate` is taken from
/// the authoritative top-level [`Config`] field: both the builder and the
/// connection-string parser set it, but the parser does not mirror it into
/// `config.tls`, so reading it here is what keeps `TrustServerCertificate=...`
/// connection strings working.
///
/// `strict` selects TDS 8.0 strict mode (TLS-first) and adds the `tds/8.0`
/// ALPN protocol; TDS 7.x leaves both off (its TLS is wrapped in PreLogin).
///
/// Note the asymmetry: root certificates and client auth come from
/// `config.tls`, but `trust_server_certificate` is taken from the top-level
/// field and overrides whatever `config.tls` holds. So setting *only*
/// `config.tls = TlsConfig::new().trust_server_certificate(true)` while
/// leaving the top-level field at its `false` default does not trust the
/// server — set it via the connection string (`TrustServerCertificate=true`)
/// or `Config::trust_server_certificate(true)`, which is the supported path.
#[cfg(feature = "tls")]
fn connection_tls_config(config: &Config, strict: bool) -> TlsConfig {
    let tls = config
        .tls
        .clone()
        .trust_server_certificate(config.trust_server_certificate);
    if strict {
        tls.strict_mode(true)
            .with_alpn_protocols(vec![b"tds/8.0".to_vec()])
    } else {
        tls
    }
}

#[cfg(all(test, feature = "tls"))]
mod tls_config_tests {
    use super::*;
    use mssql_tls::CertificateDer;

    fn config_with_root(cert: Vec<u8>) -> Config {
        let mut config = Config::new();
        config.tls = config
            .tls
            .clone()
            .add_root_certificate(CertificateDer::from(cert));
        config
    }

    #[test]
    fn custom_root_certificate_reaches_connector_config() {
        // The bug: connect built a fresh TlsConfig and dropped config.tls,
        // so a custom CA was unreachable. Assert it survives into the
        // connection's TLS config, in both strict and non-strict paths.
        let config = config_with_root(vec![0xCA; 32]);

        for strict in [true, false] {
            let tls = connection_tls_config(&config, strict);
            assert_eq!(
                tls.root_certificates.len(),
                1,
                "custom root must reach the connector (strict={strict})"
            );
            assert_eq!(tls.root_certificates[0].as_ref(), &[0xCA; 32][..]);
        }
    }

    #[test]
    fn trust_server_certificate_taken_from_top_level_field() {
        // Mirrors the connection-string path, which sets the top-level field
        // without updating config.tls.
        let mut config = Config::new();
        config.trust_server_certificate = true;
        // config.tls still has the default (false) trust flag.
        assert!(!config.tls.trust_server_certificate);

        let tls = connection_tls_config(&config, false);
        assert!(
            tls.trust_server_certificate,
            "top-level trust flag must win"
        );
    }

    #[test]
    fn strict_mode_adds_tds8_alpn() {
        let config = Config::new();
        let strict = connection_tls_config(&config, true);
        assert!(strict.strict_mode);
        assert!(strict.alpn_protocols.iter().any(|p| p == b"tds/8.0"));

        let non_strict = connection_tls_config(&config, false);
        assert!(!non_strict.strict_mode);
    }
}