diff --git a/proto/src/writer/backend.rs b/proto/src/writer/backend.rs new file mode 100644 index 0000000..4cbbdc8 --- /dev/null +++ b/proto/src/writer/backend.rs @@ -0,0 +1,59 @@ +use crate::message::backend::BackendMessage; +use crate::writer::oneway::OneWayProtoWriter; +use crate::writer::protowriter::ProtoWriter; +use async_trait::async_trait; +use tokio::io::{AsyncWrite, AsyncWriteExt}; + +#[async_trait] +pub trait BackendProtoWriter: OneWayProtoWriter { + async fn write_ssl_reject(&mut self) -> anyhow::Result<()>; +} + +#[async_trait] +impl BackendProtoWriter for ProtoWriter +where + W: AsyncWrite + Unpin + Send, +{ + async fn write_ssl_reject(&mut self) -> anyhow::Result<()> { + self.inner.write_u8(b'N').await?; + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::message::backend::AuthenticationOkData; + use crate::writer::protowriter::ProtoWriter; + use tokio::io::BufWriter; + + #[tokio::test] + async fn test_message_sequence() { + let writer = BufWriter::new(Vec::new()); + let mut writer = ProtoWriter::new(writer); + + writer + .write_proto(BackendMessage::AuthenticationOk(AuthenticationOkData { + status: 123, + })) + .await + .unwrap(); + + writer.write_proto(BackendMessage::NoData).await.unwrap(); + + assert_eq!( + writer.inner.buffer(), + vec![b'R', 0, 0, 0, 8, 0, 0, 0, 123, b'n', 0, 0, 0, 4] + ); + } + + #[tokio::test] + async fn test_ssl_reject() { + let writer = BufWriter::new(Vec::new()); + let mut writer = ProtoWriter::new(writer); + + writer.write_ssl_reject().await.unwrap(); + + assert_eq!(writer.inner.buffer(), vec![b'N']); + } +} diff --git a/proto/src/writer/mod.rs b/proto/src/writer/mod.rs index fac68e8..58e55a6 100644 --- a/proto/src/writer/mod.rs +++ b/proto/src/writer/mod.rs @@ -1,3 +1,4 @@ +pub mod backend; pub mod oneway; pub mod protowriter;