minisql/proto/src/handshake/server.rs
2024-01-22 21:59:39 +01:00

64 lines
2.4 KiB
Rust

use crate::handshake::errors::ServerHandshakeError;
use crate::handshake::request::HandshakeRequest;
use crate::handshake::response::HandshakeResponse;
use crate::message::backend::{AuthenticationOkData, BackendMessage, ReadyForQueryData};
use crate::message::special::SpecialMessage;
use crate::reader::frontend::FrontendProtoReader;
use crate::writer::backend::BackendProtoWriter;
use crate::writer::protowriter::ProtoFlush;
/// Performs server-side handshake with the client until ending it with `ReadyForQuery` message.
/// Client can send `CancelRequest` message instead of `StartupMessage` to cancel the request.
/// For more info visit the [`55.2.1. Start-up`](https://www.postgresql.org/docs/current/protocol-flow.html#PROTOCOL-FLOW-START-UP)
pub async fn do_server_handshake(
writer: &mut (impl BackendProtoWriter + ProtoFlush),
reader: &mut impl FrontendProtoReader,
response: HandshakeResponse,
) -> Result<HandshakeRequest, ServerHandshakeError> {
// Check if client requested SSL
match &reader.peek_special_message().await? {
Some(msg @ SpecialMessage::SSLRequest) => {
reader.consume_special_message(msg).await?;
writer.write_ssl_reject().await?;
writer.flush().await?;
}
_ => {
// No SSL request
}
}
// Wait for mandatory StartupMessage or CancelRequest
let startup_message = match &reader.peek_special_message().await? {
Some(msg @ SpecialMessage::StartupMessage(data)) => {
reader.consume_special_message(msg).await?;
data.clone()
}
Some(msg @ SpecialMessage::CancelRequest(data)) => {
reader.consume_special_message(msg).await?;
return Err(ServerHandshakeError::IsCancelRequest(data.clone()));
}
_ => {
return Err(ServerHandshakeError::MissingStartupMessage);
}
};
// Authenticate client
writer
.write_proto(BackendMessage::from(AuthenticationOkData { status: 0 }))
.await?;
// Send server parameters
let messages: Vec<BackendMessage> = response.into();
for message in messages {
writer.write_proto(message).await?;
}
// Finish the handshake
writer
.write_proto(BackendMessage::from(ReadyForQueryData { status: b'I' }))
.await?;
writer.flush().await?;
Ok(startup_message.into())
}