exarrow-rs 0.28.0

//! WebSocket transport implementation for Exasol.
//!
//! This module provides a WebSocket-based transport for communicating with
//! Exasol databases using the Exasol WebSocket protocol.

use std::collections::HashMap;

use async_trait::async_trait;
use aws_lc_rs::rand::{SecureRandom, SystemRandom};
use base64::engine::general_purpose::STANDARD;
use base64::Engine;
use futures_util::{SinkExt, StreamExt};
use num_bigint::BigUint;
use tokio::net::TcpStream;
use tokio_tungstenite::{
    tungstenite::{protocol::WebSocketConfig, Message},
    MaybeTlsStream, WebSocketStream,
};

use crate::error::TransportError;
use crate::DRIVER_NAME;

use super::deadline::SetupStep;
use super::lifecycle::{self, ConnectionLifecycle, ConnectionState, LifecycleSteps};
use super::messages::{
    AttributedResponse, AuthRequest, ClosePreparedStatementRequest, ClosePreparedStatementResponse,
    CloseResultSetRequest, CloseResultSetResponse, CreatePreparedStatementRequest,
    CreatePreparedStatementResponse, DisconnectRequest, DisconnectResponse,
    ExecutePreparedStatementRequest, ExecuteRequest, ExecuteResponse, FetchRequest, FetchResponse,
    GetAttributesRequest, GetAttributesResponse, LoginInitRequest, LoginResponse,
    PublicKeyResponse, ResponseAttributes, ResultData, ResultEntryKind, ResultPayload,
    ResultSetHandle, SessionInfo, SetAttributesRequest, SetAttributesResponse,
};
use super::protocol::{
    ConnectionParams, Credentials, PreparedStatementHandle, QueryResult, TransportProtocol,
};
use super::tls;

/// WebSocket transport implementation.
///
/// This transport uses Exasol's WebSocket API for communication.
pub struct WebSocketTransport {
    /// WebSocket connection (None if not connected)
    ws_stream: Option<WebSocketStream<MaybeTlsStream<TcpStream>>>,
    /// Current session information (None if not authenticated)
    session_info: Option<SessionInfo>,
    lifecycle: ConnectionLifecycle,
    /// Next fetch start position per open result set handle
    fetch_positions: HashMap<i32, i64>,
    current_schema: Option<String>,
}

impl WebSocketTransport {
    /// Create a new WebSocket transport.
    pub fn new() -> Self {
        Self {
            ws_stream: None,
            session_info: None,
            lifecycle: ConnectionLifecycle::new(),
            fetch_positions: HashMap::new(),
            current_schema: None,
        }
    }

    /// Send a message and receive a response.
    ///
    /// Exasol may send intermediate status messages (like "EXECUTING") before
    /// the actual JSON response. This method skips those messages and returns
    /// only the final JSON response. Every response passes through here, so this
    /// also records the current schema when the response reports it.
    async fn send_receive<T, R>(&mut self, request: &T) -> Result<R, TransportError>
    where
        T: serde::Serialize,
        R: serde::de::DeserializeOwned + AttributedResponse,
    {
        // Serialize request
        let request_json = serde_json::to_string(request)?;

        // Get WebSocket stream
        let ws_stream = self
            .ws_stream
            .as_mut()
            .ok_or_else(|| TransportError::ProtocolError("Not connected".to_string()))?;

        // Send request
        ws_stream
            .send(Message::Text(request_json.into()))
            .await
            .map_err(|e| TransportError::SendError(e.to_string()))?;

        // Receive response, skipping intermediate status messages
        loop {
            let response_msg = ws_stream
                .next()
                .await
                .ok_or_else(|| TransportError::ReceiveError("Connection closed".to_string()))?
                .map_err(|e| TransportError::ReceiveError(e.to_string()))?;

            // Parse response
            let response_text = response_msg.to_text().map_err(|e| {
                TransportError::ProtocolError(format!("Invalid message format: {}", e))
            })?;

            // Skip intermediate status messages (e.g., "EXECUTING", "FETCHING")
            // These are plain text, not JSON objects starting with '{'
            if !response_text.starts_with('{') {
                continue;
            }

            let response: R = serde_json::from_str(response_text).map_err(|e| {
                // Include context in error for debugging
                let preview: String = response_text.chars().take(200).collect();
                TransportError::InvalidResponse(format!(
                    "JSON parse error: {}. Response preview: '{}'",
                    e, preview
                ))
            })?;

            self.record_current_schema(response.attributes());
            return Ok(response);
        }
    }

    /// Record a reported current schema, where an empty value means none; a
    /// response that does not report it leaves the recorded value unchanged.
    fn record_current_schema(&mut self, attributes: Option<&ResponseAttributes>) {
        if let Some(schema) = attributes.and_then(|a| a.current_schema.as_deref()) {
            self.current_schema = Some(schema.to_owned()).filter(|name| !name.is_empty());
        }
    }

    /// Send `setAttributes` and fail with the server's message when it rejects a value.
    async fn set_attributes(
        &mut self,
        request: &SetAttributesRequest,
    ) -> Result<(), TransportError> {
        self.lifecycle.require(
            ConnectionState::Authenticated,
            "Must authenticate before setting attributes",
        )?;

        let response: SetAttributesResponse = self.send_receive(request).await?;
        self.check_status(&response.status, &response.exception)
    }

    /// Check response status and return error if not ok.
    fn check_status(
        &self,
        status: &str,
        exception: &Option<super::messages::ExceptionInfo>,
    ) -> Result<(), TransportError> {
        if status != "ok" {
            let error_msg = exception
                .as_ref()
                .map(|e| {
                    format!(
                        "{} (SQL code: {})",
                        e.text,
                        e.sql_code.as_deref().unwrap_or("unknown")
                    )
                })
                .unwrap_or_else(|| "Unknown error".to_string());
            return Err(TransportError::ProtocolError(error_msg));
        }
        Ok(())
    }

    /// Turn an `execute` / `executePreparedStatement` response into a
    /// [`QueryResult`].
    ///
    /// Both statement paths receive the same `ExecuteResponse` shape and owe
    /// callers the same `QueryResult`, so the status check, the first-result
    /// selection, and the result-type dispatch live here once. A result set
    /// handle starts its fetch position after the rows this response delivered.
    fn query_result_from_response(
        &mut self,
        response: ExecuteResponse,
    ) -> Result<QueryResult, TransportError> {
        self.check_status(&response.status, &response.exception)?;

        let response_data = response
            .response_data
            .ok_or_else(|| TransportError::InvalidResponse("Missing response data".to_string()))?;

        if response_data.results.is_empty() {
            return Err(TransportError::InvalidResponse(
                "No results returned".to_string(),
            ));
        }

        // Process first result (multi-result statements not fully supported yet)
        let result = &response_data.results[0];

        match result.kind() {
            ResultEntryKind::ResultSet => {
                let result_set = result.result_set.as_ref().ok_or_else(|| {
                    TransportError::InvalidResponse(format!(
                        "Missing result_set data. Result: {:?}",
                        result
                    ))
                })?;

                let columns = result_set.columns.clone().ok_or_else(|| {
                    TransportError::InvalidResponse("Missing columns".to_string())
                })?;

                // Data is in column-major format from Exasol
                let data_values = result_set.data.clone().unwrap_or_default();
                let total_rows = result_set.num_rows.unwrap_or(0);

                let data = ResultData {
                    columns,
                    data: ResultPayload::Json(data_values),
                    total_rows,
                };

                // Handle may be None when all data fits in one response
                let handle = result_set.result_set_handle.map(ResultSetHandle::new);
                if let Some(handle) = handle {
                    self.fetch_positions
                        .insert(handle.as_i32(), data.data.num_rows() as i64);
                }

                Ok(QueryResult::result_set(handle, data))
            }
            ResultEntryKind::RowCount => {
                let count = result.row_count.unwrap_or(0);
                Ok(QueryResult::row_count(count))
            }
            ResultEntryKind::Unknown => Err(TransportError::InvalidResponse(format!(
                "Unknown result type: {}",
                result.result_type
            ))),
        }
    }

