br_pgsql/

connect.rs

use crate::config::Config;
use crate::error::PgsqlError;
use crate::packet::{AuthStatus, Packet, SuccessMessage};
use std::io::{Read, Write};
use std::net::TcpStream;
use std::sync::Arc;
use std::time::Duration;

#[derive(Clone, Debug)]
pub struct Connect {
    pub(crate) stream: Arc<TcpStream>,
    packet: Packet,
    auth_status: AuthStatus,
}

impl Connect {
    pub fn is_valid(&mut self) -> bool {
        self.query("SELECT 1").is_ok()
    }

    pub fn _close(&mut self) {
        let _ = self.stream.as_ref().write_all(&Packet::pack_terminate());
        let _ = self.stream.shutdown(std::net::Shutdown::Both);
    }

    pub fn new(mut config: Config) -> Result<Connect, PgsqlError> {
        let stream =
            TcpStream::connect(config.url()).map_err(|e| PgsqlError::Connection(e.to_string()))?;

        stream
            .set_read_timeout(Some(Duration::from_secs(30)))
            .map_err(|e| PgsqlError::Connection(format!("设置读取超时失败: {}", e)))?;
        stream
            .set_write_timeout(Some(Duration::from_secs(30)))
            .map_err(|e| PgsqlError::Connection(format!("设置写入超时失败: {}", e)))?;

        let _ = stream.peer_addr();

        let mut connect = Self {
            stream: Arc::new(stream),
            packet: Packet::new(config),
            auth_status: AuthStatus::None,
        };

        connect.authenticate()?;

        Ok(connect)
    }

    fn authenticate(&mut self) -> Result<(), PgsqlError> {
        self.stream
            .as_ref()
            .write_all(&self.packet.pack_first())
            .map_err(|e| PgsqlError::Auth(format!("发送 startup message 失败: {}", e)))?;

        let data = self.read()?;
        self.packet.unpack(data, 0)?;

        if !self.packet.md5_salt.is_empty() {
            self.md5_auth()?;
        } else if self.packet.auth_mechanism.is_empty() && self.packet.md5_salt.is_empty() {
            self.cleartext_auth()?;
        } else {
            self.scram_auth()?;
        }

        self.auth_status = AuthStatus::AuthenticationOk;
        Ok(())
    }

    fn md5_auth(&mut self) -> Result<(), PgsqlError> {
        self.stream
            .as_ref()
            .write_all(&self.packet.pack_md5_password())
            .map_err(|e| PgsqlError::Auth(format!("发送 MD5 密码失败: {}", e)))?;

        let data = self.read()?;
        self.packet.unpack(data, 0)?;
        Ok(())
    }

    fn cleartext_auth(&mut self) -> Result<(), PgsqlError> {
        self.stream
            .as_ref()
            .write_all(&self.packet.pack_cleartext_password())
            .map_err(|e| PgsqlError::Auth(format!("发送明文密码失败: {}", e)))?;

        let data = self.read()?;
        self.packet.unpack(data, 0)?;
        Ok(())
    }

    fn scram_auth(&mut self) -> Result<(), PgsqlError> {
        self.stream
            .as_ref()
            .write_all(&self.packet.pack_auth())
            .map_err(|e| PgsqlError::Auth(format!("发送 SASL Initial Response 失败: {}", e)))?;

        let data = self.read()?;
        self.packet.unpack(data, 0)?;

        self.stream
            .as_ref()
            .write_all(&self.packet.pack_auth_verify())
            .map_err(|e| PgsqlError::Auth(format!("发送 SASL Verify 失败: {}", e)))?;

        let data = self.read()?;
        self.packet.unpack(data, 0)?;
        Ok(())
    }

    fn read(&mut self) -> Result<Vec<u8>, PgsqlError> {
        let mut msg = Vec::new();
        let mut buf = [0u8; 4096];
        let mut retry_count = 0;

        #[cfg(not(test))]
        const MAX_RETRIES: u32 = 100;
        #[cfg(test)]
        const MAX_RETRIES: u32 = 3;

        #[cfg(not(test))]
        const MAX_MESSAGE_SIZE: usize = 256 * 1024 * 1024;
        #[cfg(test)]
        const MAX_MESSAGE_SIZE: usize = 128;

        #[cfg(not(test))]
        let deadline = std::time::Instant::now() + Duration::from_secs(300);
        #[cfg(test)]
        let deadline = std::time::Instant::now() + Duration::from_millis(200);

        loop {
            if std::time::Instant::now() >= deadline {
                return Err(PgsqlError::Timeout("读取总超时".into()));
            }

            match self.stream.as_ref().read(&mut buf) {
                Ok(0) => return Err(PgsqlError::Connection("连接已关闭或服务端断开".into())),
                Ok(n) => {
                    if msg.len() + n > MAX_MESSAGE_SIZE {
                        return Err(PgsqlError::Protocol("消息超过最大允许大小".into()));
                    }
                    msg.extend_from_slice(&buf[..n]);
                    retry_count = 0;
                }
                Err(ref e)
                    if e.kind() == std::io::ErrorKind::WouldBlock
                        || e.kind() == std::io::ErrorKind::TimedOut =>
                {
                    retry_count += 1;
                    if retry_count > MAX_RETRIES {
                        return Err(PgsqlError::Timeout("读取超时，已达最大重试次数".into()));
                    }
                    std::thread::sleep(Duration::from_millis(10));
                    continue;
                }
                Err(e) => return Err(PgsqlError::Io(e)),
            };

            if let AuthStatus::AuthenticationOk = self.auth_status {
                if msg.ends_with(&[90, 0, 0, 0, 5, 73])
                    || msg.ends_with(&[90, 0, 0, 0, 5, 84])
                    || msg.ends_with(&[90, 0, 0, 0, 5, 69])
                {
                    break;
                }
            } else if msg.len() >= 5 {
                let len_bytes = &msg[1..=4];
                if let Ok(len) = len_bytes.try_into().map(u32::from_be_bytes) {
                    if msg.len() > len as usize {
                        break;
                    }
                }
            }
        }

        Ok(msg)
    }

