写一个迷你 Postgres 客户端

// Append one message to the write buffer; nothing is sent until flush().
// `ty: None` is the StartupMessage, the only message without a type byte.
pub fn put(&mut self, ty: Option<u8>, body: impl FnOnce(&mut BytesMut)) {
  if let Some(ty) = ty {
    self.write_buf.put_u8(ty);
  }
  let at = self.write_buf.len();
  self.write_buf.put_u32(0);
  body(&mut self.write_buf);
  let len = (self.write_buf.len() - at) as u32;
  self.write_buf[at..at + 4].copy_from_slice(&len.to_be_bytes());
}

pub async fn recv(&mut self) -> Result<(u8, Bytes), Error> {
  loop {
    if let [ty, a, b, c, d, ..] = self.read_buf[..] {
      let len = u32::from_be_bytes([a, b, c, d]) as usize;
      // ... len < 4 报协议错误
      if self.read_buf.len() > len {
        let mut msg = self.read_buf.split_to(1 + len).freeze();
        msg.advance(5);
        return Ok((ty, msg));
      }
    }
    if self.stream.read_buf(&mut self.read_buf).await? == 0 {
      return protocol("server closed the connection");
    }
  }
}

wire.put(None, |b| {
  b.put_i32(196_608); // protocol 3.0
  for (k, v) in [("user", &cfg.user), ("database", &cfg.database)] {
    put_cstr(b, k);
    put_cstr(b, v);
  }
  for (k, v) in &cfg.params {
    put_cstr(b, k);
    put_cstr(b, v);
  }
  b.put_u8(0);
});
wire.flush().await?;

let mut backend_pid = 0;
loop {
  let (ty, mut body) = wire.recv().await?;
  match ty {
    b'R' => match body.try_get_i32()? {
      0 => {} // AuthenticationOk
      10 => scram::authenticate(&mut wire, &cfg.user, &cfg.password, body).await?,
      n => return protocol(format!("unsupported authentication method {n}")),
    },
    b'K' => backend_pid = body.try_get_i32()?,
    b'S' | b'N' => {} // ParameterStatus, NoticeResponse
    b'E' => return Err(parse_error(body)),
    b'Z' => return Ok(Conn { wire, backend_pid, busy: false, tx_status: body.try_get_u8()? }),
    // ...
  }
}

// SaltedPassword := Hi(password, salt, i), Hi is PBKDF2-HMAC-SHA256
let mut salted = [0u8; 32];
pbkdf2::pbkdf2_hmac::<Sha256>(password.as_bytes(), &salt, iterations, &mut salted);
let client_key = hmac(&salted, b"Client Key")?;
let stored_key = Sha256::digest(client_key);
let without_proof = format!("c=biws,r={nonce}"); // biws = base64("n,,")
let auth_message = format!("{client_first_bare},{server_first},{without_proof}");
let client_signature = hmac(&stored_key, auth_message.as_bytes())?;
let proof: Vec<u8> = client_key.iter().zip(client_signature).map(|(k, s)| k ^ s).collect();
let client_final = format!("{without_proof},p={}", BASE64.encode(proof));
wire.put(Some(b'p'), |b| b.put_slice(client_final.as_bytes()));
wire.flush().await?;

// server-final-message: v=<ServerSignature>; checking it proves the server knows the password too
let server_final = recv_sasl(wire, AUTH_SASL_FINAL).await?;
let server_key = hmac(&salted, b"Server Key")?;
let expected = BASE64.encode(hmac(&server_key, auth_message.as_bytes())?);
if server_final.strip_prefix("v=") != Some(expected.as_str()) {
  return protocol("server signature mismatch");
}

debug
4096 rounds: 34.19925ms
4096 rounds: 35.401167ms
4096 rounds: 34.959834ms
release
4096 rounds: 1.451708ms
4096 rounds: 1.454417ms
4096 rounds: 1.446417ms