    /// Encrypt password using RSA with PKCS#1 v1.5 padding.
    ///
    /// The password is encrypted using the server's public key and then
    /// base64-encoded for transmission. Supports RSA key sizes from 1024
    /// to 8192 bits (Exasol servers may emit 1024-bit keys during login).
    fn encrypt_password(password: &str, public_key_pem: &str) -> Result<String, TransportError> {
        let pkcs1_der = Self::pem_to_pkcs1_der(public_key_pem).map_err(|e| {
            TransportError::ProtocolError(format!("Failed to parse RSA public key PEM: {}", e))
        })?;

        let (n, e) = Self::parse_rsa_public_key(&pkcs1_der).map_err(|err| {
            TransportError::ProtocolError(format!("Failed to parse RSA public key: {}", err))
        })?;

        let bit_len = n.bits() as usize;
        if !(1024..=8192).contains(&bit_len) {
            return Err(TransportError::ProtocolError(format!(
                "Unsupported RSA key size: {} bits (expected 1024..=8192)",
                bit_len
            )));
        }

        let k = bit_len.div_ceil(8);

        let em = Self::pkcs1_v15_pad(password.as_bytes(), k).map_err(|err| {
            TransportError::ProtocolError(format!("Failed to pad password: {}", err))
        })?;

        let m = BigUint::from_bytes_be(&em);
        let c = m.modpow(&e, &n);

        let mut ciphertext = vec![0u8; k];
        let c_bytes = c.to_bytes_be();
        ciphertext[k - c_bytes.len()..].copy_from_slice(&c_bytes);

        Ok(STANDARD.encode(&ciphertext))
    }

    /// Strip PEM headers and base64-decode to obtain the raw PKCS#1 DER bytes.
    ///
    /// Exasol sends RSA public keys in PKCS#1 PEM format (`BEGIN RSA PUBLIC KEY`).
    /// Returns the raw PKCS#1 RSAPublicKey DER (SEQUENCE { INTEGER n, INTEGER e }).
    fn pem_to_pkcs1_der(pem: &str) -> Result<Vec<u8>, &'static str> {
        let start_marker = "-----BEGIN RSA PUBLIC KEY-----";
        let end_marker = "-----END RSA PUBLIC KEY-----";

        let start = pem.find(start_marker).ok_or("Missing PEM start marker")? + start_marker.len();
        let end = pem.find(end_marker).ok_or("Missing PEM end marker")?;

        let base64_content: String = pem[start..end]
            .chars()
            .filter(|c| !c.is_whitespace())
            .collect();

        STANDARD
            .decode(&base64_content)
            .map_err(|_| "Invalid base64 in PEM")
    }

    /// Parse a PKCS#1 RSAPublicKey DER structure into (n, e).
    ///
    /// Reads `SEQUENCE { INTEGER n, INTEGER e }` from the raw DER bytes.
    fn parse_rsa_public_key(pkcs1_der: &[u8]) -> Result<(BigUint, BigUint), &'static str> {
        let mut pos = 0;

        // Outer SEQUENCE
        if pkcs1_der.get(pos) != Some(&0x30) {
            return Err("Invalid DER: expected SEQUENCE");
        }
        pos += 1;

        let (seq_len, len_bytes) = Self::read_der_length(&pkcs1_der[pos..])?;
        pos += len_bytes;

        let seq_end = pos + seq_len;
        if seq_end > pkcs1_der.len() {
            return Err("Invalid DER: truncated");
        }

        // First INTEGER: n (modulus)
        let (n_bytes, consumed) = Self::read_der_integer(&pkcs1_der[pos..seq_end])?;
        pos += consumed;

        // Second INTEGER: e (public exponent)
        let (e_bytes, consumed) = Self::read_der_integer(&pkcs1_der[pos..seq_end])?;
        pos += consumed;

        if pos != seq_end {
            return Err("Invalid DER: trailing bytes after SEQUENCE");
        }

        Ok((
            BigUint::from_bytes_be(n_bytes),
            BigUint::from_bytes_be(e_bytes),
        ))
    }

    /// Read a DER length field, returning (length_value, bytes_consumed).
    fn read_der_length(data: &[u8]) -> Result<(usize, usize), &'static str> {
        if data.is_empty() {
            return Err("Invalid DER: truncated");
        }

        let first = data[0];
        if first < 0x80 {
            Ok((first as usize, 1))
        } else {
            let num_bytes = (first & 0x7F) as usize;
            if num_bytes == 0 || num_bytes > 3 {
                return Err("Invalid DER: truncated");
            }
            if data.len() < 1 + num_bytes {
                return Err("Invalid DER: truncated");
            }
            let mut len: usize = 0;
            for &b in &data[1..1 + num_bytes] {
                len = (len << 8) | (b as usize);
            }
            Ok((len, 1 + num_bytes))
        }
    }

    /// Read a DER INTEGER, returning (value_bytes_without_sign_padding, total_bytes_consumed).
    fn read_der_integer(data: &[u8]) -> Result<(&[u8], usize), &'static str> {
        if data.is_empty() || data[0] != 0x02 {
            return Err("Invalid DER: expected INTEGER");
        }

        let (int_len, len_bytes) = Self::read_der_length(&data[1..])?;
        let header_len = 1 + len_bytes;

        if data.len() < header_len + int_len {
            return Err("Invalid DER: truncated");
        }

        let mut value = &data[header_len..header_len + int_len];

        // Strip leading 0x00 sign-bit padding
        if value.len() > 1 && value[0] == 0x00 {
            value = &value[1..];
        }

        Ok((value, header_len + int_len))
    }

    /// Build a PKCS#1 v1.5 encryption block: 0x00 || 0x02 || PS || 0x00 || message.
    ///
    /// PS consists of random non-zero bytes from `SystemRandom`.
    fn pkcs1_v15_pad(message: &[u8], k: usize) -> Result<Vec<u8>, &'static str> {
        if message.len() > k.saturating_sub(11) {
            return Err("Message too long for RSA modulus");
        }

        let ps_len = k - message.len() - 3;
        let rng = SystemRandom::new();

        let mut em = Vec::with_capacity(k);
        em.push(0x00);
        em.push(0x02);

        // Fill PS with non-zero random bytes
        for _ in 0..ps_len {
            let mut byte = [0u8; 1];
            loop {
                rng.fill(&mut byte)
                    .map_err(|_| "Failed to generate random bytes")?;
                if byte[0] != 0 {
                    break;
                }
            }
            em.push(byte[0]);
        }

        em.push(0x00);
        em.extend_from_slice(message);

        Ok(em)
    }
}

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

fn auth_request(credentials: &Credentials, encrypted_password: String) -> AuthRequest {
    AuthRequest::new(
        credentials.username.clone(),
        encrypted_password,
        credentials.client_name.clone(),
    )
    .with_driver_name(DRIVER_NAME.to_string())
    .with_client_version(credentials.client_version.clone())
}

#[async_trait]
impl LifecycleSteps for WebSocketTransport {
    fn lifecycle_mut(&mut self) -> &mut ConnectionLifecycle {
        &mut self.lifecycle
    }

    async fn login_exchange(
        &mut self,
        credentials: &Credentials,
    ) -> Result<SessionInfo, TransportError> {
        // Step 1: Send login init request to get the server's public key
        let init_request = LoginInitRequest::new();
        let key_response: PublicKeyResponse = self.send_receive(&init_request).await?;

        // Check response status
        self.check_status(&key_response.status, &key_response.exception)?;

        // Extract public key data
        let key_data = key_response.response_data.ok_or_else(|| {
            TransportError::InvalidResponse("Missing public key data in response".to_string())
        })?;

        // Step 2: Encrypt password using the server's RSA public key
        let encrypted_password =
            Self::encrypt_password(&credentials.password, &key_data.public_key_pem)?;

        // Step 3: Send authentication request with encrypted password
        let auth_request = auth_request(credentials, encrypted_password);
        let login_response: LoginResponse = self.send_receive(&auth_request).await?;

        // Check response status
        self.check_status(&login_response.status, &login_response.exception)?;

        // Step 4: Extract session info from the final response
        let session_data = login_response
            .response_data
            .ok_or_else(|| TransportError::InvalidResponse("Missing response data".to_string()))?;

        let session_info: SessionInfo = session_data.into();

        self.session_info = Some(session_info.clone());
        Ok(session_info)
    }

    async fn send_disconnect(&mut self) {
        // The server may close the connection before it answers, so errors are ignored.
        if self.lifecycle.state() == ConnectionState::Authenticated {
            let request = DisconnectRequest::new();
            let _ = self.send_receive::<_, DisconnectResponse>(&request).await;
        }

        if let Some(mut ws_stream) = self.ws_stream.take() {
            use tokio_tungstenite::tungstenite::protocol::frame::coding::CloseCode;
            use tokio_tungstenite::tungstenite::protocol::CloseFrame;

            let close_frame = CloseFrame {
                code: CloseCode::Normal,
                reason: "Client closing connection".into(),
            };
            let _ = ws_stream.close(Some(close_frame)).await;

            // Drain any remaining messages to complete the close handshake
            while let Ok(Some(_)) =
                tokio::time::timeout(std::time::Duration::from_millis(100), ws_stream.next()).await
            {
            }
        }
    }

