diff --git a/proto/src/handshake/errors.rs b/proto/src/handshake/errors.rs index 0811790..cd2a8c4 100644 --- a/proto/src/handshake/errors.rs +++ b/proto/src/handshake/errors.rs @@ -4,6 +4,7 @@ use crate::reader::errors::{ProtoConsumeError, ProtoPeekError, ProtoReadError}; use crate::writer::errors::ProtoWriteError; use thiserror::Error; use tokio::io; +use crate::message::special::CancelRequestData; #[derive(Debug, Error)] pub enum ClientHandshakeError { @@ -23,6 +24,8 @@ pub enum ClientHandshakeError { pub enum ServerHandshakeError { #[error("startup message not found")] MissingStartupMessage, + #[error("cancel request found instead of startup message")] + IsCancelRequest(CancelRequestData), #[error("socket communication failed")] Io(#[from] io::Error), #[error("deserialization of inner data failed")] diff --git a/proto/src/handshake/server.rs b/proto/src/handshake/server.rs index 6c8deb2..d1a332f 100644 --- a/proto/src/handshake/server.rs +++ b/proto/src/handshake/server.rs @@ -8,6 +8,7 @@ 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), @@ -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? { 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); }