pub async fn simple_query(&mut self, sql: &str) -> Result<Vec<SimpleResult>, Error> {
  self.busy = true;
  self.wire.put(Some(b'Q'), |b| put_cstr(b, sql));
  self.wire.flush().await?;
  let mut results = Vec::new();
  let mut current = SimpleResult::default();
  let mut error = None;
  loop {
    let (ty, mut body) = self.wire.recv().await?;
    match ty {
      b'T' => current.columns = parse_row_description(body)?,
      b'D' => { /* 每列 i32 长度加字节，-1 是 NULL，按 UTF-8 转成 String */ }
      b'C' => {
        current.tag = get_cstr(&mut body)?;
        results.push(std::mem::take(&mut current));
      }
      b'I' => {} // EmptyQueryResponse
      b'E' => error = Some(parse_error(body)),
      b'S' | b'N' => {}
      b'Z' => {
        self.busy = false;
        self.tx_status = body.try_get_u8()?;
        return error.map_or(Ok(results), Err);
      }
      // ...
    }
  }
}

self.wire.put(Some(b'P'), |b| {
  put_cstr(b, ""); // unnamed statement
  put_cstr(b, sql);
  b.put_i16(1); // one parameter type
  b.put_u32(INT8_OID);
});
self.wire.put(Some(b'B'), |b| {
  put_cstr(b, ""); // unnamed portal
  put_cstr(b, ""); // unnamed statement
  b.put_i16(1); // one parameter format: binary
  b.put_i16(BINARY);
  b.put_i16(1); // one parameter value
  b.put_i32(8);
  b.put_i64(param);
  b.put_i16(1); // one result format, applies to all columns: binary
  b.put_i16(BINARY);
});
self.wire.put(Some(b'E'), |b| {
  put_cstr(b, "");
  b.put_i32(0); // no row limit
});
self.wire.put(Some(b'S'), |_| {});
self.wire.flush().await?;

match value.as_deref() {
  None => row.push(None),
  Some(&[a, b, c, d, e, f, g, h]) => row.push(Some(i64::from_be_bytes([a, b, c, d, e, f, g, h]))),
  // keep reading up to ReadyForQuery so the connection stays usable
  Some(v) => {
    let msg = format!("column {i}: {} bytes is not an int8", v.len());
    error.get_or_insert(Error::Decode(msg));
  }
}

pub fn parse_error(mut body: Bytes) -> Error {
  let (mut severity, mut code, mut message) = (String::new(), String::new(), String::new());
  while let Ok(field) = body.try_get_u8() {
    let Ok(value) = get_cstr(&mut body) else { break };
    match field {
      b'V' => severity = value,
      b'C' => code = value,
      b'M' => message = value,
      _ => {}
    }
  }
  Error::Db { severity, code, message }
}

pub async fn acquire(&self) -> Result<PooledConn, Error> {
  let shared = &self.0;
  let attempt = async {
    // Semaphore is FIFO: waiters are served in arrival order
    let permit = shared.permits.clone().acquire_owned().await.map_err(|_| Error::PoolTimedOut)?;
    let idle = shared.idle.lock().map(|mut idle| idle.pop()).unwrap_or(None);
    let conn = match idle {
      Some(conn) => conn,
      None => Conn::connect(&shared.config).await?,
    };
    Ok(PooledConn { conn: Some(conn), pool: shared.clone(), _permit: permit })
  };
  tokio::time::timeout(shared.acquire_timeout, attempt).await.map_err(|_| Error::PoolTimedOut)?
}

impl Drop for PooledConn {
  fn drop(&mut self) {
    // a busy connection or one inside a transaction is dropped (closing the socket);
    // the permit still goes back, so the next acquire opens a new one
    if let Some(conn) = self.conn.take().filter(|c| c.reusable())
      && let Ok(mut idle) = self.pool.idle.lock()
    {
      idle.push(conn);
    }
  }
}

        // NOTE: we need to make sure we drop the permit *after* we push to the idle queue
        // don't decrease the size
        guard.release_permit();

--- connect ---
  connected, backend pid 735, 39.628042ms
--- simple query: setup (4 statements in one Query) ---
  "CREATE SCHEMA" columns []
  "DROP TABLE" columns []
  "CREATE TABLE" columns []
  "INSERT 0 2" columns []
--- simple query: select ---
  "SELECT 2" columns [id:20, title:25, note:25]
    [Some("1"), Some("hello sqlx"), None]
    [Some("2"), Some("mini client"), Some("no sqlx")]
--- extended query: one int8 parameter, binary result ---
  [[Some(1), Some(10)], [Some(2), Some(20)]]
--- errors ---
  simple:   ERROR 42601: syntax error at or near "selec"
  extended: ERROR 22012: division by zero
  decode:   decode: column 0: 4 bytes is not an int8
  same connection afterwards: [[Some(42)]]
  wrong password: FATAL 28P01: password authentication failed for user "lab"