    fn release_connection(&mut self) {
        self.ws_stream = None;
        self.session_info = None;
        self.current_schema = None;
    }
}

#[async_trait]
impl TransportProtocol for WebSocketTransport {
    async fn connect(&mut self, params: &ConnectionParams) -> Result<(), TransportError> {
        let deadline = self.lifecycle.begin_connect(params.timeout_ms)?;

        let tcp_stream = deadline
            .run(SetupStep::TcpConnect, async {
                TcpStream::connect((params.host.as_str(), params.port))
                    .await
                    .map_err(|e| TransportError::WebSocketError(e.to_string()))
            })
            .await?;

        let stream = if params.use_tls {
            let config = tls::client_config(
                params.certificate_fingerprint.as_deref(),
                params.validate_server_certificate,
            );
            let tls_stream = deadline
                .run(
                    SetupStep::TlsHandshake,
                    tls::client_handshake(tcp_stream, &params.host, config),
                )
                .await?;
            MaybeTlsStream::Rustls(tls_stream)
        } else {
            MaybeTlsStream::Plain(tcp_stream)
        };

        // Exasol sends result set frames that can exceed the tungstenite default of 16 MiB.
        let mut ws_config = WebSocketConfig::default();
        ws_config.max_frame_size = None;
        ws_config.max_message_size = None;
        let (ws_stream, _) = deadline
            .run(SetupStep::WebSocketUpgrade, async {
                tokio_tungstenite::client_async_with_config(
                    params.to_websocket_url(),
                    stream,
                    Some(ws_config),
                )
                .await
                .map_err(|e| TransportError::WebSocketError(e.to_string()))
            })
            .await?;

        self.ws_stream = Some(ws_stream);
        self.lifecycle.connected(deadline);
        Ok(())
    }

    async fn authenticate(
        &mut self,
        credentials: &Credentials,
    ) -> Result<SessionInfo, TransportError> {
        lifecycle::authenticate_within_deadline(self, credentials).await
    }

    async fn execute_query(&mut self, sql: &str) -> Result<QueryResult, TransportError> {
        self.lifecycle.require(
            ConnectionState::Authenticated,
            "Must authenticate before executing queries",
        )?;

        // Send execute request
        let request = ExecuteRequest::new(sql.to_string());
        let response: ExecuteResponse = self.send_receive(&request).await?;

        self.query_result_from_response(response)
    }

    async fn fetch_results(
        &mut self,
        handle: ResultSetHandle,
    ) -> Result<ResultData, TransportError> {
        self.lifecycle.require(
            ConnectionState::Authenticated,
            "Must authenticate before fetching results",
        )?;

        // Get max data message size from session info
        let max_bytes = self
            .session_info
            .as_ref()
            .map(|s| s.max_data_message_size)
            .unwrap_or(1024 * 1024); // 1MB default

        // Send fetch request
        let start_position = self
            .fetch_positions
            .get(&handle.as_i32())
            .copied()
            .unwrap_or(0);
        let request = FetchRequest::new(handle.as_i32(), start_position, max_bytes);
        let response: FetchResponse = self.send_receive(&request).await?;

        // Check response status
        self.check_status(&response.status, &response.exception)?;

        // Extract fetch data
        let fetch_data = response
            .response_data
            .ok_or_else(|| TransportError::InvalidResponse("Missing response data".to_string()))?;

        self.fetch_positions.insert(
            handle.as_i32(),
            start_position + fetch_data.data.len() as i64,
        );

        // Note: columns are not included in fetch response,
        // they should be cached from the initial execute response
        // Data is in column-major format from Exasol
        Ok(ResultData {
            columns: vec![],
            data: ResultPayload::Json(fetch_data.data),
            total_rows: fetch_data.num_rows,
        })
    }

    async fn close_result_set(&mut self, handle: ResultSetHandle) -> Result<(), TransportError> {
        self.lifecycle.require(
            ConnectionState::Authenticated,
            "Must authenticate before closing result sets",
        )?;

        self.fetch_positions.remove(&handle.as_i32());

        // Send close result set request
        let request = CloseResultSetRequest::new(vec![handle.as_i32()]);
        let response: CloseResultSetResponse = self.send_receive(&request).await?;

        // Check response status
        self.check_status(&response.status, &response.exception)?;

        Ok(())
    }

    async fn create_prepared_statement(
        &mut self,
        sql: &str,
    ) -> Result<PreparedStatementHandle, TransportError> {
        self.lifecycle.require(
            ConnectionState::Authenticated,
            "Must authenticate before creating prepared statements",
        )?;

        // Send create prepared statement request
        let request = CreatePreparedStatementRequest::new(sql);
        let response: CreatePreparedStatementResponse = self.send_receive(&request).await?;

        // Check response status
        self.check_status(&response.status, &response.exception)?;

        // Extract response data
        let response_data = response
            .response_data
            .ok_or_else(|| TransportError::InvalidResponse("Missing response data".to_string()))?;

        let result_columns = response_data.result_set_columns();

        // Extract parameter types from parameter_data if present
        let (num_params, parameter_types, parameter_names) =
            if let Some(param_data) = response_data.parameter_data {
                let mut types = Vec::with_capacity(param_data.columns.len());
                let mut names = Vec::with_capacity(param_data.columns.len());
                for p in param_data.columns {
                    types.push(p.data_type);
                    names.push(p.name);
                }
                (param_data.num_columns, types, names)
            } else {
                (0, vec![], vec![])
            };

        Ok(PreparedStatementHandle::new(
            response_data.statement_handle,
            num_params,
            parameter_types,
            parameter_names,
        )
        .with_result_columns(result_columns))
    }

    async fn execute_prepared_statement(
        &mut self,
        handle: &PreparedStatementHandle,
        parameters: Option<Vec<Vec<serde_json::Value>>>,
    ) -> Result<QueryResult, TransportError> {
        self.lifecycle.require(
            ConnectionState::Authenticated,
            "Must authenticate before executing prepared statements",
        )?;

        // Build execution request
        let mut request = ExecutePreparedStatementRequest::new(handle.handle);

        // Add parameters if provided
        if let Some(ref data) = parameters {
            // Build column info from the handle's parameter types, or infer from values
            let columns: Vec<_> = if handle.parameter_types.is_empty() {
                // Infer types from the first row of parameter values
                data.iter()
                    .enumerate()
                    .map(|(i, col_values)| {
                        let data_type = col_values
                            .first()
                            .map(super::messages::DataType::infer_from_json)
                            .unwrap_or_else(|| super::messages::DataType::varchar(2_000_000));
                        super::messages::ColumnInfo {
                            name: format!("param{}", i),
                            data_type,
                        }
                    })
                    .collect()
            } else {
                handle
                    .parameter_types
                    .iter()
                    .enumerate()
                    .map(|(i, dt)| super::messages::ColumnInfo {
                        name: format!("param{}", i),
                        data_type: dt.clone(),
                    })
                    .collect()
            };

            request = request.with_data(columns, data.clone());
        }

        // Send execute prepared statement request
        let response: ExecuteResponse = self.send_receive(&request).await?;

        self.query_result_from_response(response)
    }

    async fn close_prepared_statement(
        &mut self,
        handle: &PreparedStatementHandle,
    ) -> Result<(), TransportError> {
        self.lifecycle.require(
            ConnectionState::Authenticated,
            "Must authenticate before closing prepared statements",
        )?;

        // Send close prepared statement request
        let request = ClosePreparedStatementRequest::new(handle.handle);
        let response: ClosePreparedStatementResponse = self.send_receive(&request).await?;

        // Check response status
        self.check_status(&response.status, &response.exception)?;

        Ok(())
    }

    async fn close(&mut self) -> Result<(), TransportError> {
        lifecycle::close_gracefully(self).await
    }

    fn terminate(&mut self) {
        lifecycle::terminate(self);
    }

    fn is_connected(&self) -> bool {
        self.lifecycle.is_open()
    }

    async fn set_autocommit(&mut self, enabled: bool) -> Result<(), TransportError> {
        self.set_attributes(&SetAttributesRequest::autocommit(enabled))
            .await
    }

