Plugins!!

Author: levkk

Cargo.lock

@@ -14,7 +14,7 @@ rand = "0.8"
 chrono = "0.4"
 sha-1 = "0.10"
 toml = "0.7"
-serde = "1"
+serde = { version = "1", features = ["derive"] }
 serde_derive = "1"
 regex = "1"
 num_cpus = "1"
@@ -44,6 +44,7 @@ webpki-roots = "0.23"
 rustls = { version = "0.21", features = ["dangerous_configuration"] }
 trust-dns-resolver = "0.22.0"
 tokio-test = "0.4.2"
+serde_json = "1"
 
 [target.'cfg(not(target_env = "msvc"))'.dependencies]
 jemallocator = "0.5.0"

Cargo.toml

@@ -77,6 +77,9 @@ admin_username = "admin_user"
 # Password to access the virtual administrative database
 admin_password = "admin_pass"
 
+# Plugins!!
+# plugins = ["pg_table_access", "intercept"]
+
 # pool configs are structured as pool.<pool_name>
 # the pool_name is what clients use as database name when connecting.
 # For a pool named `sharded_db`, clients access that pool using connection string like

pgcat.toml

@@ -12,9 +12,9 @@ use tokio::time::Instant;
 use crate::config::{get_config, reload_config, VERSION};
 use crate::errors::Error;
 use crate::messages::*;
+use crate::pool::ClientServerMap;
 use crate::pool::{get_all_pools, get_pool};
 use crate::stats::{get_client_stats, get_pool_stats, get_server_stats, ClientState, ServerState};
