diff --git a/src/server.rs b/src/server.rs index b242ec1..5305e33 100644 --- a/src/server.rs +++ b/src/server.rs @@ -1,3 +1,6 @@ +#![allow(dead_code)] +#![allow(unused_variables)] + use bytes::{Buf, BufMut, BytesMut}; use tokio::io::{AsyncReadExt, BufReader}; use tokio::net::tcp::{OwnedReadHalf, OwnedWriteHalf}; @@ -10,6 +13,9 @@ pub struct Server { read: BufReader, write: OwnedWriteHalf, buffer: BytesMut, + server_info: BytesMut, + backend_id: i32, + secret_key: i32, in_transaction: bool, bad: bool, } @@ -32,6 +38,10 @@ impl Server { startup(&mut stream, user, database).await?; + let mut server_info = BytesMut::with_capacity(25); + let mut backend_id: i32 = 0; + let mut secret_key: i32 = 0; + loop { let code = match stream.read_u8().await { Ok(code) => code as char, @@ -88,14 +98,22 @@ impl Server { Ok(_) => (), Err(_) => return Err(Error::SocketError), }; + + server_info.put_u8(b'S'); + server_info.put_i32(len); + server_info.put_slice(¶m[..]); } 'K' => { - // TODO: save cancellation secret - let mut cancel_secret = vec![0u8; len as usize - 4]; - match stream.read_exact(&mut cancel_secret).await { - Ok(_) => (), - Err(_) => return Err(Error::SocketError), + // Query cancellation data. + backend_id = match stream.read_i32().await { + Ok(id) => id, + Err(err) => return Err(Error::SocketError), + }; + + secret_key = match stream.read_i32().await { + Ok(id) => id, + Err(err) => return Err(Error::SocketError), }; } @@ -114,6 +132,9 @@ impl Server { read: BufReader::new(read), write: write, buffer: BytesMut::with_capacity(8196), + server_info: server_info, + backend_id: backend_id, + secret_key: secret_key, in_transaction: false, bad: false, }); @@ -204,6 +225,10 @@ impl Server { self.bad } + pub fn server_info(&self) -> BytesMut { + self.server_info.clone() + } + pub fn mark_bad(&mut self) { println!(">> Server marked bad"); self.bad = true;