--- pool: max 2, acquire timeout 300ms ---
  holding pids 737 and 738
  third acquire: pool timed out after 302ms
  dropped pid 737; waiting acquire got pid 737 after 101ms
  idle now 2
--- pool: 20 tasks, max 2 ---
  queries per backend pid: {"737": 10, "738": 10}
  idle now 2

sqlx: select $1::int4 as i64 -> error occurred while decoding column 0: mismatched types; Rust type `i64` (as SQL type `INT8`) is not compatible with SQL type `INT4`
sqlx: column type from Describe: INT4
same simple columns [id:20, title:25, note:25]
same simple row [Some("1"), Some("hello sqlx"), None]
same simple row [Some("2"), Some("mini client"), Some("no sqlx")]
same extended [[Some(1), Some(10)], [Some(2), Some(20)]]
same error ERROR 42601: syntax error at or near "selec"
same error ERROR 22012: division by zero
same error FATAL 28P01: password authentication failed for user "lab"
lines: mini 7, sqlx 7

=== mini: connect ===
  [conn 1] C->S StartupMessage(v3.0)
  [conn 1] S->C Authentication
  [conn 1] C->S Password/SASL
  [conn 1] S->C Authentication
  [conn 1] C->S Password/SASL
  [conn 1] S->C Authentication, Authentication, ParameterStatus, ... BackendKeyData, ReadyForQuery(I)
=== mini: simple query ===
  [conn 1] C->S Query("select id, title from posts where id = 1")
  [conn 1] S->C RowDescription, DataRow(列长 1/10), CommandComplete("SELECT 1"), ReadyForQuery(I)
=== mini: extended query ===
  [conn 1] C->S Parse(unnamed, "select id, id * $1 from posts order by id"), Bind(portal unnamed, stmt unnamed), Execute, Sync
  [conn 1] S->C ParseComplete, BindComplete, DataRow(列长 8/8), DataRow(列长 8/8), CommandComplete("SELECT 2"), ReadyForQuery(I)
[[Some(1), Some(10)], [Some(2), Some(20)]]
=== sqlx: connect ===
  [conn 2] C->S StartupMessage(v3.0)
  [conn 2] S->C Authentication
  [conn 2] C->S Password/SASL
  [conn 2] S->C Authentication
  [conn 2] C->S Password/SASL
  [conn 2] S->C Authentication, Authentication, ParameterStatus, ... BackendKeyData, ReadyForQuery(I)
=== sqlx: query + bind, first ===
  [conn 2] C->S Parse(sqlx_s_1, "select id, id * $1 from posts order by id"), Describe(S, sqlx_s_1), Sync
  [conn 2] S->C ParseComplete, ParameterDescription, RowDescription, ReadyForQuery(I)
  [conn 2] C->S Bind(portal unnamed, stmt sqlx_s_1), Execute, Close(P, unnamed), Sync
  [conn 2] S->C BindComplete, DataRow(列长 8/8), DataRow(列长 8/8), CommandComplete("SELECT 2"), CloseComplete, ReadyForQuery(I)
[(1, 10), (2, 20)]
=== sqlx: query + bind, second ===
  [conn 2] C->S Bind(portal unnamed, stmt sqlx_s_1), Execute, Close(P, unnamed), Sync
  [conn 2] S->C BindComplete, DataRow(列长 8/8), DataRow(列长 8/8), CommandComplete("SELECT 2"), CloseComplete, ReadyForQuery(I)
[(1, 10), (2, 20)]
  [conn 2] C->S Terminate

pid 856: start a 200ms query, give up after 50ms
  timeout: true
idle after drop: 0
pid 857: select 41 + 1 -> [[Some(42)]]
pid 857: select 42 + 1 -> [[Some(43)]]
pid 857: begin, then drop without commit
idle after drop: 0
next acquire: [Some("858"), Some("t")]

pid 866: start a 200ms query, give up after 50ms
  timeout: true
idle after drop: 1
pid 866: select 41 + 1 -> [[Some(7)]]
pid 866: select 42 + 1 -> [[Some(42)]]
Error: Protocol("unexpected message '1'")

pid 810: start a 200ms query, give up after 50ms
  timeout: true
idle after drop: 1
pid 810: select 41 + 1 -> 42
pid 810: select 42 + 1 -> 43