#![forbid(unsafe_code)]
use std::pin::Pin;
use std::task::{Context, Poll};
use bytes::BytesMut;
use fallible_iterator_02::FallibleIterator;
use postgres_protocol::authentication::md5_hash;
use postgres_protocol::authentication::sasl::{ChannelBinding, ScramSha256, SCRAM_SHA_256};
use postgres_protocol::message::backend::{ErrorResponseBody, Message};
use postgres_protocol::message::frontend;
use rustls_pki_types::ServerName;
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, ReadBuf};
use tokio::net::TcpStream;
use tokio_rustls::client::TlsStream;
use tokio_rustls::TlsConnector;
use crate::builder::TlsMode;
use crate::connection::parse_pg_conn_str;
use crate::error::PgError;
#[derive(Debug)]
pub enum MaybeTlsStream {
Plain(TcpStream),
Tls(Box<TlsStream<TcpStream>>),
}
impl AsyncRead for MaybeTlsStream {
fn poll_read(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut ReadBuf<'_>,
) -> Poll<std::io::Result<()>> {
match self.get_mut() {
MaybeTlsStream::Plain(s) => Pin::new(s).poll_read(cx, buf),
MaybeTlsStream::Tls(s) => Pin::new(&mut **s).poll_read(cx, buf),
}
}
}
impl AsyncWrite for MaybeTlsStream {
fn poll_write(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<std::io::Result<usize>> {
match self.get_mut() {
MaybeTlsStream::Plain(s) => Pin::new(s).poll_write(cx, buf),
MaybeTlsStream::Tls(s) => Pin::new(&mut **s).poll_write(cx, buf),
}
}
fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
match self.get_mut() {
MaybeTlsStream::Plain(s) => Pin::new(s).poll_flush(cx),
MaybeTlsStream::Tls(s) => Pin::new(&mut **s).poll_flush(cx),
}
}
fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
match self.get_mut() {
MaybeTlsStream::Plain(s) => Pin::new(s).poll_shutdown(cx),
MaybeTlsStream::Tls(s) => Pin::new(&mut **s).poll_shutdown(cx),
}
}
}
pub async fn connect_replication(conn_str: &str, tls: TlsMode) -> Result<MaybeTlsStream, PgError> {
let parts = parse_pg_conn_str(conn_str)
.map_err(|e| PgError::Replication(format!("invalid replication connection string: {e}")))?;
let user = match parts.user.as_deref() {
Some(u) if !u.is_empty() => u,
_ => {
return Err(PgError::Replication(
"connection string must specify a non-empty user for a replication connection"
.to_string(),
));
}
};
let password = extract_password(conn_str);
let tcp = TcpStream::connect((parts.host.as_str(), parts.port)).await?;
tcp.set_nodelay(true)?;
let mut stream = upgrade_tls(tcp, &parts.host, tls).await?;
let params = build_startup_params(user, parts.dbname.as_deref());
let mut out = BytesMut::new();
frontend::startup_message(params, &mut out)
.map_err(|e| PgError::Protocol(format!("failed to build StartupMessage: {e}")))?;
stream.write_all(&out).await?;
let mut in_buf = BytesMut::new();
authenticate(&mut stream, &mut in_buf, user, password.as_deref()).await?;
drain_until_ready(&mut stream, &mut in_buf).await?;
Ok(stream)
}
async fn upgrade_tls(
mut tcp: TcpStream,
host: &str,
tls: TlsMode,
) -> Result<MaybeTlsStream, PgError> {
let cfg = match tls {
TlsMode::Disabled => return Ok(MaybeTlsStream::Plain(tcp)),
TlsMode::Rustls(cfg) => cfg,
};
let mut out = BytesMut::new();
frontend::ssl_request(&mut out);
tcp.write_all(&out).await?;
let resp = tcp.read_u8().await?;
classify_ssl_response(resp)?;
let domain = ServerName::try_from(host.to_string())
.map_err(|e| PgError::Tls(format!("invalid server name {host:?} for TLS SNI: {e}")))?;
let tls_stream = TlsConnector::from(cfg)
.connect(domain, tcp)
.await
.map_err(|e| PgError::Tls(format!("TLS handshake failed: {e}")))?;
Ok(MaybeTlsStream::Tls(Box::new(tls_stream)))
}
fn classify_ssl_response(byte: u8) -> Result<(), PgError> {
match byte {
b'S' => Ok(()),
b'N' => Err(PgError::Tls(
"server does not support TLS (replied 'N' to SSLRequest)".to_string(),
)),
other => Err(PgError::Protocol(format!(
"unexpected byte {other:#04x} in response to SSLRequest (expected 'S' or 'N')"
))),
}
}
async fn authenticate<S>(
stream: &mut S,
buf: &mut BytesMut,
user: &str,
password: Option<&str>,
) -> Result<(), PgError>
where
S: AsyncRead + AsyncWrite + Unpin,
{
loop {
match read_message(stream, buf).await? {
Message::AuthenticationOk => return Ok(()),
Message::AuthenticationCleartextPassword => {
let response = build_password_response(user, require_password(password)?, None);
send_password(stream, &response).await?;
}
Message::AuthenticationMd5Password(body) => {
let response =
build_password_response(user, require_password(password)?, Some(body.salt()));
send_password(stream, &response).await?;
}
Message::AuthenticationSasl(body) => {
let mechanisms: Vec<&str> = body.mechanisms().collect().map_err(|e| {
PgError::Protocol(format!("malformed AuthenticationSASL mechanism list: {e}"))
})?;
select_scram_mechanism(&mechanisms)?;
run_scram_exchange(stream, buf, require_password(password)?).await?;
}
Message::ErrorResponse(body) => return Err(server_error_response(&body)),
_ => {
return Err(PgError::Protocol(
"unexpected message during replication authentication".to_string(),
));
}
}
}
}
async fn run_scram_exchange<S>(
stream: &mut S,
buf: &mut BytesMut,
password: &str,
) -> Result<(), PgError>
where
S: AsyncRead + AsyncWrite + Unpin,
{
let mut scram = ScramSha256::new(password.as_bytes(), ChannelBinding::unsupported());
let mut out = BytesMut::new();
frontend::sasl_initial_response(SCRAM_SHA_256, scram.message(), &mut out)
.map_err(|e| PgError::Protocol(format!("failed to build SASLInitialResponse: {e}")))?;
stream.write_all(&out).await?;
match read_message(stream, buf).await? {
Message::AuthenticationSaslContinue(body) => scram.update(body.data()).map_err(|e| {
PgError::Replication(format!(
"SCRAM-SHA-256 exchange failed while processing the server's first message: {e}"
))
})?,
Message::ErrorResponse(body) => return Err(server_error_response(&body)),
_ => {
return Err(PgError::Protocol(
"expected AuthenticationSASLContinue after SASLInitialResponse".to_string(),
));
}
}
let mut out = BytesMut::new();
frontend::sasl_response(scram.message(), &mut out)
.map_err(|e| PgError::Protocol(format!("failed to build SASLResponse: {e}")))?;
stream.write_all(&out).await?;
match read_message(stream, buf).await? {
Message::AuthenticationSaslFinal(body) => scram.finish(body.data()).map_err(|e| {
PgError::Replication(format!(
"SCRAM-SHA-256 exchange failed while verifying the server's final message: {e}"
))
}),
Message::ErrorResponse(body) => Err(server_error_response(&body)),
_ => Err(PgError::Protocol(
"expected AuthenticationSASLFinal after SASLResponse".to_string(),
)),
}
}
async fn drain_until_ready<S>(stream: &mut S, buf: &mut BytesMut) -> Result<(), PgError>
where
S: AsyncRead + Unpin,
{
loop {
match read_message(stream, buf).await? {
Message::ReadyForQuery(_) => return Ok(()),
Message::BackendKeyData(_)
| Message::ParameterStatus(_)
| Message::NoticeResponse(_) => {}
Message::ErrorResponse(body) => return Err(server_error_response(&body)),
_ => {
return Err(PgError::Protocol(
"unexpected message while waiting for ReadyForQuery after authentication"
.to_string(),
));
}
}
}
}
pub(crate) fn server_error_response(body: &ErrorResponseBody) -> PgError {
let mut sqlstate: Option<String> = None;
let mut message: Option<String> = None;
let mut fields = body.fields();
loop {
match fields.next() {
Ok(Some(field)) => {
let value = String::from_utf8_lossy(field.value_bytes()).into_owned();
match field.type_() {
b'C' => sqlstate = Some(value),
b'M' => message = Some(value),
_ => {}
}
}
Ok(None) => break,
Err(e) => {
return PgError::Protocol(format!(
"malformed ErrorResponse field during replication connection setup: {e}"
));
}
}
}
match (sqlstate, message) {
(Some(sqlstate), Some(message)) => {
PgError::Replication(format!("server error [{sqlstate}]: {message}"))
}
(Some(sqlstate), None) => PgError::Replication(format!(
"server error [{sqlstate}] (ErrorResponse had no message ('M') field)"
)),
(None, Some(message)) => PgError::Replication(format!(
"server error: {message} (ErrorResponse had no SQLSTATE ('C') field)"
)),
(None, None) => PgError::Replication(
"server sent an ErrorResponse during replication connection startup or \
authentication with neither a SQLSTATE nor a message field"
.to_string(),
),
}
}
fn select_scram_mechanism<'a>(mechanisms: &[&'a str]) -> Result<&'a str, PgError> {
mechanisms
.iter()
.find(|&&mechanism| mechanism == SCRAM_SHA_256)
.copied()
.ok_or_else(|| {
PgError::Replication(format!(
"server does not support SCRAM-SHA-256 authentication; offered: {mechanisms:?}"
))
})
}
fn require_password(password: Option<&str>) -> Result<&str, PgError> {
password.ok_or_else(|| {
PgError::Replication(
"server requires a password but none was supplied in the connection string".to_string(),
)
})
}
async fn send_password<S>(stream: &mut S, response: &str) -> Result<(), PgError>
where
S: AsyncWrite + Unpin,
{
let mut out = BytesMut::new();
frontend::password_message(response.as_bytes(), &mut out)
.map_err(|e| PgError::Protocol(format!("failed to build PasswordMessage: {e}")))?;
stream.write_all(&out).await?;
Ok(())
}
pub(crate) async fn read_message<S>(stream: &mut S, buf: &mut BytesMut) -> Result<Message, PgError>
where
S: AsyncRead + Unpin,
{
loop {
if let Some(msg) = Message::parse(buf).map_err(|e| {
PgError::Protocol(format!(
"malformed backend message during replication connection setup: {e}"
))
})? {
return Ok(msg);
}
let n = stream.read_buf(buf).await?;
if n == 0 {
return Err(PgError::Connection(
"server closed the connection during replication connection setup".to_string(),
));
}
}
}
fn build_startup_params<'a>(user: &'a str, dbname: Option<&'a str>) -> Vec<(&'a str, &'a str)> {
let mut params = vec![("user", user), ("replication", "database")];
if let Some(db) = dbname {
params.push(("database", db));
}
params
}
fn build_password_response(user: &str, password: &str, salt: Option<[u8; 4]>) -> String {
match salt {
Some(salt) => md5_hash(user.as_bytes(), password.as_bytes(), salt),
None => password.to_string(),
}
}
fn extract_password(conn_str: &str) -> Option<String> {
let trimmed = conn_str.trim();
if let Some(rest) = trimmed
.strip_prefix("postgresql://")
.or_else(|| trimmed.strip_prefix("postgres://"))
{
let (rest, _query) = rest.split_once('?').unwrap_or((rest, ""));
let user_info = &rest[..rest.rfind('@')?];
let (_user, password) = user_info.split_once(':')?;
if password.is_empty() {
None
} else {
Some(password.to_string())
}
} else {
trimmed.split_whitespace().find_map(|token| {
let (k, v) = token.split_once('=')?;
(k == "password" && !v.is_empty()).then(|| v.to_string())
})
}
}
#[cfg(test)]
mod tests {
use bytes::BufMut;
use tokio::io::duplex;
use super::*;
#[test]
fn cleartext_password_response_is_unmodified() {
assert_eq!(build_password_response("alice", "hunter2", None), "hunter2");
}
#[test]
fn md5_password_response_matches_postgres_protocol_reference_vector() {
let salt = [0x2a, 0x3d, 0x8f, 0xe0];
let response = build_password_response("md5_user", "password", Some(salt));
assert_eq!(response, "md562af4dd09bbb41884907a838a3233294");
assert!(response.starts_with("md5"));
}
#[test]
fn md5_password_response_changes_with_salt() {
let a = build_password_response("bob", "s3cret", Some([0, 0, 0, 0]));
let b = build_password_response("bob", "s3cret", Some([1, 2, 3, 4]));
assert_ne!(a, b);
}
#[test]
fn md5_password_response_changes_with_user() {
let a = build_password_response("alice", "s3cret", Some([1, 2, 3, 4]));
let b = build_password_response("bob", "s3cret", Some([1, 2, 3, 4]));
assert_ne!(a, b);
}
#[test]
fn startup_params_include_replication_database_and_dbname_when_present() {
let params = build_startup_params("alice", Some("mydb"));
assert_eq!(
params,
vec![
("user", "alice"),
("replication", "database"),
("database", "mydb"),
]
);
}
#[test]
fn startup_params_omit_database_key_when_dbname_absent() {
let params = build_startup_params("alice", None);
assert_eq!(params, vec![("user", "alice"), ("replication", "database")]);
assert!(!params.iter().any(|(k, _)| *k == "database"));
}
#[test]
fn extract_password_from_kv_form() {
assert_eq!(
extract_password("host=x port=5432 user=bob password=hunter2"),
Some("hunter2".to_string())
);
}
#[test]
fn extract_password_missing_in_kv_form() {
assert_eq!(extract_password("host=x user=bob"), None);
}
#[test]
fn extract_password_empty_value_in_kv_form_is_none() {
assert_eq!(extract_password("host=x user=bob password="), None);
}
#[test]
fn extract_password_from_uri_form() {
assert_eq!(
extract_password("postgres://bob:hunter2@host/db"),
Some("hunter2".to_string())
);
}
#[test]
fn extract_password_from_postgresql_scheme_uri_form() {
assert_eq!(
extract_password("postgresql://bob:hunter2@host:5433/db"),
Some("hunter2".to_string())
);
}
#[test]
fn extract_password_absent_in_uri_form_without_colon() {
assert_eq!(extract_password("postgres://bob@host/db"), None);
}
#[test]
fn extract_password_absent_in_uri_form_without_userinfo() {
assert_eq!(extract_password("postgres://host/db"), None);
}
#[test]
fn extract_password_ignores_query_string() {
assert_eq!(
extract_password("postgres://bob:hunter2@host/db?sslmode=require"),
Some("hunter2".to_string())
);
}
#[test]
fn require_password_ok_when_present() {
assert_eq!(require_password(Some("x")).unwrap(), "x");
}
#[test]
fn require_password_err_when_absent() {
assert!(matches!(
require_password(None),
Err(PgError::Replication(_))
));
}
#[test]
fn classify_ssl_response_s_means_proceed() {
assert!(classify_ssl_response(b'S').is_ok());
}
#[test]
fn classify_ssl_response_n_means_tls_unsupported() {
assert!(matches!(classify_ssl_response(b'N'), Err(PgError::Tls(_))));
}
#[test]
fn classify_ssl_response_other_byte_is_protocol_error() {
assert!(matches!(
classify_ssl_response(b'X'),
Err(PgError::Protocol(_))
));
}
#[test]
fn select_scram_mechanism_picks_scram_sha_256_when_only_option() {
let mechanisms = ["SCRAM-SHA-256"];
assert_eq!(
select_scram_mechanism(&mechanisms).unwrap(),
"SCRAM-SHA-256"
);
}
#[test]
fn select_scram_mechanism_picks_scram_sha_256_among_multiple() {
let mechanisms = ["SCRAM-SHA-256-PLUS", "SCRAM-SHA-256"];
assert_eq!(
select_scram_mechanism(&mechanisms).unwrap(),
"SCRAM-SHA-256"
);
}
#[test]
fn select_scram_mechanism_rejects_list_without_scram_sha_256() {
let mechanisms = ["SCRAM-SHA-256-PLUS"];
let err = select_scram_mechanism(&mechanisms).unwrap_err();
let PgError::Replication(text) = err else {
panic!("expected PgError::Replication, got a different variant");
};
assert!(
text.contains("SCRAM-SHA-256"),
"error text should name the missing mechanism: {text}"
);
assert!(
text.contains("SCRAM-SHA-256-PLUS"),
"error text should list what the server did offer: {text}"
);
}
#[test]
fn select_scram_mechanism_rejects_empty_list() {
let mechanisms: [&str; 0] = [];
assert!(matches!(
select_scram_mechanism(&mechanisms),
Err(PgError::Replication(_))
));
}
fn encode_message(tag: u8, body: &[u8]) -> BytesMut {
let mut buf = BytesMut::new();
buf.put_u8(tag);
let len = u32::try_from(body.len() + 4).expect("test message body fits in u32");
buf.put_u32(len);
buf.put_slice(body);
buf
}
fn encode_error_fields(fields: &[(u8, &str)]) -> BytesMut {
let mut body = BytesMut::new();
for (type_, value) in fields {
body.put_u8(*type_);
body.put_slice(value.as_bytes());
body.put_u8(0);
}
body.put_u8(0); body
}
fn encode_sasl_mechanisms(mechanisms: &[&str]) -> BytesMut {
let mut body = BytesMut::new();
for mechanism in mechanisms {
body.put_slice(mechanism.as_bytes());
body.put_u8(0);
}
body.put_u8(0); body
}
async fn read_raw_frontend_message<S: AsyncRead + Unpin>(stream: &mut S) -> (u8, Vec<u8>) {
let mut header = [0_u8; 5];
stream
.read_exact(&mut header)
.await
.expect("read frontend message header");
let tag = header[0];
let len = u32::from_be_bytes([header[1], header[2], header[3], header[4]]) as usize;
let mut body = vec![0_u8; len - 4];
stream
.read_exact(&mut body)
.await
.expect("read frontend message body");
(tag, body)
}
#[test]
fn authentication_ok_wire_bytes_parse_to_the_expected_variant() {
let mut body = BytesMut::new();
body.put_i32(0);
let mut buf = encode_message(b'R', &body);
let msg = Message::parse(&mut buf).unwrap().unwrap();
assert!(matches!(msg, Message::AuthenticationOk));
}
#[test]
fn authentication_md5_password_wire_bytes_feed_build_password_response() {
let salt = [0x2a, 0x3d, 0x8f, 0xe0];
let mut body = BytesMut::new();
body.put_i32(5);
body.put_slice(&salt);
let mut buf = encode_message(b'R', &body);
let msg = Message::parse(&mut buf).unwrap().unwrap();
let Message::AuthenticationMd5Password(md5_body) = msg else {
panic!("expected Message::AuthenticationMd5Password, got a different variant");
};
assert_eq!(md5_body.salt(), salt);
let response = build_password_response("md5_user", "password", Some(md5_body.salt()));
assert_eq!(response, "md562af4dd09bbb41884907a838a3233294");
}
#[test]
fn error_response_wire_bytes_feed_server_error_response_with_sqlstate_and_message() {
let body = encode_error_fields(&[
(b'S', "FATAL"),
(b'C', "28000"),
(b'M', "no pg_hba.conf entry for replication connection"),
]);
let mut buf = encode_message(b'E', &body);
let msg = Message::parse(&mut buf).unwrap().unwrap();
let Message::ErrorResponse(error_body) = msg else {
panic!("expected Message::ErrorResponse, got a different variant");
};
let err = server_error_response(&error_body);
let PgError::Replication(text) = err else {
panic!("expected PgError::Replication, got a different variant");
};
assert!(
text.contains("28000"),
"error text missing SQLSTATE: {text}"
);
assert!(
text.contains("no pg_hba.conf entry for replication connection"),
"error text missing message: {text}"
);
}
#[test]
fn error_response_with_malformed_field_list_yields_protocol_error() {
let mut buf = encode_message(b'E', &[]);
let msg = Message::parse(&mut buf).unwrap().unwrap();
let Message::ErrorResponse(error_body) = msg else {
panic!("expected Message::ErrorResponse, got a different variant");
};
let err = server_error_response(&error_body);
assert!(matches!(err, PgError::Protocol(_)));
}
#[test]
fn message_parse_returns_none_when_header_incomplete() {
let mut buf = BytesMut::new();
buf.put_u8(b'R');
buf.put_slice(&[0, 0]);
assert!(Message::parse(&mut buf).unwrap().is_none());
}
#[test]
fn message_parse_returns_none_when_body_incomplete() {
let mut buf = BytesMut::new();
buf.put_u8(b'R');
buf.put_u32(8);
buf.put_slice(&[0, 0]);
assert!(Message::parse(&mut buf).unwrap().is_none());
}
#[tokio::test]
async fn authenticate_cleartext_password_flow_end_to_end() {
let (mut client, mut server) = duplex(4096);
let server_task = tokio::spawn(async move {
let mut body = BytesMut::new();
body.put_i32(3); server
.write_all(&encode_message(b'R', &body))
.await
.expect("write AuthenticationCleartextPassword");
let (tag, body) = read_raw_frontend_message(&mut server).await;
assert_eq!(tag, b'p');
assert_eq!(&body[..body.len() - 1], b"hunter2");
assert_eq!(
body.last(),
Some(&0_u8),
"PasswordMessage must be NUL-terminated"
);
let mut ok_body = BytesMut::new();
ok_body.put_i32(0);
server
.write_all(&encode_message(b'R', &ok_body))
.await
.expect("write AuthenticationOk");
});
let mut buf = BytesMut::new();
let result = authenticate(&mut client, &mut buf, "alice", Some("hunter2")).await;
assert!(result.is_ok(), "authenticate() failed: {result:?}");
server_task.await.expect("server task panicked");
}
#[tokio::test]
async fn authenticate_md5_password_flow_end_to_end() {
let (mut client, mut server) = duplex(4096);
let salt: [u8; 4] = [0x2a, 0x3d, 0x8f, 0xe0];
let server_task = tokio::spawn(async move {
let mut body = BytesMut::new();
body.put_i32(5); body.put_slice(&salt);
server
.write_all(&encode_message(b'R', &body))
.await
.expect("write AuthenticationMD5Password");
let (tag, body) = read_raw_frontend_message(&mut server).await;
assert_eq!(tag, b'p');
let expected = build_password_response("md5_user", "password", Some(salt));
assert_eq!(&body[..body.len() - 1], expected.as_bytes());
let mut ok_body = BytesMut::new();
ok_body.put_i32(0);
server
.write_all(&encode_message(b'R', &ok_body))
.await
.expect("write AuthenticationOk");
});
let mut buf = BytesMut::new();
let result = authenticate(&mut client, &mut buf, "md5_user", Some("password")).await;
assert!(result.is_ok(), "authenticate() failed: {result:?}");
server_task.await.expect("server task panicked");
}
#[tokio::test]
async fn authenticate_fails_clearly_when_password_required_but_missing() {
let (mut client, mut server) = duplex(4096);
let server_task = tokio::spawn(async move {
let mut body = BytesMut::new();
body.put_i32(3); server
.write_all(&encode_message(b'R', &body))
.await
.expect("write AuthenticationCleartextPassword");
});
let mut buf = BytesMut::new();
let result = authenticate(&mut client, &mut buf, "alice", None).await;
assert!(matches!(result, Err(PgError::Replication(_))));
server_task.await.expect("server task panicked");
}
#[tokio::test]
async fn authenticate_stops_cleanly_on_error_response() {
let (mut client, mut server) = duplex(4096);
let server_task = tokio::spawn(async move {
server
.write_all(&encode_message(b'E', &encode_error_fields(&[])))
.await
.expect("write ErrorResponse");
});
let mut buf = BytesMut::new();
let result = authenticate(&mut client, &mut buf, "alice", Some("hunter2")).await;
assert!(matches!(result, Err(PgError::Replication(_))));
server_task.await.expect("server task panicked");
}
#[tokio::test]
async fn authenticate_surfaces_structured_error_response_fields() {
let (mut client, mut server) = duplex(4096);
let server_task = tokio::spawn(async move {
let body = encode_error_fields(&[
(b'S', "FATAL"),
(b'C', "28000"),
(b'M', "no pg_hba.conf entry for replication connection"),
]);
server
.write_all(&encode_message(b'E', &body))
.await
.expect("write ErrorResponse");
});
let mut buf = BytesMut::new();
let result = authenticate(&mut client, &mut buf, "alice", Some("hunter2")).await;
let Err(PgError::Replication(text)) = result else {
panic!("expected Err(PgError::Replication(_)), got a different result");
};
assert!(
text.contains("28000"),
"error text missing SQLSTATE: {text}"
);
assert!(
text.contains("no pg_hba.conf entry for replication connection"),
"error text missing message: {text}"
);
server_task.await.expect("server task panicked");
}
#[tokio::test]
async fn authenticate_fails_clearly_when_server_does_not_offer_scram_sha_256() {
let (mut client, mut server) = duplex(4096);
let server_task = tokio::spawn(async move {
let mut body = BytesMut::new();
body.put_i32(10); body.put_slice(&encode_sasl_mechanisms(&["SCRAM-SHA-256-PLUS"]));
server
.write_all(&encode_message(b'R', &body))
.await
.expect("write AuthenticationSASL");
});
let mut buf = BytesMut::new();
let result = authenticate(&mut client, &mut buf, "alice", Some("hunter2")).await;
let Err(PgError::Replication(text)) = result else {
panic!("expected Err(PgError::Replication(_)), got a different result");
};
assert!(
text.contains("SCRAM-SHA-256"),
"error text should name the required mechanism: {text}"
);
server_task.await.expect("server task panicked");
}
#[tokio::test]
async fn authenticate_proceeds_with_scram_sha_256_when_offered_among_multiple_mechanisms() {
let (mut client, mut server) = duplex(4096);
let server_task = tokio::spawn(async move {
let mut body = BytesMut::new();
body.put_i32(10); body.put_slice(&encode_sasl_mechanisms(&[
"SCRAM-SHA-256-PLUS",
"SCRAM-SHA-256",
]));
server
.write_all(&encode_message(b'R', &body))
.await
.expect("write AuthenticationSASL");
let (tag, req_body) = read_raw_frontend_message(&mut server).await;
assert_eq!(tag, b'p', "expected a SASLInitialResponse");
assert!(
req_body.starts_with(b"SCRAM-SHA-256\0"),
"expected the initial response to name SCRAM-SHA-256"
);
server
.write_all(&encode_message(b'E', &encode_error_fields(&[])))
.await
.expect("write ErrorResponse");
});
let mut buf = BytesMut::new();
let result = authenticate(&mut client, &mut buf, "alice", Some("hunter2")).await;
assert!(
matches!(result, Err(PgError::Replication(_))),
"authenticate() should reach the SCRAM exchange and stop on ErrorResponse, got {result:?}"
);
server_task.await.expect("server task panicked");
}
#[tokio::test]
async fn authenticate_fails_on_unexpected_message() {
let (mut client, mut server) = duplex(4096);
let server_task = tokio::spawn(async move {
server
.write_all(&encode_message(b'Z', b"I"))
.await
.expect("write ReadyForQuery");
});
let mut buf = BytesMut::new();
let result = authenticate(&mut client, &mut buf, "alice", Some("hunter2")).await;
assert!(matches!(result, Err(PgError::Protocol(_))));
server_task.await.expect("server task panicked");
}
#[tokio::test]
async fn scram_exchange_stops_cleanly_when_server_sends_error_after_initial_response() {
let (mut client, mut server) = duplex(4096);
let server_task = tokio::spawn(async move {
let (tag, _body) = read_raw_frontend_message(&mut server).await;
assert_eq!(tag, b'p', "expected a SASLInitialResponse");
server
.write_all(&encode_message(b'E', &encode_error_fields(&[])))
.await
.expect("write ErrorResponse");
});
let mut buf = BytesMut::new();
let result = run_scram_exchange(&mut client, &mut buf, "password").await;
assert!(matches!(result, Err(PgError::Replication(_))));
server_task.await.expect("server task panicked");
}
#[tokio::test]
async fn scram_exchange_fails_on_unexpected_message_after_initial_response() {
let (mut client, mut server) = duplex(4096);
let server_task = tokio::spawn(async move {
let (tag, _body) = read_raw_frontend_message(&mut server).await;
assert_eq!(tag, b'p');
let mut ok_body = BytesMut::new();
ok_body.put_i32(0);
server
.write_all(&encode_message(b'R', &ok_body))
.await
.expect("write AuthenticationOk (unexpected here)");
});
let mut buf = BytesMut::new();
let result = run_scram_exchange(&mut client, &mut buf, "password").await;
assert!(matches!(result, Err(PgError::Protocol(_))));
server_task.await.expect("server task panicked");
}
#[tokio::test]
async fn drain_until_ready_consumes_backend_key_data_and_parameter_status() {
let (mut client, mut server) = duplex(4096);
let server_task = tokio::spawn(async move {
let mut key_body = BytesMut::new();
key_body.put_i32(1234);
key_body.put_i32(5678);
server
.write_all(&encode_message(b'K', &key_body))
.await
.expect("write BackendKeyData");
let mut ps_body = BytesMut::new();
ps_body.put_slice(b"server_version\0");
ps_body.put_slice(b"16.0\0");
server
.write_all(&encode_message(b'S', &ps_body))
.await
.expect("write ParameterStatus");
server
.write_all(&encode_message(b'Z', b"I"))
.await
.expect("write ReadyForQuery");
});
let mut buf = BytesMut::new();
let result = drain_until_ready(&mut client, &mut buf).await;
assert!(result.is_ok(), "drain_until_ready() failed: {result:?}");
server_task.await.expect("server task panicked");
}
#[tokio::test]
async fn drain_until_ready_stops_cleanly_on_error_response() {
let (mut client, mut server) = duplex(4096);
let server_task = tokio::spawn(async move {
server
.write_all(&encode_message(b'E', &encode_error_fields(&[])))
.await
.expect("write ErrorResponse");
});
let mut buf = BytesMut::new();
let result = drain_until_ready(&mut client, &mut buf).await;
assert!(matches!(result, Err(PgError::Replication(_))));
server_task.await.expect("server task panicked");
}
#[test]
fn maybe_tls_stream_is_send() {
fn assert_send<T: Send>() {}
assert_send::<MaybeTlsStream>();
}
}