    pub fn query(&mut self, sql: &str) -> Result<SuccessMessage, PgsqlError> {
        self.stream
            .as_ref()
            .write_all(&self.packet.pack_query(sql))
            .map_err(PgsqlError::Io)?;

        let data = self.read()?;

        self.packet.unpack(data, 0)
    }

    pub fn execute(&mut self, sql: &str) -> Result<SuccessMessage, PgsqlError> {
        self.stream
            .as_ref()
            .write_all(&self.packet.pack_execute(sql))
            .map_err(PgsqlError::Io)?;
        let data = self.read()?;
        self.packet.unpack(data, 0)
    }

    /// 参数化查询
    pub fn query_params(
        &mut self,
        sql: &str,
        params: &[Option<&str>],
    ) -> Result<SuccessMessage, PgsqlError> {
        self.stream
            .as_ref()
            .write_all(&self.packet.pack_query_params(sql, params))
            .map_err(PgsqlError::Io)?;

        let data = self.read()?;
        self.packet.unpack(data, 0)
    }

    /// 参数化执行
    pub fn execute_params(
        &mut self,
        sql: &str,
        params: &[Option<&str>],
    ) -> Result<SuccessMessage, PgsqlError> {
        self.stream
            .as_ref()
            .write_all(&self.packet.pack_execute_params(sql, params))
            .map_err(PgsqlError::Io)?;
        let data = self.read()?;
        self.packet.unpack(data, 0)
    }

    /// 参数化查询（便捷版，所有参数非 NULL）
    pub fn query_str(&mut self, sql: &str, params: &[&str]) -> Result<SuccessMessage, PgsqlError> {
        let opts: Vec<Option<&str>> = params.iter().map(|s| Some(*s)).collect();
        self.query_params(sql, &opts)
    }

    /// 参数化执行（便捷版，所有参数非 NULL）
    pub fn execute_str(
        &mut self,
        sql: &str,
        params: &[&str],
    ) -> Result<SuccessMessage, PgsqlError> {
        let opts: Vec<Option<&str>> = params.iter().map(|s| Some(*s)).collect();
        self.execute_params(sql, &opts)
    }
}