-use crate::ClientServerMap;
 
 pub fn generate_server_info_for_admin() -> BytesMut {
     let mut server_info = BytesMut::new();

src/admin.rs

@@ -16,6 +16,7 @@ use crate::auth_passthrough::refetch_auth_hash;
 use crate::config::{get_config, get_idle_client_in_transaction_timeout, Address, PoolMode};
 use crate::constants::*;
 use crate::messages::*;
+use crate::plugins::PluginOutput;
 use crate::pool::{get_pool, ClientServerMap, ConnectionPool};
 use crate::query_router::{Command, QueryRouter};
 use crate::server::Server;
@@ -765,6 +766,9 @@ where
 
         self.stats.register(self.stats.clone());
 
+        // Error returned by one of the plugins.
+        let mut plugin_output = None;
+
         // Our custom protocol loop.
         // We expect the client to either start a transaction with regular queries
         // or issue commands for our sharding and server selection protocol.
@@ -816,6 +820,22 @@ where
                 'Q' => {
                     if query_router.query_parser_enabled() {
                         if let Ok(ast) = QueryRouter::parse(&message) {
+                            let plugin_result = query_router.execute_plugins(&ast).await;
+
+                            match plugin_result {
+                                Ok(PluginOutput::Deny(error)) => {
+                                    error_response(&mut self.write, &error).await?;
+                                    continue;
+                                }
+
+                                Ok(PluginOutput::Intercept(result)) => {
+                                    write_all(&mut self.write, result).await?;
+                                    continue;
+                                }
+
+                                _ => (),
+                            };
+
                             let _ = query_router.infer(&ast);
                         }
                     }
@@ -826,6 +846,10 @@ where
 
                     if query_router.query_parser_enabled() {
                         if let Ok(ast) = QueryRouter::parse(&message) {
+                            if let Ok(output) = query_router.execute_plugins(&ast).await {
+                                plugin_output = Some(output);
+                            }
+
                             let _ = query_router.infer(&ast);
                         }
                     }
@@ -861,6 +885,18 @@ where
                 continue;
             }
 
+            // Check on plugin results.
+            match plugin_output {
+                Some(PluginOutput::Deny(error)) => {
+                    self.buffer.clear();
+                    error_response(&mut self.write, &error).await?;
+                    plugin_output = None;
+                    continue;
+                }
+
+                _ => (),
+            };
+
             // Get a pool instance referenced by the most up-to-date
             // pointer. This ensures we always read the latest config
             // when starting a query.
@@ -1089,6 +1125,27 @@ where
                 match code {
                     // Query
                     'Q' => {
+                        if query_router.query_parser_enabled() {
+                            if let Ok(ast) = QueryRouter::parse(&message) {
+                                let plugin_result = query_router.execute_plugins(&ast).await;
+
+                                match plugin_result {
+                                    Ok(PluginOutput::Deny(error)) => {
+                                        error_response(&mut self.write, &error).await?;
+                                        continue;
+                                    }
+
+                                    Ok(PluginOutput::Intercept(result)) => {
+                                        write_all(&mut self.write, result).await?;
+                                        continue;
+                                    }
+
+                                    _ => (),
+                                };
+
+                                let _ = query_router.infer(&ast);
+                            }
+                        }
                         debug!("Sending query to server");
 
                         self.send_and_receive_loop(
@@ -1128,6 +1185,14 @@ where
                     // Parse
                     // The query with placeholders is here, e.g. `SELECT * FROM users WHERE email = $1 AND active = $2`.
                     'P' => {
+                        if query_router.query_parser_enabled() {
+                            if let Ok(ast) = QueryRouter::parse(&message) {
+                                if let Ok(output) = query_router.execute_plugins(&ast).await {
+                                    plugin_output = Some(output);
+                                }
+                            }
+                        }
+
                         self.buffer.put(&message[..]);
                     }
 
@@ -1159,6 +1224,24 @@ where
                     'S' => {
                         debug!("Sending query to server");
 
+                        match plugin_output {
+                            Some(PluginOutput::Deny(error)) => {
+                                error_response(&mut self.write, &error).await?;
+                                plugin_output = None;
+                                self.buffer.clear();
+                                continue;
+                            }
+
+                            Some(PluginOutput::Intercept(result)) => {
+                                write_all(&mut self.write, result).await?;
+                                plugin_output = None;
+                                self.buffer.clear();
+                                continue;
+                            }
+
+                            _ => (),
+                        };
+
                         self.buffer.put(&message[..]);
 
                         let first_message_code = (*self.buffer.get(0).unwrap_or(&0)) as char;

src/client.rs

@@ -302,6 +302,8 @@ pub struct General {
     pub auth_query: Option<String>,
     pub auth_query_user: Option<String>,
     pub auth_query_password: Option<String>,
+
+    pub query_router_plugins: Option<Vec<String>>,
 }
 
 impl General {
@@ -402,6 +404,7 @@ impl Default for General {
             auth_query_user: None,
             auth_query_password: None,
             server_lifetime: 1000 * 3600 * 24, // 24 hours,
+            query_router_plugins: None,
         }
     }
 }

src/config.rs

@@ -1,7 +1,7 @@
 //! Errors.
 
 /// Various errors.
-#[derive(Debug, PartialEq)]
+#[derive(Debug, PartialEq, Clone)]
 pub enum Error {
     SocketError(String),
     ClientSocketError(String, ClientIdentifier),
@@ -26,7 +26,9 @@ pub enum Error {
     AuthPassthroughError(String),
     UnsupportedStatement,
     QueryRouterParserError(String),
+    PermissionDenied(String),
     PermissionDeniedTable(String),
+    QueryDenied(String),
 }
 
 #[derive(Clone, PartialEq, Debug)]

src/errors.rs

@@ -1,12 +1,17 @@
+pub mod admin;
 pub mod auth_passthrough;
+pub mod client;
 pub mod config;
 pub mod constants;
 pub mod dns_cache;
 pub mod errors;
 pub mod messages;
 pub mod mirrors;
 pub mod multi_logger;
+pub mod plugins;
 pub mod pool;
+pub mod prometheus;
+pub mod query_router;
 pub mod scram;
 pub mod server;
 pub mod sharding;

src/lib.rs

@@ -61,37 +61,19 @@ use std::str::FromStr;
 use std::sync::Arc;
 use tokio::sync::broadcast;
 
-mod admin;
-mod auth_passthrough;
-mod client;
-mod config;
-mod constants;
-mod dns_cache;
-mod errors;
-mod messages;
-mod mirrors;
-mod multi_logger;
-mod pool;
-mod prometheus;
-mod query_router;
-mod scram;
-mod server;
-mod sharding;
-mod stats;
-mod tls;
-
-use crate::config::{get_config, reload_config, VERSION};
-use crate::messages::configure_socket;
-use crate::pool::{ClientServerMap, ConnectionPool};
-use crate::prometheus::start_metric_server;
-use crate::stats::{Collector, Reporter, REPORTER};
+use pgcat::config::{get_config, reload_config, VERSION};
+use pgcat::messages::configure_socket;
+use pgcat::pool::{ClientServerMap, ConnectionPool};
+use pgcat::prometheus::start_metric_server;
+use pgcat::stats::{Collector, Reporter, REPORTER};
+use pgcat::dns_cache;
 
 fn main() -> Result<(), Box<dyn std::error::Error>> {
-    multi_logger::MultiLogger::init().unwrap();
+    pgcat::multi_logger::MultiLogger::init().unwrap();
 
     info!("Welcome to PgCat! Meow. (Version {})", VERSION);
 
-    if !query_router::QueryRouter::setup() {
+    if !pgcat::query_router::QueryRouter::setup() {
         error!("Could not setup query router");
         std::process::exit(exitcode::CONFIG);
     }
@@ -109,7 +91,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
         let runtime = Builder::new_multi_thread().worker_threads(1).build()?;
 
         runtime.block_on(async {
-            match config::parse(&config_file).await {
+            match pgcat::config::parse(&config_file).await {
                 Ok(_) => (),
                 Err(err) => {
                     error!("Config parse error: {:?}", err);
@@ -168,14 +150,14 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
         // Statistics reporting.
         REPORTER.store(Arc::new(Reporter::default()));
 
-	// Starts (if enabled) dns cache before pools initialization
-	match dns_cache::CachedResolver::from_config().await {
-            Ok(_) => (),
-            Err(err) => error!("DNS cache initialization error: {:?}", err),
-	};
+        // Starts (if enabled) dns cache before pools initialization
+        match dns_cache::CachedResolver::from_config().await {
+                Ok(_) => (),
+                Err(err) => error!("DNS cache initialization error: {:?}", err),
+        };
 
-	// Connection pool that allows to query all shards and replicas.
-	match ConnectionPool::from_config(client_server_map.clone()).await {
+        // Connection pool that allows to query all shards and replicas.
+        match ConnectionPool::from_config(client_server_map.clone()).await {
             Ok(_) => (),
             Err(err) => {
                 error!("Pool error: {:?}", err);
@@ -303,7 +285,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
                     tokio::task::spawn(async move {
                         let start = chrono::offset::Utc::now().naive_utc();
 
-                        match client::client_entrypoint(
+                        match pgcat::client::client_entrypoint(
                             socket,
                             client_server_map,
                             shutdown_rx,
@@ -334,7 +316,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
 
                             Err(err) => {
                                 match err {
-                                    errors::Error::ClientBadStartup => debug!("Client disconnected with error {:?}", err),
+                                    pgcat::errors::Error::ClientBadStartup => debug!("Client disconnected with error {:?}", err),
                                     _ => warn!("Client disconnected with error {:?}", err),
                                 }