feat(proto): add client handshake implementation

This commit is contained in:
Jindřich Moravec 2023-12-23 00:29:40 +01:00
parent df4c4166d9
commit 7b2dce4dfb
6 changed files with 116 additions and 0 deletions

View file

@ -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),
}

View file

@ -0,0 +1,4 @@
pub mod request;
pub mod response;
pub mod shaker;
pub mod error;

View file

@ -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<ClientHandshakeRequest> for StartupMessageData {
fn from(request: ClientHandshakeRequest) -> Self {
Self { version: request.version, params: request.parameters }
}
}

View file

@ -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<Self, ClientHandshakeError> {
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),
}
}
}

View file

@ -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<ClientHandshakeResponse, ClientHandshakeError> {
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)
}

View file

@ -1,2 +1,3 @@
pub mod server; pub mod server;
pub mod errors; pub mod errors;
pub mod client;