diff --git a/proto/src/handshake/client/error.rs b/proto/src/handshake/client/error.rs new file mode 100644 index 0000000..ab6f23c --- /dev/null +++ b/proto/src/handshake/client/error.rs @@ -0,0 +1,16 @@ +use thiserror::Error; +use crate::message::backend::BackendMessage; +use crate::reader::errors::ProtoReadError; +use crate::writer::errors::ProtoWriteError; + +#[derive(Debug, Error)] +pub enum ClientHandshakeError { + #[error("unexpected response from server")] + UnexpectedResponse, + #[error("unexpected auth response")] + UnexpectedAuthResponse(BackendMessage), + #[error("writing message to socket failed")] + Write(#[from] ProtoWriteError), + #[error("reading message from socket failed")] + Read(#[from] ProtoReadError), +} \ No newline at end of file diff --git a/proto/src/handshake/client/mod.rs b/proto/src/handshake/client/mod.rs new file mode 100644 index 0000000..060973a --- /dev/null +++ b/proto/src/handshake/client/mod.rs @@ -0,0 +1,4 @@ +pub mod request; +pub mod response; +pub mod shaker; +pub mod error; \ No newline at end of file diff --git a/proto/src/handshake/client/request.rs b/proto/src/handshake/client/request.rs new file mode 100644 index 0000000..54998b5 --- /dev/null +++ b/proto/src/handshake/client/request.rs @@ -0,0 +1,24 @@ +use crate::message::primitive::pgstring::PgString; +use crate::message::special::StartupMessageData; + +pub struct ClientHandshakeRequest { + version: i32, + parameters: Vec<(PgString, PgString)>, +} + +impl ClientHandshakeRequest { + pub fn new(version: i32) -> Self { + Self { version, parameters: Vec::new() } + } + + pub fn parameter(mut self, key: &str, value: &str) -> Self { + self.parameters.push((key.into(), value.into())); + self + } +} + +impl From for StartupMessageData { + fn from(request: ClientHandshakeRequest) -> Self { + Self { version: request.version, params: request.parameters } + } +} diff --git a/proto/src/handshake/client/response.rs b/proto/src/handshake/client/response.rs new file mode 100644 index 0000000..777cb81 --- /dev/null +++ b/proto/src/handshake/client/response.rs @@ -0,0 +1,37 @@ +use crate::handshake::client::error::ClientHandshakeError; +use crate::message::backend::BackendMessage; + +pub struct ClientHandshakeResponse { + pub version: String, + pub process_id: i32, + pub secret_key: i32, +} + +impl ClientHandshakeResponse { + pub fn from_backend_messages(message: &[BackendMessage]) -> Result { + let mut version = None; + let mut process_id = None; + let mut secret_key = None; + for message in message { + match message { + BackendMessage::ParameterStatus(data) => { + if data.name.as_str() == "server_version" { + version = Some(String::from(data.value.as_str())); + } + } + BackendMessage::BackendKeyData(data) => { + process_id = Some(data.process); + secret_key = Some(data.secret); + } + _ => {} + } + } + + match (version, process_id, secret_key) { + (Some(version), Some(process_id), Some(secret_key)) => { + Ok(Self { version, process_id, secret_key }) + } + _ => Err(ClientHandshakeError::UnexpectedResponse), + } + } +} \ No newline at end of file diff --git a/proto/src/handshake/client/shaker.rs b/proto/src/handshake/client/shaker.rs new file mode 100644 index 0000000..893f0b7 --- /dev/null +++ b/proto/src/handshake/client/shaker.rs @@ -0,0 +1,34 @@ +use crate::handshake::client::error::ClientHandshakeError; +use crate::handshake::client::request::ClientHandshakeRequest; +use crate::handshake::client::response::ClientHandshakeResponse; +use crate::message::backend::{AuthenticationOkData, BackendMessage}; +use crate::message::special::StartupMessageData; +use crate::reader::backend::BackendProtoReader; +use crate::writer::frontend::FrontendProtoWriter; +use crate::writer::protowriter::ProtoFlush; + +pub async fn do_client_handshake( + writer: &mut (impl FrontendProtoWriter + ProtoFlush), + reader: &mut impl BackendProtoReader, + request: ClientHandshakeRequest, +) -> Result { + let startup_message: StartupMessageData = request.into(); + writer.write_startup_message(startup_message).await?; + + let auth = reader.read_proto().await?; + if !matches!(auth, BackendMessage::AuthenticationOk(AuthenticationOkData { status: 0 })) { + return Err(ClientHandshakeError::UnexpectedAuthResponse(auth)); + } + + let mut messages = Vec::new(); + loop { + let msg = reader.read_proto().await?; + if matches!(msg, BackendMessage::ReadyForQuery(_)) { + break; + } + + messages.push(msg); + } + + ClientHandshakeResponse::from_backend_messages(&messages) +} \ No newline at end of file diff --git a/proto/src/handshake/mod.rs b/proto/src/handshake/mod.rs index 24a7408..0b12ab2 100644 --- a/proto/src/handshake/mod.rs +++ b/proto/src/handshake/mod.rs @@ -1,2 +1,3 @@ pub mod server; pub mod errors; +pub mod client;