omedusyo/minisql
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

feat: support for cancel requests in handshake

Browse source
This commit is contained in:
Jindřich Moravec 2024-01-22 21:59:39 +01:00
parent f9fb8f0670
commit ebabf50291
2 changed files with 9 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

3
proto/src/handshake/errors.rs
View file

@ -4,6 +4,7 @@ use crate::reader::errors::{ProtoConsumeError, ProtoPeekError, ProtoReadError};
use crate::writer::errors::ProtoWriteError; use crate::writer::errors::ProtoWriteError;
use thiserror::Error; use thiserror::Error;
use tokio::io; use tokio::io;
use crate::message::special::CancelRequestData;
#[derive(Debug, Error)] #[derive(Debug, Error)]
pub enum ClientHandshakeError { pub enum ClientHandshakeError {
@ -23,6 +24,8 @@ pub enum ClientHandshakeError {
pub enum ServerHandshakeError { pub enum ServerHandshakeError {
#[error("startup message not found")] #[error("startup message not found")]
MissingStartupMessage, MissingStartupMessage,
#[error("cancel request found instead of startup message")]
IsCancelRequest(CancelRequestData),
#[error("socket communication failed")] #[error("socket communication failed")]
Io(#[from] io::Error), Io(#[from] io::Error),
#[error("deserialization of inner data failed")] #[error("deserialization of inner data failed")]

7
proto/src/handshake/server.rs
View file

@ -8,6 +8,7 @@ use crate::writer::backend::BackendProtoWriter;
use crate::writer::protowriter::ProtoFlush; use crate::writer::protowriter::ProtoFlush;
/// Performs server-side handshake with the client until ending it with `ReadyForQuery` message. /// 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) /// 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( pub async fn do_server_handshake(
writer: &mut (impl BackendProtoWriter + ProtoFlush), writer: &mut (impl BackendProtoWriter + ProtoFlush),
@ -27,12 +28,16 @@ pub async fn do_server_handshake(
} }
} }
// Wait for mandatory StartupMessage // Wait for mandatory StartupMessage or CancelRequest
let startup_message = match &reader.peek_special_message().await? { let startup_message = match &reader.peek_special_message().await? {
Some(msg @ SpecialMessage::StartupMessage(data)) => { Some(msg @ SpecialMessage::StartupMessage(data)) => {
reader.consume_special_message(msg).await?; reader.consume_special_message(msg).await?;
data.clone() data.clone()
} }
Some(msg @ SpecialMessage::CancelRequest(data)) => {
reader.consume_special_message(msg).await?;
return Err(ServerHandshakeError::IsCancelRequest(data.clone()));
}
_ => { _ => {
return Err(ServerHandshakeError::MissingStartupMessage); return Err(ServerHandshakeError::MissingStartupMessage);
} }