feat: encoded secret key

Signed-off-by: Ning Sun <[email protected]>

Author: sunng87

Cargo.lock

@@ -43,6 +43,7 @@ moka = { workspace = true, features = ["future", "sync"] }
 partition.workspace = true
 paste.workspace = true
 prometheus.workspace = true
+rand.workspace = true
 rustc-hash.workspace = true
 serde_json.workspace = true
 session.workspace = true

src/catalog/Cargo.toml

@@ -15,6 +15,7 @@
 use std::collections::hash_map::Entry;
 use std::collections::HashMap;
 use std::fmt::{Debug, Formatter};
+use std::hash::{DefaultHasher, Hash, Hasher};
 use std::sync::atomic::{AtomicU32, Ordering};
 use std::sync::{Arc, RwLock};
 
@@ -101,6 +102,27 @@ impl ProcessManager {
         self.next_id.fetch_add(1, Ordering::Relaxed)
     }
 
+    /// Generate a Postgres specific secret key
+    ///
+    /// According to Postgres this secret key is a 32-bit long integer for
+    /// identify a connection. This implementation uses first 16bits hash value
+    /// of server address and 16bit random data for this key.
+    pub fn generate_secret_key(&self) -> i32 {
+        let mut hasher = DefaultHasher::new();
+        self.server_addr.hash(&mut hasher);
+        let hostname_hash = hasher.finish();
+
+        // Take lower 16 bits of the hash for first part
+        let first_part = (hostname_hash & 0xFFFF) as u16;
+
+        // Generate random 16 bits for second part
+        let second_part = rand::random::<u16>();
+
+        // Combine both parts into i32
+        let result = ((first_part as u32) << 16) | (second_part as u32);
+        result as i32
+    }
+
     /// De-register a query from process list.
     pub fn deregister_query(&self, catalog: String, id: ProcessId) {
         if let Entry::Occupied(mut o) = self.catalogs.write().unwrap().entry(catalog) {

src/catalog/src/process_manager.rs

@@ -119,12 +119,13 @@ impl PgWireServerHandlers for PostgresServerHandler {
 impl MakePostgresServerHandler {
     fn make(&self, addr: Option<SocketAddr>) -> PostgresServerHandler {
         let process_id = self.process_manager.next_id();
-        let session = Arc::new(Session::new(
-            addr,
-            Channel::Postgres,
-            Default::default(),
-            process_id,
-        ));
+        // generate cluster unique secret key
+        let secret_key = self.process_manager.generate_secret_key();
+
+        let session = Arc::new(
+            Session::new(addr, Channel::Postgres, Default::default(), process_id)
+                .with_secret_key(secret_key),
+        );
         let handler = PostgresServerHandlerInner {
             query_handler: self.query_handler.clone(),
             process_manager: self.process_manager.clone(),

src/servers/src/postgres.rs

@@ -126,7 +126,10 @@ where
 
     // pass generated process id and secret key to client, this information will
     // be sent to postgres client for query cancellation.
-    client.set_pid_and_secret_key(session.process_id() as i32, rand::random::<i32>());
+    client.set_pid_and_secret_key(
+        session.process_id() as i32,
+        session.secret_key().unwrap_or_default(),
+    );
     // set userinfo outside
 }
 

src/servers/src/postgres/auth_handler.rs

@@ -43,6 +43,8 @@ pub struct Session {
     configuration_variables: Arc<ConfigurationVariables>,
     // the process id to use when killing the query
     process_id: u32,
+    // a postgres specific key for cancel request
+    secret_key: Option<i32>,
 }
 
 pub type SessionRef = Arc<Session>;
@@ -85,9 +87,15 @@ impl Session {
             configuration_variables: Arc::new(configuration_variables),
             mutable_inner: Arc::new(RwLock::new(MutableInner::default())),
             process_id,
+            secret_key: None,
         }
     }
 
+    pub fn with_secret_key(mut self, secret_key: i32) -> Self {
+        self.secret_key = Some(secret_key);
+        self
+    }
+
     pub fn new_query_context(&self) -> QueryContextRef {
         QueryContextBuilder::default()
             // catalog is not allowed for update in query context so we use
@@ -155,4 +163,8 @@ impl Session {
     pub fn process_id(&self) -> u32 {
         self.process_id
     }
+
+    pub fn secret_key(&self) -> Option<i32> {
+        self.secret_key
+    }
 }