mirror of
https://github.com/postgresml/pgcat.git
synced 2026-03-27 02:36:29 +00:00
parse startup client parameters (#16)
This commit is contained in:
@@ -8,6 +8,8 @@ use tokio::io::{AsyncReadExt, BufReader};
|
|||||||
use tokio::net::tcp::{OwnedReadHalf, OwnedWriteHalf};
|
use tokio::net::tcp::{OwnedReadHalf, OwnedWriteHalf};
|
||||||
use tokio::net::TcpStream;
|
use tokio::net::TcpStream;
|
||||||
|
|
||||||
|
use std::collections::HashMap;
|
||||||
|
|
||||||
use crate::config::Role;
|
use crate::config::Role;
|
||||||
use crate::errors::Error;
|
use crate::errors::Error;
|
||||||
use crate::messages::*;
|
use crate::messages::*;
|
||||||
@@ -52,6 +54,9 @@ pub struct Client {
|
|||||||
// Unless client specifies, route queries to the servers that have this role,
|
// Unless client specifies, route queries to the servers that have this role,
|
||||||
// e.g. primary or replicas or any.
|
// e.g. primary or replicas or any.
|
||||||
default_server_role: Option<Role>,
|
default_server_role: Option<Role>,
|
||||||
|
|
||||||
|
// Client parameters, e.g. user, client_encoding, etc.
|
||||||
|
parameters: HashMap<String, String>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl Client {
|
impl Client {
|
||||||
@@ -96,7 +101,7 @@ impl Client {
|
|||||||
// Regular startup message.
|
// Regular startup message.
|
||||||
196608 => {
|
196608 => {
|
||||||
// TODO: perform actual auth.
|
// TODO: perform actual auth.
|
||||||
// TODO: record startup parameters client sends over.
|
let parameters = parse_startup(bytes.clone())?;
|
||||||
|
|
||||||
// Generate random backend ID and secret key
|
// Generate random backend ID and secret key
|
||||||
let process_id: i32 = rand::random();
|
let process_id: i32 = rand::random();
|
||||||
@@ -121,6 +126,7 @@ impl Client {
|
|||||||
secret_key: secret_key,
|
secret_key: secret_key,
|
||||||
client_server_map: client_server_map,
|
client_server_map: client_server_map,
|
||||||
default_server_role: default_server_role,
|
default_server_role: default_server_role,
|
||||||
|
parameters: parameters,
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -141,6 +147,7 @@ impl Client {
|
|||||||
secret_key: secret_key,
|
secret_key: secret_key,
|
||||||
client_server_map: client_server_map,
|
client_server_map: client_server_map,
|
||||||
default_server_role: default_server_role,
|
default_server_role: default_server_role,
|
||||||
|
parameters: HashMap::new(),
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -1,9 +1,11 @@
|
|||||||
use bytes::{BufMut, BytesMut};
|
use bytes::{Buf, BufMut, BytesMut};
|
||||||
use md5::{Digest, Md5};
|
use md5::{Digest, Md5};
|
||||||
use tokio::io::{AsyncReadExt, AsyncWriteExt, BufReader};
|
use tokio::io::{AsyncReadExt, AsyncWriteExt, BufReader};
|
||||||
use tokio::net::tcp::{OwnedReadHalf, OwnedWriteHalf};
|
use tokio::net::tcp::{OwnedReadHalf, OwnedWriteHalf};
|
||||||
use tokio::net::TcpStream;
|
use tokio::net::TcpStream;
|
||||||
|
|
||||||
|
use std::collections::HashMap;
|
||||||
|
|
||||||
use crate::errors::Error;
|
use crate::errors::Error;
|
||||||
|
|
||||||
// This is a funny one. `psql` parses this to figure out which
|
// This is a funny one. `psql` parses this to figure out which
|
||||||
@@ -105,6 +107,51 @@ pub async fn startup(stream: &mut TcpStream, user: &str, database: &str) -> Resu
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Parse StartupMessage parameters.
|
||||||
|
/// e.g. user, database, application_name, etc.
|
||||||
|
pub fn parse_startup(mut bytes: BytesMut) -> Result<HashMap<String, String>, Error> {
|
||||||
|
let mut result = HashMap::new();
|
||||||
|
let mut buf = Vec::new();
|
||||||
|
let mut tmp = String::new();
|
||||||
|
|
||||||
|
while bytes.has_remaining() {
|
||||||
|
let mut c = bytes.get_u8();
|
||||||
|
|
||||||
|
// Null-terminated C-strings.
|
||||||
|
while c != 0 {
|
||||||
|
tmp.push(c as char);
|
||||||
|
c = bytes.get_u8();
|
||||||
|
}
|
||||||
|
|
||||||
|
if tmp.len() > 0 {
|
||||||
|
buf.push(tmp.clone());
|
||||||
|
tmp.clear();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Expect pairs of name and value
|
||||||
|
// and at least one pair to be present.
|
||||||
|
if buf.len() % 2 != 0 && buf.len() >= 2 {
|
||||||
|
return Err(Error::ClientBadStartup);
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut i = 0;
|
||||||
|
while i < buf.len() {
|
||||||
|
let name = buf[i].clone();
|
||||||
|
let value = buf[i + 1].clone();
|
||||||
|
let _ = result.insert(name, value);
|
||||||
|
i += 2;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Minimum required parameters
|
||||||
|
// I want to have the user at the very minimum, according to the protocol spec.
|
||||||
|
if !result.contains_key("user") {
|
||||||
|
return Err(Error::ClientBadStartup);
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(result)
|
||||||
|
}
|
||||||
|
|
||||||
/// Send password challenge response to the server.
|
/// Send password challenge response to the server.
|
||||||
/// This is the MD5 challenge.
|
/// This is the MD5 challenge.
|
||||||
pub async fn md5_password(
|
pub async fn md5_password(
|
||||||
|
|||||||
Reference in New Issue
Block a user