pgdog.git / summary / log / commit / refs
commit eb99384f4698a42dcfd08c8aaeb9ee00a370bc63
Author: Lev Kokotov <lev.kokotov@gmail.com>
Date: Fri Jan 03 00:27:41 2025 +0000
save
Cargo.toml | 5 ++-
src/backend/error.rs | 3 ++
src/backend/pool/config.rs | 2 +-
src/backend/pool/connection.rs | 3 +-
src/backend/pool/pool.rs | 6 ++-
src/backend/server.rs | 46 ++++++++++++++++----
src/frontend/buffer.rs | 6 +--
src/frontend/listener.rs | 6 ++-
src/main.rs | 8 ++--
src/net/error.rs | 7 +++
src/net/stream.rs | 7 ++-
src/net/tls.rs | 97 ++++++++++++++++++++++++++++++++++--------
12 files changed, 153 insertions(+), 43 deletions(-)
diff --git a/Cargo.toml b/Cargo.toml
index 1dab2add..439812ce 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -9,7 +9,7 @@ authors = ["Lev Kokotov <lev.kokotov@gmail.com>"]
pin-project = "1"
tokio = { version = "1", features = ["full"] }
tracing = "0.1"
-tracing-subscriber = "0.3"
+tracing-subscriber = { version = "0.3", features = ["env-filter", "std"] }
parking_lot = "0.12"
thiserror = "2"
bytes = "1"
@@ -19,3 +19,6 @@ async-trait = "*"
rand = "*"
tokio-native-tls = "0.3"
once_cell = "1"
+tokio-rustls = "0.26"
+rustls-native-certs = "0.8"
+rustls-pki-types = "*"
diff --git a/src/backend/error.rs b/src/backend/error.rs
index 61e6b198..f7ca4887 100644
--- a/src/backend/error.rs
+++ b/src/backend/error.rs
@@ -10,6 +10,9 @@ pub enum Error {
#[error("{0}")]
Tls(#[from] tokio_native_tls::native_tls::Error),
+ #[error("{0}")]
+ InvalidTlsDnsName(#[from] rustls_pki_types::InvalidDnsNameError),
+
#[error("net: {0}")]
Net(#[from] crate::net::Error),
diff --git a/src/backend/pool/config.rs b/src/backend/pool/config.rs
index 84f5cda4..aef40cbd 100644
--- a/src/backend/pool/config.rs
+++ b/src/backend/pool/config.rs
@@ -20,7 +20,7 @@ pub struct Config {
impl Default for Config {
fn default() -> Self {
Self {
- min: 0,
+ min: 1,
max: 10,
checkout_timeout: 5_000,
idle_timeout: 60_000,
diff --git a/src/backend/pool/connection.rs b/src/backend/pool/connection.rs
index 6a08578a..555faf05 100644
--- a/src/backend/pool/connection.rs
+++ b/src/backend/pool/connection.rs
@@ -38,7 +38,8 @@ impl Connection {
/// Read a message from the server connection.
pub async fn read(&mut self) -> Result<Message, Error> {
if let Some(ref mut server) = self.server {
- server.read().await
+ let message = server.read().await?;
+ Ok(message)
} else {
// Suspend the future until select! cancels it.
loop {
diff --git a/src/backend/pool/pool.rs b/src/backend/pool/pool.rs
index a6a4dc63..5bec0e17 100644
--- a/src/backend/pool/pool.rs
+++ b/src/backend/pool/pool.rs
@@ -2,7 +2,7 @@
use std::collections::VecDeque;
use std::sync::Arc;
-use std::time::Duration;
+use std::time::{Duration, Instant};
use once_cell::sync::OnceCell;
use parking_lot::Mutex;
@@ -33,6 +33,7 @@ struct Inner {
taken: Vec<Mapping>,
config: Config,
waiting: usize,
+ banned_at: Option<Instant>,
}
struct Comms {
@@ -74,6 +75,7 @@ impl Pool {
taken: Vec::new(),
config: Config::default(),
waiting: 0,
+ banned_at: None,
})),
comms: Arc::new(Comms {
ready: Notify::new(),
@@ -176,6 +178,8 @@ impl Pool {
if server.done() {
guard.conns.push_back(server);
+ } else if server.error() {
+ guard.banned_at = Some(Instant::now());
}
let index = guard
diff --git a/src/backend/server.rs b/src/backend/server.rs
index f79741a1..fbbcf6ea 100644
--- a/src/backend/server.rs
+++ b/src/backend/server.rs
@@ -1,7 +1,8 @@
//! PostgreSQL serer connection.
-use std::time::Instant;
+use std::time::{Duration, Instant};
use bytes::{BufMut, BytesMut};
+use rustls_pki_types::ServerName;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpStream;
use tracing::{debug, info};
@@ -44,7 +45,12 @@ impl Server {
let connector = connector()?;
let plain = stream.take()?;
- stream = Stream::tls(connector.connect(addr, plain).await?);
+ let server_name = ServerName::try_from(addr.to_string())?;
+
+ let cipher =
+ tokio_rustls::TlsStream::Client(connector.connect(server_name, plain).await?);
+
+ stream = Stream::tls(cipher);
}
stream.write_all(&Startup::new().to_bytes()?).await?;
@@ -109,8 +115,6 @@ impl Server {
/// Request query cancellation for the given backend server identifier.
pub async fn cancel(addr: &str, id: &BackendKeyData) -> Result<(), Error> {
- debug!("cancelling query");
-
let mut stream = TcpStream::connect(addr).await?;
stream
.write_all(
@@ -129,19 +133,33 @@ impl Server {
/// Send messages to the server.
pub async fn send(&mut self, messages: Vec<impl Protocol + ToBytes>) -> Result<(), Error> {
self.state = State::Active;
- self.stream.send_many(messages).await?;
- Ok(())
+ if let Err(err) = self.stream.send_many(messages).await {
+ self.state = State::Error;
+ Err(err.into())
+ } else {
+ Ok(())
+ }
}
/// Flush all pending messages making sure they are sent to the server immediately.
pub async fn flush(&mut self) -> Result<(), Error> {
- self.stream.flush().await?;
- Ok(())
+ if let Err(err) = self.stream.flush().await {
+ self.state = State::Error;
+ Err(err.into())
+ } else {
+ Ok(())
+ }
}
/// Read a single message from the server.
pub async fn read(&mut self) -> Result<Message, Error> {
- let message = self.stream.read().await?;
+ let message = match self.stream.read().await {
+ Ok(message) => message,
+ Err(err) => {
+ self.state = State::Error;
+ return Err(err.into());
+ }
+ };
match message.code() {
'Z' => {
@@ -184,6 +202,11 @@ impl Server {
)
}
+ /// The server connection permanently failed.
+ pub fn error(&self) -> bool {
+ self.state == State::Error
+ }
+
/// Server parameters.
pub fn params(&self) -> &Vec<(String, String)> {
&self.params
@@ -223,4 +246,9 @@ impl Server {
pub fn id(&self) -> &BackendKeyData {
&self.id
}
+
+ /// How old this connection is.
+ pub fn age(&self) -> Duration {
+ self.created_at.elapsed()
+ }
}
diff --git a/src/frontend/buffer.rs b/src/frontend/buffer.rs
index 0e76ec42..d0c9271d 100644
--- a/src/frontend/buffer.rs
+++ b/src/frontend/buffer.rs
@@ -29,11 +29,11 @@ impl Buffer {
}
/// The buffer is full and the client won't send any more messages
- /// until it gets a reply.
+ /// until it gets a reply, or we don't want to buffer the data in memory.
pub fn full(&self) -> bool {
if let Some(ref message) = self.buffer.last() {
- // Flush (F) | Sync (F) | Query (F)
- if matches!(message.code(), 'H' | 'S' | 'Q') {
+ // Flush (F) | Sync (F) | Query (F) | CopyData (F) | CopyDone (F)
+ if matches!(message.code(), 'H' | 'S' | 'Q' | 'd' | 'c') {
return true;
}
}
diff --git a/src/frontend/listener.rs b/src/frontend/listener.rs
index 1f588ca6..8c27e60d 100644
--- a/src/frontend/listener.rs
+++ b/src/frontend/listener.rs
@@ -42,6 +42,10 @@ impl Listener {
pub async fn listen(&mut self) -> Result<(), Error> {
info!("🐕 pgDog listening on {}", self.addr);
+ // Init the pool early in case we need to spin up some
+ // connections.
+ let _pool = pool();
+
let listener = TcpListener::bind(&self.addr).await?;
// Load TLS cert, if any.
@@ -87,7 +91,7 @@ impl Listener {
stream.send_flush(SslReply::Yes).await?;
let plain = stream.take()?;
let cipher = tls.accept(plain).await?;
- stream = Stream::tls(cipher);
+ stream = Stream::tls(tokio_rustls::TlsStream::Server(cipher));
} else {
stream.send_flush(SslReply::No).await?;
}
diff --git a/src/main.rs b/src/main.rs
index 5357f37b..616df1fe 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,5 +1,5 @@
use frontend::listener::Listener;
-use tracing::Level;
+use tracing_subscriber::{fmt, prelude::*, EnvFilter};
pub mod auth;
pub mod backend;
@@ -12,11 +12,11 @@ pub mod state;
#[tokio::main]
async fn main() {
- tracing_subscriber::fmt()
- .with_max_level(Level::DEBUG)
+ tracing_subscriber::registry()
+ .with(fmt::layer())
+ .with(EnvFilter::from_default_env())
.init();
let mut listener = Listener::new("0.0.0.0:6432");
listener.listen().await.unwrap();
- println!("Hello, world!");
}
diff --git a/src/net/error.rs b/src/net/error.rs
index e7e0b306..2491d570 100644
--- a/src/net/error.rs
+++ b/src/net/error.rs
@@ -1,6 +1,7 @@
//! Frontend errors.
use thiserror::Error;
+use tokio_rustls::rustls;
#[derive(Debug, Error)]
pub enum Error {
@@ -27,4 +28,10 @@ pub enum Error {
#[error("unexpected ssl request reply: {0}")]
UnexpectedSslReply(char),
+
+ #[error("{0}")]
+ TlsCertificate(#[from] rustls::pki_types::pem::Error),
+
+ #[error("{0}")]
+ Rustls(#[from] rustls::Error),
}
diff --git a/src/net/stream.rs b/src/net/stream.rs
index f14e4007..9c8ffe21 100644
--- a/src/net/stream.rs
+++ b/src/net/stream.rs
@@ -4,7 +4,6 @@ use bytes::{BufMut, BytesMut};
use pin_project::pin_project;
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, BufStream, ReadBuf};
use tokio::net::TcpStream;
-use tokio_native_tls::TlsStream;
use tracing::debug;
use std::io::Error;
@@ -17,7 +16,7 @@ use super::messages::{FromBytes, Message, Protocol, ToBytes};
#[pin_project(project = StreamProjection)]
pub enum Stream {
Plain(#[pin] BufStream<TcpStream>),
- Tls(#[pin] BufStream<TlsStream<TcpStream>>),
+ Tls(#[pin] BufStream<tokio_rustls::TlsStream<TcpStream>>),
}
impl AsyncRead for Stream {
@@ -77,7 +76,7 @@ impl Stream {
}
/// Wrap an encrypted TCP stream.
- pub fn tls(stream: TlsStream<TcpStream>) -> Self {
+ pub fn tls(stream: tokio_rustls::TlsStream<TcpStream>) -> Self {
Self::Tls(BufStream::new(stream))
}
@@ -91,7 +90,7 @@ impl Stream {
&mut self,
message: impl ToBytes + Protocol,
) -> Result<(), crate::net::Error> {
- debug!("📡 <= {}", message.code());
+ // debug!("📡 <= {}", message.code());
let bytes = message.to_bytes()?;
match self {
diff --git a/src/net/tls.rs b/src/net/tls.rs
index a7bd4c6d..1c57f09a 100644
--- a/src/net/tls.rs
+++ b/src/net/tls.rs
@@ -1,33 +1,40 @@
//! TLS configuration.
+use std::sync::Arc;
+
use once_cell::sync::OnceCell;
-use tokio::fs::read_to_string;
-use tokio_native_tls::{
- native_tls::{Identity, TlsAcceptor, TlsConnector},
- TlsAcceptor as TlsAcceptorAsync, TlsConnector as TlsConnectorAsync,
+use rustls::pki_types::{CertificateDer, PrivateKeyDer};
+use tokio_rustls::rustls::{
+ self,
+ client::danger::{ServerCertVerified, ServerCertVerifier},
+ pki_types::pem::PemObject,
+ server::danger::ClientCertVerifier,
+ ClientConfig,
};
+use tokio_rustls::{TlsAcceptor, TlsConnector};
use tracing::info;
use super::Error;
-static ACCEPTOR: OnceCell<TlsAcceptorAsync> = OnceCell::new();
-static CONNECTOR: OnceCell<TlsConnectorAsync> = OnceCell::new();
+static ACCEPTOR: OnceCell<TlsAcceptor> = OnceCell::new();
+static CONNECTOR: OnceCell<TlsConnector> = OnceCell::new();
/// Create a new TLS acceptor from the cert and key.
-pub async fn acceptor() -> Result<Option<TlsAcceptorAsync>, Error> {
+pub async fn acceptor() -> Result<Option<TlsAcceptor>, Error> {
if let Some(acceptor) = ACCEPTOR.get() {
return Ok(Some(acceptor.clone()));
}
- let pem = read_to_string("tests/cert.pem").await?;
- let key = read_to_string("tests/key.pem").await?;
+ let pem = CertificateDer::from_pem_file("tests/cert.pem")?;
+ let key = PrivateKeyDer::from_pem_file("tests/key.pem")?;
- let identity = Identity::from_pkcs8(pem.as_bytes(), key.as_bytes()).unwrap();
- let acceptor = TlsAcceptor::new(identity).unwrap();
+ let config = rustls::ServerConfig::builder()
+ .with_no_client_auth()
+ .with_single_cert(vec![pem], key)?;
- info!("🔑 TLS on");
+ let acceptor = tokio_rustls::TlsAcceptor::from(Arc::new(config));
- let acceptor = TlsAcceptorAsync::from(acceptor);
+ info!("🔑 TLS on");
// A bit of a race, but it's not a big deal unless this is called
// with different certificate/secret key.
@@ -37,18 +44,72 @@ pub async fn acceptor() -> Result<Option<TlsAcceptorAsync>, Error> {
}
/// Create new TLS connector.
-pub fn connector() -> Result<TlsConnectorAsync, Error> {
+pub fn connector() -> Result<TlsConnector, Error> {
if let Some(connector) = CONNECTOR.get() {
return Ok(connector.clone());
}
- let connector = TlsConnector::builder()
- .danger_accept_invalid_certs(true)
- .danger_accept_invalid_hostnames(true)
+
+ let mut roots = rustls::RootCertStore::empty();
+ for cert in rustls_native_certs::load_native_certs().expect("load native certs") {
+ roots.add(cert)?;
+ }
+
+ let verifier = rustls::server::WebPkiClientVerifier::builder(roots.clone().into())
.build()
.unwrap();
- let connector = TlsConnectorAsync::from(connector);
+ let verifier = CertificateVerifyer { verifier };
+
+ let mut config = ClientConfig::builder()
+ .with_root_certificates(roots)
+ .with_no_client_auth();
+
+ config
+ .dangerous()
+ .set_certificate_verifier(Arc::new(verifier));
+
+ let connector = TlsConnector::from(Arc::new(config));
let _ = CONNECTOR.set(connector.clone());
Ok(connector)
}
+
+#[derive(Debug)]
+struct CertificateVerifyer {
+ verifier: Arc<dyn ClientCertVerifier>,
+}
+
+impl ServerCertVerifier for CertificateVerifyer {
+ fn verify_server_cert(
+ &self,
+ _end_entity: &CertificateDer<'_>,
+ _intermediates: &[CertificateDer<'_>],
+ _server_name: &rustls::pki_types::ServerName<'_>,
+ _ocsp_response: &[u8],
+ _now: rustls::pki_types::UnixTime,
+ ) -> Result<rustls::client::danger::ServerCertVerified, rustls::Error> {
+ Ok(ServerCertVerified::assertion())
+ }
+
+ fn verify_tls12_signature(
+ &self,
+ message: &[u8],
+ cert: &CertificateDer<'_>,
+ dss: &rustls::DigitallySignedStruct,
+ ) -> Result<rustls::client::danger::HandshakeSignatureValid, rustls::Error> {
+ self.verifier.verify_tls12_signature(message, cert, dss)
+ }
+
+ fn verify_tls13_signature(
+ &self,
+ message: &[u8],
+ cert: &CertificateDer<'_>,
+ dss: &rustls::DigitallySignedStruct,
+ ) -> Result<rustls::client::danger::HandshakeSignatureValid, rustls::Error> {
+ self.verifier.verify_tls13_signature(message, cert, dss)
+ }
+
+ fn supported_verify_schemes(&self) -> Vec<rustls::SignatureScheme> {
+ self.verifier.supported_verify_schemes()
+ }
+}
[parent: ccf53ea8bcf0]