    async fn set_query_timeout(&mut self, timeout_secs: u64) -> Result<(), TransportError> {
        self.set_attributes(&SetAttributesRequest::query_timeout(timeout_secs))
            .await
    }

    async fn set_current_schema(&mut self, schema: &str) -> Result<(), TransportError> {
        self.set_attributes(&SetAttributesRequest::current_schema(schema))
            .await?;
        self.refresh_current_schema().await?;
        Ok(())
    }

    async fn refresh_current_schema(&mut self) -> Result<Option<String>, TransportError> {
        self.lifecycle.require(
            ConnectionState::Authenticated,
            "Must authenticate before reading attributes",
        )?;

        let response: GetAttributesResponse =
            self.send_receive(&GetAttributesRequest::new()).await?;
        self.check_status(&response.status, &response.exception)?;
        Ok(self.current_schema.clone())
    }

    fn current_schema(&self) -> Option<String> {
        self.current_schema.clone()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::transport::test_support::{
        assert_names_the_termination, finish_within, silent_server_params, test_credentials,
        FakeWebSocketServer, SilentServer, DISCONNECT_BOUND, LOOPBACK_TEST_BOUND,
        SILENT_SERVER_CONNECTION_TIMEOUT,
    };
    use serde_json::{json, Value};
    use std::time::{Duration, Instant};

    const REJECTION_TIMEOUT: Duration = Duration::from_secs(5);

    /// Scenario: Client name and version identify the session
    #[test]
    fn auth_request_reports_the_client_and_the_driver() {
        let credentials = Credentials::new("user".to_string(), "pass".to_string())
            .with_client("dbt".to_string(), "9.9".to_string());

        let request = auth_request(&credentials, "encrypted".to_string());
        let json = serde_json::to_value(&request).unwrap();

        assert_eq!(json["clientName"], "dbt");
        assert_eq!(json["clientVersion"], "9.9");
        assert_eq!(json["driverName"], "exarrow-rs");
        assert_eq!(json["username"], "user");
        assert_eq!(json["password"], "encrypted");
    }

    #[test]
    fn test_websocket_transport_new() {
        let transport = WebSocketTransport::new();
        assert!(!transport.is_connected());
        assert_eq!(transport.lifecycle.state(), ConnectionState::Disconnected);
    }

    #[test]
    fn test_websocket_transport_default() {
        let transport = WebSocketTransport::default();
        assert!(!transport.is_connected());
    }

    #[test]
    fn test_connection_state_transitions() {
        let mut transport = WebSocketTransport::new();
        assert_eq!(transport.lifecycle.state(), ConnectionState::Disconnected);

        // Simulate state changes
        transport.lifecycle = ConnectionLifecycle::in_state(ConnectionState::Connected);
        assert!(transport.is_connected());

        transport.lifecycle = ConnectionLifecycle::in_state(ConnectionState::Authenticated);
        assert!(transport.is_connected());

        transport.lifecycle = ConnectionLifecycle::in_state(ConnectionState::Closed);
        assert!(!transport.is_connected());
    }

    /// Scenario: Terminate a connection whose in-flight response is no longer trusted
    #[test]
    fn terminate_drops_the_session_and_reports_the_transport_disconnected() {
        let mut transport = WebSocketTransport::new();
        transport.lifecycle = ConnectionLifecycle::in_state(ConnectionState::Authenticated);
        transport.session_info = Some(SessionInfo {
            session_id: "12345".to_string(),
            protocol_version: 3,
            release_version: "7.1.0".to_string(),
            database_name: "test_db".to_string(),
            product_name: "EXASolution".to_string(),
            max_data_message_size: 1024 * 1024,
            time_zone: Some("UTC".to_string()),
        });

        transport.terminate();

        assert!(!transport.is_connected());
        assert!(transport.session_info.is_none());
    }

    // Integration-style tests would require a real or mocked WebSocket server
    // For now, we have unit tests for the basic structure

    #[tokio::test]
    async fn test_connect_requires_disconnected_state() {
        let mut transport = WebSocketTransport::new();
        transport.lifecycle = ConnectionLifecycle::in_state(ConnectionState::Connected);

        let params = ConnectionParams::new("localhost".to_string(), 8563);
        let result = transport.connect(&params).await;

        assert!(result.is_err());
        if let Err(TransportError::ProtocolError(msg)) = result {
            assert!(msg.contains("Already connected"));
        } else {
            panic!("Expected ProtocolError");
        }
    }

    /// Scenario: Server that never answers the TLS handshake
    #[tokio::test]
    async fn connect_fails_at_the_tls_handshake_when_the_server_never_answers() {
        finish_within(LOOPBACK_TEST_BOUND, async {
            let mut server = SilentServer::accepting().await;
            let mut transport = WebSocketTransport::new();
            let started = Instant::now();

            let error = transport
                .connect(&silent_server_params(&server))
                .await
                .expect_err("a server that never answers must fail the connection");

            assert!(started.elapsed() >= SILENT_SERVER_CONNECTION_TIMEOUT);
            assert!(
                error
                    .to_string()
                    .contains("Connection timeout after 300ms (TLS handshake)"),
                "{error}"
            );
            finish_within(DISCONNECT_BOUND, server.wait_for_disconnect()).await;
        })
        .await;
    }

    #[tokio::test]
    async fn connect_reports_a_rejected_server_certificate_as_a_tls_error() {
        finish_within(LOOPBACK_TEST_BOUND, async {
            let server = SilentServer::after_tls().await;
            let params = ConnectionParams::new(server.host.clone(), server.port)
                .with_timeout(REJECTION_TIMEOUT.as_millis() as u64);

            let error = WebSocketTransport::new()
                .connect(&params)
                .await
                .expect_err("an untrusted certificate must fail the connection");

            assert!(matches!(error, TransportError::TlsError(_)), "{error:?}");
        })
        .await;
    }

    /// Scenario: WebSocket server that never answers the upgrade
    #[tokio::test]
    async fn connect_fails_at_the_websocket_upgrade_when_the_server_never_answers() {
        finish_within(LOOPBACK_TEST_BOUND, async {
            let mut server = SilentServer::accepting().await;
            let mut transport = WebSocketTransport::new();
            let started = Instant::now();

            let error = transport
                .connect(&silent_server_params(&server).with_tls(false))
                .await
                .expect_err("a server that never answers the upgrade must fail the connection");

            assert!(started.elapsed() >= SILENT_SERVER_CONNECTION_TIMEOUT);
            assert!(
                error
                    .to_string()
                    .contains("Connection timeout after 300ms (WebSocket upgrade)"),
                "{error}"
            );
            finish_within(DISCONNECT_BOUND, server.wait_for_disconnect()).await;
        })
        .await;
    }

    /// Scenario: Server that never answers the login
    #[tokio::test]
    async fn authenticate_fails_at_login_when_the_server_goes_silent_after_the_upgrade() {
        finish_within(LOOPBACK_TEST_BOUND, async {
            let mut server = SilentServer::after_websocket_upgrade().await;
            let mut transport = WebSocketTransport::new();
            let started = Instant::now();

            transport
                .connect(&silent_server_params(&server).with_tls(false))
                .await
                .expect("the WebSocket upgrade completes");
            let error = transport
                .authenticate(&test_credentials())
                .await
                .expect_err("a server that never answers the login must fail it");

            assert!(started.elapsed() >= SILENT_SERVER_CONNECTION_TIMEOUT);
            assert!(
                error
                    .to_string()
                    .contains("Connection timeout after 300ms (login)"),
                "{error}"
            );
            assert!(!transport.is_connected());
            finish_within(DISCONNECT_BOUND, server.wait_for_disconnect()).await;
        })
        .await;
    }

    /// Scenario: Terminate a connection whose in-flight response is no longer trusted
    /// Scenario: Operations after an export timeout name the termination
    #[tokio::test]
    async fn operations_after_terminate_report_the_terminated_transport() {
        let mut transport = WebSocketTransport::new();
        transport.lifecycle = ConnectionLifecycle::in_state(ConnectionState::Authenticated);

        transport.terminate();
        let error = transport
            .execute_query("SELECT 1")
            .await
            .expect_err("a terminated transport refuses queries");
        transport
            .close()
            .await
            .expect("closing a terminated transport succeeds");
        let error_after_close = transport
            .execute_query("SELECT 1")
            .await
            .expect_err("a closed terminated transport still refuses queries");

        assert_names_the_termination(&error);
        assert_names_the_termination(&error_after_close);
        assert!(!transport.is_connected());
    }

    #[tokio::test]
    async fn test_authenticate_requires_connected_state() {
        let mut transport = WebSocketTransport::new();
        assert_eq!(transport.lifecycle.state(), ConnectionState::Disconnected);

        let credentials = Credentials::new("user".to_string(), "pass".to_string());
        let result = transport.authenticate(&credentials).await;

        assert!(result.is_err());
        if let Err(TransportError::ProtocolError(msg)) = result {
            assert!(msg.contains("Must connect before authenticating"));
        } else {
            panic!("Expected ProtocolError");
        }
    }

    #[tokio::test]
    async fn test_execute_query_requires_authenticated_state() {
        let mut transport = WebSocketTransport::new();

        let result = transport.execute_query("SELECT 1").await;

        assert!(result.is_err());
        if let Err(TransportError::ProtocolError(msg)) = result {
            assert!(msg.contains("Must authenticate before executing queries"));
        } else {
            panic!("Expected ProtocolError");
        }
    }

    #[tokio::test]
    async fn test_create_prepared_statement_requires_authenticated_state() {
        let mut transport = WebSocketTransport::new();

        let result = transport
            .create_prepared_statement("SELECT * FROM test WHERE id = ?")
            .await;

        assert!(result.is_err());
        if let Err(TransportError::ProtocolError(msg)) = result {
            assert!(msg.contains("Must authenticate before creating prepared statements"));
        } else {
            panic!("Expected ProtocolError");
        }
    }

    #[tokio::test]
    async fn test_execute_prepared_statement_requires_authenticated_state() {
        let mut transport = WebSocketTransport::new();
        let handle = PreparedStatementHandle::new(1, 0, vec![], vec![]);

        let result = transport.execute_prepared_statement(&handle, None).await;

        assert!(result.is_err());
        if let Err(TransportError::ProtocolError(msg)) = result {
            assert!(msg.contains("Must authenticate before executing prepared statements"));
        } else {
            panic!("Expected ProtocolError");
        }
    }

    #[tokio::test]
    async fn test_close_prepared_statement_requires_authenticated_state() {
        let mut transport = WebSocketTransport::new();
        let handle = PreparedStatementHandle::new(1, 0, vec![], vec![]);

        let result = transport.close_prepared_statement(&handle).await;

        assert!(result.is_err());
        if let Err(TransportError::ProtocolError(msg)) = result {
            assert!(msg.contains("Must authenticate before closing prepared statements"));
        } else {
            panic!("Expected ProtocolError");
        }
    }

    #[tokio::test]
    async fn test_close_idempotent() {
        let mut transport = WebSocketTransport::new();

        // Close when disconnected should be ok
        let result = transport.close().await;
        assert!(result.is_ok());
        assert_eq!(transport.lifecycle.state(), ConnectionState::Disconnected);

        // Close when already closed should be ok
        transport.lifecycle = ConnectionLifecycle::in_state(ConnectionState::Closed);
        let result = transport.close().await;
        assert!(result.is_ok());
    }

    #[test]
    fn test_check_status_ok() {
        let transport = WebSocketTransport::new();
        let result = transport.check_status("ok", &None);
        assert!(result.is_ok());
    }

    #[test]
    fn test_check_status_error() {
        use super::super::messages::ExceptionInfo;

        let transport = WebSocketTransport::new();
        let exception = Some(ExceptionInfo {
            sql_code: Some("42000".to_string()),
            text: "Syntax error".to_string(),
        });

        let result = transport.check_status("error", &exception);
        assert!(result.is_err());

        if let Err(TransportError::ProtocolError(msg)) = result {
            assert!(msg.contains("Syntax error"));
            assert!(msg.contains("42000"));
        } else {
            panic!("Expected ProtocolError");
        }
    }

    #[test]
    fn test_encrypt_password_with_valid_key() {
        // This is a test RSA public key (2048-bit) in PKCS#1 PEM format
        // Generated specifically for testing - NOT for production use
        let test_public_key_pem = r#"-----BEGIN RSA PUBLIC KEY-----
MIIBCgKCAQEAulMKxKfPd02qNEVCU1M6hG/Vc9xz+0u+N47Qqa1Y0E2A5bDiz3XA
aCg2d65C7DyuTL38zwmOtjagvvIAgRj9yDf0v1/v9e1X4l5XE6UiaKKqdcXNy6lJ
QspqkOBUptlz+2h/G8Z12++xUo/4AGAGz9ZkrRRvcTGW1GJhCROizeJhTpGMpc/v
o1G53uy2eTHwnz5S3YgJF7nfX60wjJ99ifQuQ9BhDIYLNqzwHTzExMN63v0UOBIL
vJ+yVUqh0/T2f5e9E1lDNuIqLyXe8VwwUsS72A1EGtg0s77+xUQ7KiGRbHD4bsBo
A74EI7MHQ7163wVPT0VWFRvUmmv+UO7W8wIDAQAB
-----END RSA PUBLIC KEY-----"#;

        let result = WebSocketTransport::encrypt_password("test_password", test_public_key_pem);

        assert!(
            result.is_ok(),
            "encrypt_password failed: {:?}",
            result.err()
        );
        let encrypted = result.unwrap();

        // The encrypted result should be base64-encoded
        assert!(!encrypted.is_empty());

        // Verify it's valid base64
        let decoded = STANDARD.decode(&encrypted);
        assert!(decoded.is_ok());

        // RSA 2048-bit encryption produces 256 bytes
        assert_eq!(decoded.unwrap().len(), 256);

        // Additionally verify the new parser accepts this 2048-bit key
        let der = WebSocketTransport::pem_to_pkcs1_der(test_public_key_pem).unwrap();
        let (n, e) = WebSocketTransport::parse_rsa_public_key(&der).unwrap();
        // 2048-bit modulus may report 2040..=2048 bits due to leading-bit nondeterminism
        assert!(
            n.bits() > 2040 && n.bits() <= 2048,
            "Expected ~2048-bit modulus, got {} bits",
            n.bits()
        );
        assert_eq!(
            e,
            BigUint::from(65537u32),
            "Expected standard RSA exponent 65537"
        );
    }

    #[test]
    fn test_encrypt_password_with_invalid_key() {
        let invalid_pem = "not a valid PEM key";

        let result = WebSocketTransport::encrypt_password("password", invalid_pem);

        assert!(result.is_err());
        if let Err(TransportError::ProtocolError(msg)) = result {
            assert!(msg.contains("Failed to parse RSA public key"));
        } else {
            panic!("Expected ProtocolError");
        }
    }

    #[test]
    fn test_encrypt_password_with_1024_bit_key() {
        // Test-only 1024-bit RSA public key (no private key committed). Do not use in production.
        let pem_1024 = r#"-----BEGIN RSA PUBLIC KEY-----
MIGJAoGBAMh0gXUltxbJYmwQUIvXPLl9Y8bGaGFN/urgclF3Czd7viHaAMuanebQ
62s1mLV0vXaYSZk8Zsrat2T/i7jbPE0XKpVUgmnlT/CHXv6gPdTpOr3JTpo/lop0
t6J/6xJBNQDp6OrFMtTTq2M3zxSfcomlT4Q759uuGkEdM9crb8A9AgMBAAE=
-----END RSA PUBLIC KEY-----"#;

        let result = WebSocketTransport::encrypt_password("hunter2", pem_1024);
        assert!(
            result.is_ok(),
            "encrypt_password with 1024-bit key failed: {:?}",
            result.err()
        );
        let encrypted = result.unwrap();

        // Base64-decode: length MUST equal 128 bytes (1024 bits / 8)
        let decoded = STANDARD.decode(&encrypted).expect("valid base64");
        assert_eq!(
            decoded.len(),
            128,
            "1024-bit RSA ciphertext must be 128 bytes, got {}",
            decoded.len()
        );

        // Encrypt again with a different password: ciphertexts must differ (random PS padding)
        let result2 = WebSocketTransport::encrypt_password("different_password", pem_1024);
        assert!(result2.is_ok());
        let encrypted2 = result2.unwrap();
        assert_ne!(
            encrypted, encrypted2,
            "Different passwords must produce different ciphertexts"
        );

        // Even the same password should produce different ciphertext due to random padding
        let result3 = WebSocketTransport::encrypt_password("hunter2", pem_1024);
        assert!(result3.is_ok());
        let encrypted3 = result3.unwrap();
        assert_ne!(
            encrypted, encrypted3,
            "Same password with random padding should produce different ciphertexts"
        );
    }

    #[test]
    fn test_encrypt_password_rejects_too_small_key() {
        // Test-only 512-bit RSA public key. Do not use in production.
        let pem_512 = r#"-----BEGIN RSA PUBLIC KEY-----
MEgCQQCrHBlqjk3p2boRBFTAZdqcWNU+g5LjKicoOX5UIyIanBCV5fgbtoRCCvBr
++vdlAaIAcJx5iKBMp1obShMOPwVAgMBAAE=
-----END RSA PUBLIC KEY-----"#;

        let result = WebSocketTransport::encrypt_password("pw", pem_512);
        assert!(result.is_err(), "512-bit key should be rejected");
        if let Err(TransportError::ProtocolError(msg)) = result {
            assert!(
                msg.contains("Unsupported RSA key size"),
                "Error should mention unsupported key size, got: {}",
                msg
            );
            assert!(
                msg.contains("512"),
                "Error should mention 512 bits, got: {}",
                msg
            );
        } else {
            panic!("Expected ProtocolError");
        }
    }

    #[test]
    fn test_parse_rsa_public_key_roundtrip() {
        // Use the 2048-bit key from the existing test
        let test_public_key_pem = r#"-----BEGIN RSA PUBLIC KEY-----
MIIBCgKCAQEAulMKxKfPd02qNEVCU1M6hG/Vc9xz+0u+N47Qqa1Y0E2A5bDiz3XA
aCg2d65C7DyuTL38zwmOtjagvvIAgRj9yDf0v1/v9e1X4l5XE6UiaKKqdcXNy6lJ
QspqkOBUptlz+2h/G8Z12++xUo/4AGAGz9ZkrRRvcTGW1GJhCROizeJhTpGMpc/v
o1G53uy2eTHwnz5S3YgJF7nfX60wjJ99ifQuQ9BhDIYLNqzwHTzExMN63v0UOBIL
vJ+yVUqh0/T2f5e9E1lDNuIqLyXe8VwwUsS72A1EGtg0s77+xUQ7KiGRbHD4bsBo
A74EI7MHQ7163wVPT0VWFRvUmmv+UO7W8wIDAQAB
-----END RSA PUBLIC KEY-----"#;

        let der = WebSocketTransport::pem_to_pkcs1_der(test_public_key_pem).unwrap();
        let (n, e) = WebSocketTransport::parse_rsa_public_key(&der).unwrap();

        // Verify standard RSA exponent
        assert_eq!(
            e,
            BigUint::from(65537u32),
            "Expected standard RSA exponent 65537"
        );

        // Verify modulus size: n.to_bytes_be() should be 256 bytes for a 2048-bit key
        let n_bytes = n.to_bytes_be();
        assert_eq!(
            n_bytes.len(),
            256,
            "2048-bit modulus should be 256 bytes in big-endian, got {}",
            n_bytes.len()
        );
        // The high byte should have the top bit set (modulus is close to 2^2048)
        assert!(
            n_bytes[0] & 0x80 != 0,
            "High byte of 2048-bit modulus should have top bit set, got 0x{:02x}",
            n_bytes[0]
        );

        // Verify bit length
        assert!(
            n.bits() > 2040 && n.bits() <= 2048,
            "Expected ~2048-bit modulus, got {} bits",
            n.bits()
        );
    }

    #[test]
    fn test_encrypt_password_with_1024_bit_key_decrypts_correctly() {
        // Test-only 1024-bit RSA keypair. Do not use in production.
        let pub_pem = r#"-----BEGIN RSA PUBLIC KEY-----
MIGJAoGBAM5D4w7VjlVI8hUmJj8MXx4glxP3dqroATH1hito7CzfGGD/Ss73lp/n
0JLwI4oT6g0wBiMlLkPY6C+hfTI2x/UhJE8gKhz6URld6uQ29d0PAq4OvsZW5Rhz
nuIWmHv9WKvS/5DVcNbBtkiTJ9BDnsSQW/YqA4DQGmzph4ThqEkJAgMBAAE=
-----END RSA PUBLIC KEY-----"#;

        // PKCS#1 private key DER fields: SEQUENCE of 9 INTEGERs
        // [version, n, e, d, p, q, dp, dq, qinv]
        // We only need n and d for decryption.
        let priv_pem = r#"-----BEGIN RSA PRIVATE KEY-----
MIICXQIBAAKBgQDOQ+MO1Y5VSPIVJiY/DF8eIJcT93aq6AEx9YYraOws3xhg/0rO
95af59CS8COKE+oNMAYjJS5D2OgvoX0yNsf1ISRPICoc+lEZXerkNvXdDwKuDr7G
VuUYc57iFph7/Vir0v+Q1XDWwbZIkyfQQ57EkFv2KgOA0Bps6YeE4ahJCQIDAQAB
AoGAVYDAu+J86Q+fAnNZAWPAfj2mQumfMIOSE0KjBpWs6YDlmzfYq+jocIro5DBV
myRcLnFM6f68qfVdcnkv68PXqSA7acyTtAKSIJAgN1xiYELuRVWMk/+UVgGhRpcH
rY4sTIwPM5b9r6JA++6PX13b8qqybPijf/Lz5urEbU3oPu0CQQDmZrbz623uijm4
ifEhk6f+Gq4spF5tUHVwY/GlfdtaDPr/wSTBkeKweZIAd9LJh8iv7Il2tYNKzzot
BTMut3VnAkEA5S6sKNaeLQevYy7N2zKm2SdSKKo1Sh48+oLYd2CZQGG5jDRom6p/
e+4hol1jGDBwvx0sMrFCFsoQXRiP2etYDwJBANSDj2LjF837Xwww5+IhkMVXlKoG
njZUDU6yUPRlZwrjiCyY2S9WQXKnX5zg6OMMRHbIRW7iM4ywIafe8Pu5KicCQQCE
rB8nyQ5qfP9wSGENWuYx4cxzFA2jaZvdXa/Yc8hj9+7FFnXUX8BLSxCXgL5j+27Z
hBbZBbp/nNwaOKTV/6LLAkAjNOwM7VST6Medyd2lNpsCjmYAoF0eUf8Z8ZJPaZeo
El6NrMeFybqeqwjPHPG1oCwg4YIeaT8ZB2qUW143brUB
-----END RSA PRIVATE KEY-----"#;

        let password = "exasol_secret_123";
        let encrypted = WebSocketTransport::encrypt_password(password, pub_pem)
            .expect("encryption should succeed");

        // Decode ciphertext
        let ciphertext = STANDARD.decode(&encrypted).expect("valid base64");
        assert_eq!(
            ciphertext.len(),
            128,
            "1024-bit RSA ciphertext must be 128 bytes"
        );

        // Parse the private key to extract n and d
        let (n, d) = parse_pkcs1_private_key_n_d(priv_pem);

        // RSA decrypt: m = c^d mod n
        let c = BigUint::from_bytes_be(&ciphertext);
        let m = c.modpow(&d, &n);

        // Convert to fixed-width bytes and strip PKCS#1 v1.5 padding
        let k = 128; // 1024 bits / 8
        let mut em = vec![0u8; k];
        let m_bytes = m.to_bytes_be();
        em[k - m_bytes.len()..].copy_from_slice(&m_bytes);

        // EM format: 0x00 || 0x02 || PS (non-zero) || 0x00 || M
        assert_eq!(em[0], 0x00, "EM must start with 0x00");
        assert_eq!(em[1], 0x02, "EM block type must be 0x02");

        // Find the 0x00 separator after PS
        let separator_pos = em[2..]
            .iter()
            .position(|&b| b == 0x00)
            .expect("Must find 0x00 separator after PS padding");
        let message_start = 2 + separator_pos + 1;

        // Verify PS is all non-zero
        for (i, &b) in em[2..2 + separator_pos].iter().enumerate() {
            assert_ne!(b, 0x00, "PS byte at position {} must be non-zero", i);
        }

        // Verify PS is at least 8 bytes (PKCS#1 v1.5 requirement)
        assert!(
            separator_pos >= 8,
            "PS must be at least 8 bytes, got {}",
            separator_pos
        );

        // Extract and verify the decrypted message matches the original password
        let decrypted = &em[message_start..];
        assert_eq!(
            decrypted,
            password.as_bytes(),
            "Decrypted message must match original password"
        );
    }

    /// Parse a PKCS#1 RSA private key PEM to extract (n, d).
    /// Only used in tests. The PKCS#1 private key is:
    /// SEQUENCE { version, n, e, d, p, q, dp, dq, qinv }
    fn parse_pkcs1_private_key_n_d(pem: &str) -> (BigUint, BigUint) {
        // Strip PEM headers
        let start_marker = "-----BEGIN RSA PRIVATE KEY-----";
        let end_marker = "-----END RSA PRIVATE KEY-----";
        let start = pem.find(start_marker).unwrap() + start_marker.len();
        let end = pem.find(end_marker).unwrap();
        let base64_content: String = pem[start..end]
            .chars()
            .filter(|c| !c.is_whitespace())
            .collect();
        let der = STANDARD.decode(&base64_content).unwrap();

        let mut pos = 0;

        // Outer SEQUENCE
        assert_eq!(der[pos], 0x30);
        pos += 1;
        let (_, len_bytes) = read_test_der_length(&der[pos..]);
        pos += len_bytes;

        // version INTEGER (skip)
        let (_, consumed) = read_test_der_integer(&der[pos..]);
        pos += consumed;

        // n INTEGER
        let (n_bytes, consumed) = read_test_der_integer(&der[pos..]);
        pos += consumed;
        let n = BigUint::from_bytes_be(n_bytes);

        // e INTEGER (skip)
        let (_, consumed) = read_test_der_integer(&der[pos..]);
        pos += consumed;

        // d INTEGER
        let (d_bytes, _) = read_test_der_integer(&der[pos..]);
        let d = BigUint::from_bytes_be(d_bytes);

        (n, d)
    }

    fn read_test_der_length(data: &[u8]) -> (usize, usize) {
        let first = data[0];
        if first < 0x80 {
            (first as usize, 1)
        } else {
            let num_bytes = (first & 0x7F) as usize;
            let mut len: usize = 0;
            for &b in &data[1..1 + num_bytes] {
                len = (len << 8) | (b as usize);
            }
            (len, 1 + num_bytes)
        }
    }

    fn read_test_der_integer(data: &[u8]) -> (&[u8], usize) {
        assert_eq!(data[0], 0x02, "Expected INTEGER tag");
        let (int_len, len_bytes) = read_test_der_length(&data[1..]);
        let header_len = 1 + len_bytes;
        let mut value = &data[header_len..header_len + int_len];
        // Strip leading 0x00 sign-bit padding
        if value.len() > 1 && value[0] == 0x00 {
            value = &value[1..];
        }
        (value, header_len + int_len)
    }

    #[test]
    fn test_check_status_error_with_none_exception() {
        let transport = WebSocketTransport::new();

        let result = transport.check_status("error", &None);

        assert!(result.is_err());
        if let Err(TransportError::ProtocolError(msg)) = result {
            assert_eq!(msg, "Unknown error");
        } else {
            panic!("Expected ProtocolError with 'Unknown error' message");
        }
    }

    #[test]
    fn test_check_status_error_with_exception_missing_sql_code() {
        use super::super::messages::ExceptionInfo;

        let transport = WebSocketTransport::new();
        let exception = Some(ExceptionInfo {
            sql_code: None,
            text: "Some database error".to_string(),
        });

        let result = transport.check_status("error", &exception);

        assert!(result.is_err());
        if let Err(TransportError::ProtocolError(msg)) = result {
            assert!(msg.contains("Some database error"));
            assert!(msg.contains("unknown")); // sql_code defaults to "unknown"
        } else {
            panic!("Expected ProtocolError");
        }
    }

    #[tokio::test]
    async fn test_fetch_results_requires_authenticated_state() {
        let mut transport = WebSocketTransport::new();
        let handle = super::super::messages::ResultSetHandle::new(1);

        let result = transport.fetch_results(handle).await;

        assert!(result.is_err());
        if let Err(TransportError::ProtocolError(msg)) = result {
            assert!(msg.contains("Must authenticate before fetching results"));
        } else {
            panic!("Expected ProtocolError");
        }
    }

    #[tokio::test]
    async fn test_fetch_results_requires_authenticated_state_from_connected() {
        let mut transport = WebSocketTransport::new();
        transport.lifecycle = ConnectionLifecycle::in_state(ConnectionState::Connected);
        let handle = super::super::messages::ResultSetHandle::new(1);

        let result = transport.fetch_results(handle).await;

        assert!(result.is_err());
        if let Err(TransportError::ProtocolError(msg)) = result {
            assert!(msg.contains("Must authenticate before fetching results"));
        } else {
            panic!("Expected ProtocolError");
        }
    }

    #[tokio::test]
    async fn test_close_result_set_requires_authenticated_state() {
        let mut transport = WebSocketTransport::new();
        let handle = super::super::messages::ResultSetHandle::new(1);

        let result = transport.close_result_set(handle).await;

        assert!(result.is_err());
        if let Err(TransportError::ProtocolError(msg)) = result {
            assert!(msg.contains("Must authenticate before closing result sets"));
        } else {
            panic!("Expected ProtocolError");
        }
    }

    #[tokio::test]
    async fn test_close_result_set_requires_authenticated_state_from_connected() {
        let mut transport = WebSocketTransport::new();
        transport.lifecycle = ConnectionLifecycle::in_state(ConnectionState::Connected);
        let handle = super::super::messages::ResultSetHandle::new(1);

        let result = transport.close_result_set(handle).await;

        assert!(result.is_err());
        if let Err(TransportError::ProtocolError(msg)) = result {
            assert!(msg.contains("Must authenticate before closing result sets"));
        } else {
            panic!("Expected ProtocolError");
        }
    }

    #[tokio::test]
    async fn test_close_from_connected_state_succeeds() {
        let mut transport = WebSocketTransport::new();
        transport.lifecycle = ConnectionLifecycle::in_state(ConnectionState::Connected);

        // Close when connected (but not authenticated) should succeed
        // and should NOT try to send disconnect request
        let result = transport.close().await;

        assert!(result.is_ok());
        assert_eq!(transport.lifecycle.state(), ConnectionState::Closed);
    }

    #[tokio::test]
    async fn test_close_clears_session_info() {
        use super::super::messages::SessionInfo;

        let mut transport = WebSocketTransport::new();
        transport.lifecycle = ConnectionLifecycle::in_state(ConnectionState::Connected);
        transport.session_info = Some(SessionInfo {
            session_id: "12345".to_string(),
            protocol_version: 3,
            release_version: "7.1.0".to_string(),
            database_name: "test_db".to_string(),
            product_name: "EXASolution".to_string(),
            max_data_message_size: 1024 * 1024,
            time_zone: Some("UTC".to_string()),
        });

        let result = transport.close().await;

        assert!(result.is_ok());
        assert!(transport.session_info.is_none());
    }

    fn decimal_column() -> Value {
        json!({"name": "ID", "dataType": {"type": "DECIMAL", "precision": 18, "scale": 0}})
    }

    fn execute_response(handle: i32, total_rows: i64, inline: &[i64]) -> Value {
        json!({
            "status": "ok",
            "responseData": {
                "numResults": 1,
                "results": [{
                    "resultType": "resultSet",
                    "resultSet": {
                        "resultSetHandle": handle,
                        "numColumns": 1,
                        "numRows": total_rows,
                        "numRowsInMessage": inline.len(),
                        "columns": [decimal_column()],
                        "data": [inline],
                    }
                }]
            }
        })
    }

    fn fetch_response(rows: &[i64]) -> Value {
        json!({"status": "ok", "responseData": {"numRows": rows.len(), "data": [rows]}})
    }

    fn fetch_start_positions(server: &FakeWebSocketServer) -> Vec<i64> {
        server
            .requests()
            .iter()
            .filter(|request| request["command"] == "fetch")
            .map(|request| request["startPosition"].as_i64().expect("startPosition"))
            .collect()
    }

    async fn authenticated_transport(server: &FakeWebSocketServer) -> WebSocketTransport {
        let mut transport = WebSocketTransport::new();
        let params = ConnectionParams::new("127.0.0.1".to_string(), server.port).with_tls(false);
        transport.connect(&params).await.expect("connect");
        transport.lifecycle = ConnectionLifecycle::in_state(ConnectionState::Authenticated);
        transport.session_info = Some(SessionInfo {
            session_id: "1".to_string(),
            protocol_version: 3,
            release_version: "2026.1.0".to_string(),
            database_name: "exa".to_string(),
            product_name: "Exasol".to_string(),
            max_data_message_size: 1024 * 1024,
            time_zone: None,
        });
        transport
    }

    /// Scenario: Fetch results command
    #[tokio::test]
    async fn test_fetch_results_starts_after_the_inline_rows_and_advances_per_page() {
        let server = FakeWebSocketServer::scripted(vec![
            execute_response(7, 5, &[1, 2]),
            execute_response(8, 5, &[]),
            fetch_response(&[3, 4]),
            fetch_response(&[1, 2]),
            fetch_response(&[5, 6]),
        ])
        .await;
        let mut transport = authenticated_transport(&server).await;

        transport.execute_query("SELECT 7").await.unwrap();
        transport.execute_query("SELECT 8").await.unwrap();
        transport
            .fetch_results(ResultSetHandle::new(7))
            .await
            .unwrap();
        transport
            .fetch_results(ResultSetHandle::new(8))
            .await
            .unwrap();
        transport
            .fetch_results(ResultSetHandle::new(7))
            .await
            .unwrap();

        assert_eq!(fetch_start_positions(&server), vec![2, 0, 4]);
    }

    /// Scenario: Fetch results command
    #[tokio::test]
    async fn test_prepared_statement_fetch_starts_after_the_inline_rows() {
        let server = FakeWebSocketServer::scripted(vec![
            execute_response(9, 10, &[1, 2, 3]),
            fetch_response(&[4, 5]),
        ])
        .await;
        let mut transport = authenticated_transport(&server).await;
        let statement = PreparedStatementHandle::new(1, 0, vec![], vec![]);

        transport
            .execute_prepared_statement(&statement, None)
            .await
            .unwrap();
        transport
            .fetch_results(ResultSetHandle::new(9))
            .await
            .unwrap();

        assert_eq!(fetch_start_positions(&server), vec![3]);
    }

    #[tokio::test]
    async fn test_close_result_set_forgets_the_fetch_position() {
        let server = FakeWebSocketServer::scripted(vec![
            execute_response(7, 5, &[1, 2]),
            json!({"status": "ok"}),
            fetch_response(&[1]),
        ])
        .await;
        let mut transport = authenticated_transport(&server).await;

        transport.execute_query("SELECT 7").await.unwrap();
        transport
            .close_result_set(ResultSetHandle::new(7))
            .await
            .unwrap();
        transport
            .fetch_results(ResultSetHandle::new(7))
            .await
            .unwrap();

        assert_eq!(fetch_start_positions(&server), vec![0]);
    }

    fn row_count_response(attributes: Option<Value>) -> Value {
        let mut response = json!({
            "status": "ok",
            "responseData": {
                "numResults": 1,
                "results": [{"resultType": "rowCount", "rowCount": 0}]
            }
        });
        if let Some(attributes) = attributes {
            response["attributes"] = attributes;
        }
        response
    }

    /// A `getAttributes` answer with the mixed value types the server sends.
    fn get_attributes_response(current_schema: &str) -> Value {
        json!({
            "status": "ok",
            "attributes": {
                "autocommit": true,
                "queryTimeout": 0,
                "currentSchema": current_schema,
                "snapshotTransactionsEnabled": false
            }
        })
    }

    async fn execute_each(transport: &mut WebSocketTransport, statements: &[&str]) {
        for sql in statements {
            transport.execute_query(sql).await.expect("execute");
        }
    }

    /// Scenario: Set the current schema attribute
    #[tokio::test]
    async fn set_current_schema_sends_set_attributes_then_records_the_servers_name() {
        let server = FakeWebSocketServer::scripted(vec![
            json!({"status": "ok"}),
            get_attributes_response("MY_SCHEMA"),
        ])
        .await;
        let mut transport = authenticated_transport(&server).await;

        transport.set_current_schema("my_schema").await.unwrap();

        assert_eq!(
            server.requests(),
            vec![
                json!({"command": "setAttributes", "attributes": {"currentSchema": "my_schema"}}),
                json!({"command": "getAttributes"}),
            ]
        );
        assert_eq!(transport.current_schema(), Some("MY_SCHEMA".to_string()));
    }

    /// Scenario: Set the current schema attribute
    #[tokio::test]
    async fn set_current_schema_returns_the_servers_rejection() {
        let server = FakeWebSocketServer::scripted(vec![
            json!({"status": "ok"}),
            get_attributes_response("SALES"),
            json!({
                "status": "error",
                "exception": {"sqlCode": "42000", "text": "schema ZZ_TYPO not found"}
            }),
            get_attributes_response("ZZ_TYPO"),
        ])
        .await;
        let mut transport = authenticated_transport(&server).await;
        transport.set_current_schema("SALES").await.unwrap();

        let error = transport.set_current_schema("ZZ_TYPO").await.unwrap_err();

        assert!(
            error.to_string().contains("schema ZZ_TYPO not found"),
            "{error}"
        );
        assert_eq!(transport.current_schema(), Some("SALES".to_string()));
        assert_eq!(
            server.requests().len(),
            3,
            "a rejected set must not send getAttributes"
        );
    }

    #[tokio::test]
    async fn refresh_current_schema_reports_an_empty_value_as_no_schema() {
        let server = FakeWebSocketServer::scripted(vec![
            row_count_response(Some(json!({"currentSchema": "S1"}))),
            get_attributes_response(""),
        ])
        .await;
        let mut transport = authenticated_transport(&server).await;
        execute_each(&mut transport, &["OPEN SCHEMA S1"]).await;

        let refreshed = transport.refresh_current_schema().await.unwrap();

        assert_eq!(refreshed, None);
        assert_eq!(transport.current_schema(), None);
    }

    #[tokio::test]
    async fn refresh_current_schema_returns_the_servers_error() {
        let server = FakeWebSocketServer::scripted(vec![
            row_count_response(Some(json!({"currentSchema": "S1"}))),
            json!({
                "status": "error",
                "exception": {"sqlCode": "42000", "text": "getAttributes rejected"}
            }),
        ])
        .await;
        let mut transport = authenticated_transport(&server).await;
        execute_each(&mut transport, &["OPEN SCHEMA S1"]).await;

        let error = transport.refresh_current_schema().await.unwrap_err();

        assert!(
            error.to_string().contains("getAttributes rejected"),
            "{error}"
        );
        assert_eq!(transport.current_schema(), Some("S1".to_string()));
    }

    /// Scenario: Track the current schema attribute from responses
    #[tokio::test]
    async fn response_attributes_update_the_current_schema() {
        let server = FakeWebSocketServer::scripted(vec![row_count_response(Some(
            json!({"currentSchema": "S1"}),
        ))])
        .await;
        let mut transport = authenticated_transport(&server).await;

        execute_each(&mut transport, &["OPEN SCHEMA S1"]).await;

        assert_eq!(transport.current_schema(), Some("S1".to_string()));
    }

    /// Scenario: Track the current schema attribute from responses
    #[tokio::test]
    async fn response_without_current_schema_leaves_the_current_schema_unchanged() {
        let server = FakeWebSocketServer::scripted(vec![
            row_count_response(Some(json!({"currentSchema": "S1"}))),
            row_count_response(None),
            row_count_response(Some(json!({"autocommit": true}))),
        ])
        .await;
        let mut transport = authenticated_transport(&server).await;

        execute_each(&mut transport, &["OPEN SCHEMA S1", "SELECT 1", "COMMIT"]).await;

        assert_eq!(transport.current_schema(), Some("S1".to_string()));
    }

    /// Scenario: Track the current schema attribute from responses
    #[tokio::test]
    async fn empty_current_schema_in_a_response_means_no_current_schema() {
        let server = FakeWebSocketServer::scripted(vec![
            row_count_response(Some(json!({"currentSchema": "S1"}))),
            row_count_response(Some(json!({"currentSchema": ""}))),
        ])
        .await;
        let mut transport = authenticated_transport(&server).await;

        execute_each(&mut transport, &["OPEN SCHEMA S1", "CLOSE SCHEMA"]).await;

        assert_eq!(transport.current_schema(), None);
    }

    #[tokio::test]
    async fn closing_the_transport_forgets_the_current_schema() {
        let server = FakeWebSocketServer::scripted(vec![
            row_count_response(Some(json!({"currentSchema": "S1"}))),
            json!({"status": "ok"}),
        ])
        .await;
        let mut transport = authenticated_transport(&server).await;
        execute_each(&mut transport, &["OPEN SCHEMA S1"]).await;

        transport.close().await.unwrap();

        assert_eq!(transport.current_schema(), None);
    }
}