impl Drop for Connect {
    fn drop(&mut self) {
        let _ = self.stream.as_ref().write_all(&Packet::pack_terminate());
        let _ = self.stream.shutdown(std::net::Shutdown::Both);
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::net::TcpListener;
    use std::thread;

    // ── wire-protocol helpers ──────────────────────────────────────────

    /// Build a single PG backend message: type_byte | len(4) | payload
    fn pg_msg(tag: u8, payload: &[u8]) -> Vec<u8> {
        let mut m = Vec::with_capacity(5 + payload.len());
        m.push(tag);
        m.extend(&((payload.len() as u32 + 4).to_be_bytes()));
        m.extend_from_slice(payload);
        m
    }

    /// Build an Authentication message (tag 'R') with given auth_type + extra payload
    fn pg_auth(auth_type: u32, extra: &[u8]) -> Vec<u8> {
        let mut body = Vec::new();
        body.extend(&auth_type.to_be_bytes());
        body.extend_from_slice(extra);
        pg_msg(b'R', &body)
    }

    /// AuthenticationOk (R, type=0)
    fn auth_ok() -> Vec<u8> {
        pg_auth(0, &[])
    }

    /// ParameterStatus for server_version=15.0
    fn param_status() -> Vec<u8> {
        pg_msg(b'S', b"server_version\x0015.0\x00")
    }

    /// BackendKeyData (process_id=1, secret_key=2)
    fn backend_key() -> Vec<u8> {
        let mut p = Vec::new();
        p.extend(&1u32.to_be_bytes());
        p.extend(&2u32.to_be_bytes());
        pg_msg(b'K', &p)
    }

    /// ReadyForQuery (status = Idle)
    fn ready_for_query() -> Vec<u8> {
        pg_msg(b'Z', b"I")
    }

    /// The standard tail sent after auth succeeds: AuthOk + param + key + ready
    fn post_auth_ok() -> Vec<u8> {
        let mut v = Vec::new();
        v.extend(auth_ok());
        v.extend(param_status());
        v.extend(backend_key());
        v.extend(ready_for_query());
        v
    }

    /// Build a simple query response: ParseComplete + BindComplete +
    /// RowDescription(1 int4 col "c") + DataRow("1") + CommandComplete("SELECT 1") + ReadyForQuery
    fn simple_query_response() -> Vec<u8> {
        let mut r = Vec::new();
        // ParseComplete
        r.extend(pg_msg(b'1', &[]));
        // BindComplete
        r.extend(pg_msg(b'2', &[]));
        // RowDescription – 1 field "c", type_oid=23 (int4)
        let mut rd = Vec::new();
        rd.extend(&1u16.to_be_bytes()); // field count
        rd.extend(b"c\x00"); // name
        rd.extend(&0u32.to_be_bytes()); // table oid
        rd.extend(&1u16.to_be_bytes()); // column index
        rd.extend(&23u32.to_be_bytes()); // type oid (int4)
        rd.extend(&4i16.to_be_bytes()); // column length
        rd.extend(&(-1i32).to_be_bytes()); // type modifier
        rd.extend(&0u16.to_be_bytes()); // format (text)
        r.extend(pg_msg(b'T', &rd));
        // DataRow – 1 field, value "1"
        let mut dr = Vec::new();
        dr.extend(&1u16.to_be_bytes());
        dr.extend(&1u32.to_be_bytes()); // length of value
        dr.push(b'1');
        r.extend(pg_msg(b'D', &dr));
        // CommandComplete
        r.extend(pg_msg(b'C', b"SELECT 1\x00"));
        // ReadyForQuery
        r.extend(ready_for_query());
        r
    }

    /// Build an execute response (no rows): ParseComplete + BindComplete +
    /// NoData + CommandComplete("UPDATE 3") + ReadyForQuery
    fn execute_response() -> Vec<u8> {
        let mut r = Vec::new();
        r.extend(pg_msg(b'1', &[]));
        r.extend(pg_msg(b'2', &[]));
        r.extend(pg_msg(b'n', &[])); // NoData
        r.extend(pg_msg(b'C', b"UPDATE 3\x00"));
        r.extend(ready_for_query());
        r
    }

    /// Build a parameterized query response: ParseComplete + ParameterDescription + BindComplete +
    /// RowDescription(1 int4 col "p") + DataRow("42") + CommandComplete("SELECT 1") + ReadyForQuery
    fn query_params_response() -> Vec<u8> {
        let mut r = Vec::new();
        r.extend(pg_msg(b'1', &[]));

        let mut pd = Vec::new();
        pd.extend(&1u16.to_be_bytes());
        pd.extend(&23u32.to_be_bytes());
        r.extend(pg_msg(b't', &pd));

        r.extend(pg_msg(b'2', &[]));

        let mut rd = Vec::new();
        rd.extend(&1u16.to_be_bytes());
        rd.extend(b"p\x00");
        rd.extend(&0u32.to_be_bytes());
        rd.extend(&1u16.to_be_bytes());
        rd.extend(&23u32.to_be_bytes());
        rd.extend(&4i16.to_be_bytes());
        rd.extend(&(-1i32).to_be_bytes());
        rd.extend(&0u16.to_be_bytes());
        r.extend(pg_msg(b'T', &rd));

        let mut dr = Vec::new();
        dr.extend(&1u16.to_be_bytes());
        dr.extend(&2u32.to_be_bytes());
        dr.extend(b"42");
        r.extend(pg_msg(b'D', &dr));

        r.extend(pg_msg(b'C', b"SELECT 1\x00"));
        r.extend(ready_for_query());
        r
    }

    /// Build a parameterized execute response: ParseComplete + ParameterDescription + BindComplete +
    /// NoData + CommandComplete("UPDATE 1") + ReadyForQuery
    fn execute_params_response() -> Vec<u8> {
        let mut r = Vec::new();
        r.extend(pg_msg(b'1', &[]));

        let mut pd = Vec::new();
        pd.extend(&1u16.to_be_bytes());
        pd.extend(&23u32.to_be_bytes());
        r.extend(pg_msg(b't', &pd));

        r.extend(pg_msg(b'2', &[]));
        r.extend(pg_msg(b'n', &[]));
        r.extend(pg_msg(b'C', b"UPDATE 1\x00"));
        r.extend(ready_for_query());
        r
    }

    /// Build a parameterized query response with NULL row value.
    fn query_params_null_response() -> Vec<u8> {
        let mut r = Vec::new();
        r.extend(pg_msg(b'1', &[]));

        let mut pd = Vec::new();
        pd.extend(&1u16.to_be_bytes());
        pd.extend(&25u32.to_be_bytes());
        r.extend(pg_msg(b't', &pd));

        r.extend(pg_msg(b'2', &[]));

        let mut rd = Vec::new();
        rd.extend(&1u16.to_be_bytes());
        rd.extend(b"n\x00");
        rd.extend(&0u32.to_be_bytes());
        rd.extend(&1u16.to_be_bytes());
        rd.extend(&25u32.to_be_bytes());
        rd.extend(&(-1i16).to_be_bytes());
        rd.extend(&(-1i32).to_be_bytes());
        rd.extend(&0u16.to_be_bytes());
        r.extend(pg_msg(b'T', &rd));

        let mut dr = Vec::new();
        dr.extend(&1u16.to_be_bytes());
        dr.extend(&(-1i32).to_be_bytes());
        r.extend(pg_msg(b'D', &dr));

        r.extend(pg_msg(b'C', b"SELECT 1\x00"));
        r.extend(ready_for_query());
        r
    }

    /// Build an ErrorResponse for query phase
    fn error_response() -> Vec<u8> {
        let mut payload = Vec::new();
        payload.push(b'C');
        payload.extend(b"42601\x00");
        payload.push(b'M');
        payload.extend(b"syntax error\x00");
        payload.push(0);
        let mut r = Vec::new();
        r.extend(pg_msg(b'1', &[]));
        r.extend(pg_msg(b'2', &[]));
        r.extend(pg_msg(b'E', &payload));
        r.extend(ready_for_query());
        r
    }

    // ── mock server spawners ───────────────────────────────────────────

    /// Config pointing at 127.0.0.1:<port>
    fn mock_config(port: u16) -> Config {
        Config {
            debug: false,
            hostname: "127.0.0.1".into(),
            hostport: port as i32,
            username: "u".into(),
            userpass: "p".into(),
            database: "d".into(),
            charset: "utf8".into(),
            pool_max: 5,
        }
    }

    /// Spawn a mock PG server that does **cleartext** auth.
    /// Returns the port.  The server handles one connection.
    fn spawn_cleartext_server() -> u16 {
        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let port = listener.local_addr().unwrap().port();
        thread::spawn(move || {
            let (mut s, _) = listener.accept().unwrap();
            let mut buf = [0u8; 4096];
            // 1. read startup
            let _ = s.read(&mut buf).unwrap();
            // 2. send CleartextPassword request (auth_type=3)
            let _ = s.write_all(&pg_auth(3, &[]));
            // 3. read password message
            let _ = s.read(&mut buf).unwrap();
            // 4. send AuthOk + params + ready
            let _ = s.write_all(&post_auth_ok());
            // keep connection alive for queries
            loop {
                match s.read(&mut buf) {
                    Ok(0) | Err(_) => break,
                    Ok(_) => {
                        let _ = s.write_all(&simple_query_response());
                    }
                }
            }
        });
        thread::sleep(Duration::from_millis(30));
        port
    }

    /// Spawn a mock PG server that does **MD5** auth.
    fn spawn_md5_server() -> u16 {
        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let port = listener.local_addr().unwrap().port();
        thread::spawn(move || {
            let (mut s, _) = listener.accept().unwrap();
            let mut buf = [0u8; 4096];
            // 1. read startup
            let _ = s.read(&mut buf).unwrap();
            // 2. send MD5Password request (auth_type=5) + 4-byte salt
            let _ = s.write_all(&pg_auth(5, &[0xAA, 0xBB, 0xCC, 0xDD]));
            // 3. read md5 password
            let _ = s.read(&mut buf).unwrap();
            // 4. send AuthOk + params + ready
            let _ = s.write_all(&post_auth_ok());
            loop {
                match s.read(&mut buf) {
                    Ok(0) | Err(_) => break,
                    Ok(_) => {
                        let _ = s.write_all(&simple_query_response());
                    }
                }
            }
        });
        thread::sleep(Duration::from_millis(30));
        port
    }

    /// Spawn a mock PG server that does **SCRAM-SHA-256** auth.
    fn spawn_scram_server() -> u16 {
        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let port = listener.local_addr().unwrap().port();
        thread::spawn(move || {
            let (mut s, _) = listener.accept().unwrap();
            let mut buf = [0u8; 4096];
            // 1. read startup
            let _ = s.read(&mut buf).unwrap();
            // 2. send SASL auth (auth_type=10, mechanism)
            let _ = s.write_all(&pg_auth(10, b"SCRAM-SHA-256\x00\x00"));
            // 3. read SASLInitialResponse – extract client nonce
            let n = s.read(&mut buf).unwrap();
            let payload = &buf[..n];
            // find "n=,r=" in the payload to extract client nonce
            let text = String::from_utf8_lossy(payload);
            let client_nonce = text.split("r=").nth(1).unwrap_or("clientnonce").to_string();
            // 4. send SCRAM challenge (auth_type=11)
            let challenge = format!("r={client_nonce}SERVERNONCE,s=c2FsdA==,i=4096");
            let _ = s.write_all(&pg_auth(11, challenge.as_bytes()));
            // 5. read SCRAM client final
            let _ = s.read(&mut buf).unwrap();
            // 6. send SCRAM complete (auth_type=12) + AuthOk + params + ready
            let mut resp = Vec::new();
            resp.extend(pg_auth(12, b"v=dummyproof"));
            resp.extend(auth_ok());
            resp.extend(param_status());
            resp.extend(backend_key());
            resp.extend(ready_for_query());
            let _ = s.write_all(&resp);
            loop {
                match s.read(&mut buf) {
                    Ok(0) | Err(_) => break,
                    Ok(_) => {
                        let _ = s.write_all(&simple_query_response());
                    }
                }
            }
        });
        thread::sleep(Duration::from_millis(30));
        port
    }

    /// Spawn a server that accepts connection then immediately closes (EOF).
    fn spawn_eof_server() -> u16 {
        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let port = listener.local_addr().unwrap().port();
        thread::spawn(move || {
            let (s, _) = listener.accept().unwrap();
            drop(s); // close immediately
        });
        thread::sleep(Duration::from_millis(30));
        port
    }

    /// Spawn a server that sends an ErrorResponse after startup.
    fn spawn_auth_error_server() -> u16 {
        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let port = listener.local_addr().unwrap().port();
        thread::spawn(move || {
            let (mut s, _) = listener.accept().unwrap();
            let mut buf = [0u8; 4096];
            let _ = s.read(&mut buf).unwrap();
            // Send ErrorResponse
            let mut payload = Vec::new();
            payload.push(b'C');
            payload.extend(b"28P01\x00");
            payload.push(b'M');
            payload.extend(b"password authentication failed\x00");
            payload.push(0);
            let _ = s.write_all(&pg_msg(b'E', &payload));
        });
        thread::sleep(Duration::from_millis(30));
        port
    }

    /// Spawn a cleartext server that responds with error on query.
    fn spawn_query_error_server() -> u16 {
        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let port = listener.local_addr().unwrap().port();
        thread::spawn(move || {
            let (mut s, _) = listener.accept().unwrap();
            let mut buf = [0u8; 4096];
            // auth
            let _ = s.read(&mut buf).unwrap();
            let _ = s.write_all(&pg_auth(3, &[]));
            let _ = s.read(&mut buf).unwrap();
            let _ = s.write_all(&post_auth_ok());
            // query → error
            loop {
                match s.read(&mut buf) {
                    Ok(0) | Err(_) => break,
                    Ok(_) => {
                        let _ = s.write_all(&error_response());
                    }
                }
            }
        });
        thread::sleep(Duration::from_millis(30));
        port
    }

    /// Spawn a cleartext server that responds with execute response.
    fn spawn_execute_server() -> u16 {
        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let port = listener.local_addr().unwrap().port();
        thread::spawn(move || {
            let (mut s, _) = listener.accept().unwrap();
            let mut buf = [0u8; 4096];
            // auth
            let _ = s.read(&mut buf).unwrap();
            let _ = s.write_all(&pg_auth(3, &[]));
            let _ = s.read(&mut buf).unwrap();
            let _ = s.write_all(&post_auth_ok());
            // queries
            loop {
                match s.read(&mut buf) {
                    Ok(0) | Err(_) => break,
                    Ok(_) => {
                        let _ = s.write_all(&execute_response());
                    }
                }
            }
        });
        thread::sleep(Duration::from_millis(30));
        port
    }

    /// Spawn a cleartext server that responds with parameterized query response.
    fn spawn_query_params_server() -> u16 {
        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let port = listener.local_addr().unwrap().port();
        thread::spawn(move || {
            let (mut s, _) = listener.accept().unwrap();
            let mut buf = [0u8; 4096];
            let _ = s.read(&mut buf).unwrap();
            let _ = s.write_all(&pg_auth(3, &[]));
            let _ = s.read(&mut buf).unwrap();
            let _ = s.write_all(&post_auth_ok());
            loop {
                match s.read(&mut buf) {
                    Ok(0) | Err(_) => break,
                    Ok(_) => {
                        let _ = s.write_all(&query_params_response());
                    }
                }
            }
        });
        thread::sleep(Duration::from_millis(30));
        port
    }

    fn spawn_params_server() -> u16 {
        spawn_query_params_server()
    }

    /// Spawn a cleartext server that responds with parameterized execute response.
    fn spawn_execute_params_server() -> u16 {
        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let port = listener.local_addr().unwrap().port();
        thread::spawn(move || {
            let (mut s, _) = listener.accept().unwrap();
            let mut buf = [0u8; 4096];
            let _ = s.read(&mut buf).unwrap();
            let _ = s.write_all(&pg_auth(3, &[]));
            let _ = s.read(&mut buf).unwrap();
            let _ = s.write_all(&post_auth_ok());
            loop {
                match s.read(&mut buf) {
                    Ok(0) | Err(_) => break,
                    Ok(_) => {
                        let _ = s.write_all(&execute_params_response());
                    }
                }
            }
        });
        thread::sleep(Duration::from_millis(30));
        port
    }

    /// Spawn a cleartext server that returns NULL in parameterized query result.
    fn spawn_query_params_null_server() -> u16 {
        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let port = listener.local_addr().unwrap().port();
        thread::spawn(move || {
            let (mut s, _) = listener.accept().unwrap();
            let mut buf = [0u8; 4096];
            let _ = s.read(&mut buf).unwrap();
            let _ = s.write_all(&pg_auth(3, &[]));
            let _ = s.read(&mut buf).unwrap();
            let _ = s.write_all(&post_auth_ok());
            loop {
                match s.read(&mut buf) {
                    Ok(0) | Err(_) => break,
                    Ok(_) => {
                        let _ = s.write_all(&query_params_null_response());
                    }
                }
            }
        });
        thread::sleep(Duration::from_millis(30));
        port
    }

    // ── tests ──────────────────────────────────────────────────────────

    #[test]
    fn connect_cleartext_auth_success() {
        let port = spawn_cleartext_server();
        let conn = Connect::new(mock_config(port));
        assert!(conn.is_ok());
    }

    #[test]
    fn connect_md5_auth_success() {
        let port = spawn_md5_server();
        let conn = Connect::new(mock_config(port));
        assert!(conn.is_ok());
    }

    #[test]
    fn connect_scram_auth_success() {
        let port = spawn_scram_server();
        let conn = Connect::new(mock_config(port));
        assert!(conn.is_ok());
    }

    #[test]
    fn connect_connection_refused() {
        // port 1 is almost certainly not listening
        let cfg = mock_config(1);
        let result = Connect::new(cfg);
        assert!(result.is_err());
        match result.unwrap_err() {
            PgsqlError::Connection(_) => {}
            other => panic!("expected Connection error, got {other:?}"),
        }
    }

    #[test]
    fn connect_server_closes_immediately() {
        let port = spawn_eof_server();
        let result = Connect::new(mock_config(port));
        assert!(result.is_err());
    }

    #[test]
    fn connect_auth_error_from_server() {
        let port = spawn_auth_error_server();
        let result = Connect::new(mock_config(port));
        assert!(result.is_err());
    }

    #[test]
    fn connect_query_success() {
        let port = spawn_cleartext_server();
        let mut conn = Connect::new(mock_config(port)).unwrap();
        let result = conn.query("SELECT 1");
        assert!(result.is_ok());
        let msg = result.unwrap();
        assert_eq!(msg.rows.len(), 1);
        assert_eq!(msg.rows[0]["c"].as_i32(), Some(1));
    }

    #[test]
    fn connect_execute_success() {
        let port = spawn_execute_server();
        let mut conn = Connect::new(mock_config(port)).unwrap();
        let result = conn.execute("UPDATE t SET x=1");
        assert!(result.is_ok());
        let msg = result.unwrap();
        assert_eq!(msg.affect_count, 3);
        assert_eq!(msg.tag, "UPDATE 3");
    }

    #[test]
    fn connect_query_params_success() {
        let port = spawn_query_params_server();
        let mut conn = Connect::new(mock_config(port)).unwrap();
        let result = conn.query_params("SELECT $1::int", &[Some("42")]);
        assert!(result.is_ok());
        let msg = result.unwrap();
        assert_eq!(msg.rows.len(), 1);
        assert_eq!(msg.rows[0]["p"].as_i32(), Some(42));
    }

    #[test]
    fn connect_execute_params_success() {
        let port = spawn_execute_params_server();
        let mut conn = Connect::new(mock_config(port)).unwrap();
        let result = conn.execute_params("UPDATE t SET x=$1", &[Some("42")]);
        assert!(result.is_ok());
        let msg = result.unwrap();
        assert_eq!(msg.affect_count, 1);
        assert_eq!(msg.tag, "UPDATE 1");
    }

    #[test]
    fn connect_query_str_success() {
        let port = spawn_params_server();
        let mut conn = Connect::new(mock_config(port)).unwrap();
        let result = conn.query_str("SELECT $1::int", &["42"]);
        assert!(result.is_ok());
        let msg = result.unwrap();
        assert_eq!(msg.rows.len(), 1);
    }

    #[test]
    fn connect_execute_str_success() {
        let port = spawn_execute_params_server();
        let mut conn = Connect::new(mock_config(port)).unwrap();
        let result = conn.execute_str("UPDATE t SET x=$1", &["1"]);
        assert!(result.is_ok());
        let msg = result.unwrap();
        assert_eq!(msg.affect_count, 1);
    }

    #[test]
    fn connect_query_params_with_null() {
        let port = spawn_query_params_null_server();
        let mut conn = Connect::new(mock_config(port)).unwrap();
        let result = conn.query_params("SELECT $1::text", &[None]);
        assert!(result.is_ok());
        let msg = result.unwrap();
        assert_eq!(msg.rows.len(), 1);
        assert!(msg.rows[0]["n"].is_null());
    }

    #[test]
    fn connect_query_params_empty_string_vs_null() {
        let port = spawn_params_server();
        let mut conn = Connect::new(mock_config(port)).unwrap();

        // 空字符串参数
        let r1 = conn.query_params("SELECT $1::text", &[Some("")]);
        assert!(r1.is_ok());

        // NULL 参数
        let r2 = conn.query_params("SELECT $1::text", &[None]);
        assert!(r2.is_ok());
    }

    #[test]
    fn connect_query_returns_error() {
        let port = spawn_query_error_server();
        let mut conn = Connect::new(mock_config(port)).unwrap();
        let result = conn.query("BAD SQL");
        assert!(result.is_err());
    }

    #[test]
    fn connect_is_valid_true() {
        let port = spawn_cleartext_server();
        let mut conn = Connect::new(mock_config(port)).unwrap();
        assert!(conn.is_valid());
    }

    #[test]
    fn connect_is_valid_false_after_close() {
        let port = spawn_cleartext_server();
        let mut conn = Connect::new(mock_config(port)).unwrap();
        conn._close();
        // After closing, is_valid should return false
        assert!(!conn.is_valid());
    }

    #[test]
    fn connect_close_does_not_panic() {
        let port = spawn_cleartext_server();
        let mut conn = Connect::new(mock_config(port)).unwrap();
        conn._close();
        // calling close again should not panic
        conn._close();
    }

    #[test]
    fn connect_drop_does_not_panic() {
        let port = spawn_cleartext_server();
        let conn = Connect::new(mock_config(port)).unwrap();
        drop(conn);
    }

    fn spawn_transaction_status_server() -> u16 {
        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let port = listener.local_addr().unwrap().port();
        thread::spawn(move || {
            let (mut s, _) = listener.accept().unwrap();
            let mut buf = [0u8; 4096];
            let _ = s.read(&mut buf).unwrap();
            let _ = s.write_all(&pg_auth(3, &[]));
            let _ = s.read(&mut buf).unwrap();
            let _ = s.write_all(&post_auth_ok());
            loop {
                match s.read(&mut buf) {
                    Ok(0) | Err(_) => break,
                    Ok(_) => {
                        let mut r = Vec::new();
                        r.extend(pg_msg(b'1', &[]));
                        r.extend(pg_msg(b'2', &[]));
                        let mut rd = Vec::new();
                        rd.extend(&1u16.to_be_bytes());
                        rd.extend(b"c\x00");
                        rd.extend(&0u32.to_be_bytes());
                        rd.extend(&1u16.to_be_bytes());
                        rd.extend(&23u32.to_be_bytes());
                        rd.extend(&4i16.to_be_bytes());
                        rd.extend(&(-1i32).to_be_bytes());
                        rd.extend(&0u16.to_be_bytes());
                        r.extend(pg_msg(b'T', &rd));
                        let mut dr = Vec::new();
                        dr.extend(&1u16.to_be_bytes());
                        dr.extend(&1u32.to_be_bytes());
                        dr.push(b'1');
                        r.extend(pg_msg(b'D', &dr));
                        r.extend(pg_msg(b'C', b"SELECT 1\x00"));
                        r.extend(pg_msg(b'Z', b"T"));
                        let _ = s.write_all(&r);
                    }
                }
            }
        });
        thread::sleep(Duration::from_millis(30));
        port
    }

    fn spawn_error_status_server() -> u16 {
        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let port = listener.local_addr().unwrap().port();
        thread::spawn(move || {
            let (mut s, _) = listener.accept().unwrap();
            let mut buf = [0u8; 4096];
            let _ = s.read(&mut buf).unwrap();
            let _ = s.write_all(&pg_auth(3, &[]));
            let _ = s.read(&mut buf).unwrap();
            let _ = s.write_all(&post_auth_ok());
            loop {
                match s.read(&mut buf) {
                    Ok(0) | Err(_) => break,
                    Ok(_) => {
                        let mut r = Vec::new();
                        r.extend(pg_msg(b'1', &[]));
                        r.extend(pg_msg(b'2', &[]));
                        let mut rd = Vec::new();
                        rd.extend(&1u16.to_be_bytes());
                        rd.extend(b"c\x00");
                        rd.extend(&0u32.to_be_bytes());
                        rd.extend(&1u16.to_be_bytes());
                        rd.extend(&23u32.to_be_bytes());
                        rd.extend(&4i16.to_be_bytes());
                        rd.extend(&(-1i32).to_be_bytes());
                        rd.extend(&0u16.to_be_bytes());
                        r.extend(pg_msg(b'T', &rd));
                        let mut dr = Vec::new();
                        dr.extend(&1u16.to_be_bytes());
                        dr.extend(&1u32.to_be_bytes());
                        dr.push(b'1');
                        r.extend(pg_msg(b'D', &dr));
                        r.extend(pg_msg(b'C', b"SELECT 1\x00"));
                        r.extend(pg_msg(b'Z', b"E"));
                        let _ = s.write_all(&r);
                    }
                }
            }
        });
        thread::sleep(Duration::from_millis(30));
        port
    }

    fn spawn_slow_partial_server() -> u16 {
        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let port = listener.local_addr().unwrap().port();
        thread::spawn(move || {
            let (mut s, _) = listener.accept().unwrap();
            let mut buf = [0u8; 4096];
            let _ = s.read(&mut buf).unwrap();
            let _ = s.write_all(&pg_auth(3, &[]));
            let _ = s.read(&mut buf).unwrap();
            let _ = s.write_all(&post_auth_ok());
            match s.read(&mut buf) {
                Ok(0) | Err(_) => {}
                Ok(_) => {
                    let _ = s.write_all(&simple_query_response());
                }
            }
        });
        thread::sleep(Duration::from_millis(30));
        port
    }

    fn spawn_rst_on_query_server() -> u16 {
        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let port = listener.local_addr().unwrap().port();
        thread::spawn(move || {
            let (mut s, _) = listener.accept().unwrap();
            let mut buf = [0u8; 4096];
            let _ = s.read(&mut buf).unwrap();
            let _ = s.write_all(&pg_auth(3, &[]));
            let _ = s.read(&mut buf).unwrap();
            let _ = s.write_all(&post_auth_ok());
            match s.read(&mut buf) {
                Ok(0) | Err(_) => {}
                Ok(_) => {
                    drop(s);
                }
            }
        });
        thread::sleep(Duration::from_millis(30));
        port
    }

    #[test]
    fn connect_query_ready_for_query_transaction_status() {
        let port = spawn_transaction_status_server();
        let mut conn = Connect::new(mock_config(port)).unwrap();
        let result = conn.query("SELECT 1");
        assert!(result.is_ok());
    }

    #[test]
    fn connect_query_ready_for_query_error_status() {
        let port = spawn_error_status_server();
        let mut conn = Connect::new(mock_config(port)).unwrap();
        let result = conn.query("SELECT 1");
        assert!(result.is_ok());
    }

    #[test]
    fn connect_query_server_closes_after_partial() {
        let port = spawn_slow_partial_server();
        let mut conn = Connect::new(mock_config(port)).unwrap();
        let r1 = conn.query("SELECT 1");
        assert!(r1.is_ok());
        let r2 = conn.query("SELECT 1");
        assert!(r2.is_err());
    }

    #[test]
    fn connect_query_server_rst_returns_io_or_connection_error() {
        let port = spawn_rst_on_query_server();
        let mut conn = Connect::new(mock_config(port)).unwrap();
        let result = conn.query("SELECT 1");
        assert!(result.is_err());
    }

    #[test]
    fn connect_read_would_block_max_retries() {
        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let port = listener.local_addr().unwrap().port();
        thread::spawn(move || {
            let (mut s, _) = listener.accept().unwrap();
            let mut buf = [0u8; 4096];
            let _ = s.read(&mut buf);
            let _ = s.write_all(&pg_auth(3, &[]));
            let _ = s.read(&mut buf);
            let _ = s.write_all(&post_auth_ok());
            let _ = s.read(&mut buf);
            thread::sleep(Duration::from_secs(5));
        });
        thread::sleep(Duration::from_millis(30));

        let mut conn = Connect::new(mock_config(port)).unwrap();
        conn.stream
            .set_read_timeout(Some(Duration::from_millis(1)))
            .ok();
        let result = conn.query("SELECT 1");
        assert!(result.is_err());
        let err_str = result.unwrap_err().to_string();
        assert!(
            err_str.contains("超时") || err_str.contains("Timeout") || err_str.contains("重试"),
            "expected timeout error, got: {err_str}"
        );
    }

    #[test]
    fn connect_read_exceeds_max_message_size() {
        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let port = listener.local_addr().unwrap().port();
        thread::spawn(move || {
            let (mut s, _) = listener.accept().unwrap();
            let mut buf = [0u8; 4096];
            let _ = s.read(&mut buf);
            let _ = s.write_all(&pg_auth(3, &[]));
            let _ = s.read(&mut buf);
            let _ = s.write_all(&post_auth_ok());
            let _ = s.read(&mut buf);
            let big = vec![b'X'; 256];
            let _ = s.write_all(&big);
            thread::sleep(Duration::from_secs(2));
        });
        thread::sleep(Duration::from_millis(30));

        let mut conn = Connect::new(mock_config(port)).unwrap();
        let result = conn.query("SELECT 1");
        assert!(result.is_err());
        let err_str = result.unwrap_err().to_string();
        assert!(
            err_str.contains("最大") || err_str.contains("大小") || err_str.contains("size"),
            "expected max message size error, got: {err_str}"
        );
    }

    #[test]
    fn connect_read_deadline_timeout() {
        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let port = listener.local_addr().unwrap().port();
        thread::spawn(move || {
            let (mut s, _) = listener.accept().unwrap();
            let mut buf = [0u8; 4096];
            let _ = s.read(&mut buf);
            let _ = s.write_all(&pg_auth(3, &[]));
            let _ = s.read(&mut buf);
            let _ = s.write_all(&post_auth_ok());
            let _ = s.read(&mut buf);
            for _ in 0..200 {
                let _ = s.write_all(b"X");
                thread::sleep(Duration::from_millis(5));
            }
        });
        thread::sleep(Duration::from_millis(30));

        let mut conn = Connect::new(mock_config(port)).unwrap();
        let result = conn.query("SELECT 1");
        assert!(result.is_err());
    }

    #[test]
    fn connect_read_partial_auth_frame() {
        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let port = listener.local_addr().unwrap().port();
        thread::spawn(move || {
            let (mut s, _) = listener.accept().unwrap();
            let mut buf = [0u8; 4096];
            let _ = s.read(&mut buf);
            let auth = pg_auth(3, &[]);
            let _ = s.write_all(&auth[..5]);
            thread::sleep(Duration::from_millis(50));
            let _ = s.write_all(&auth[5..]);
            let _ = s.read(&mut buf);
            let _ = s.write_all(&post_auth_ok());
            loop {
                match s.read(&mut buf) {
                    Ok(0) | Err(_) => break,
                    Ok(_) => {
                        let _ = s.write_all(&simple_query_response());
                    }
                }
            }
        });
        thread::sleep(Duration::from_millis(30));

        let mut conn = Connect::new(mock_config(port)).unwrap();
        let result = conn.query("SELECT 1");
        assert!(result.is_ok());
    }
}