tmp

Author: yezizp2012

src/frontend/src/session.rs

@@ -1577,15 +1577,15 @@ impl SessionManagerImpl {
             }
 
             // Check HBA configuration for LDAP authentication
-            let connection_type = match peer_addr.as_ref() {
-                Address::Tcp(_) => ConnectionType::Host,
-                Address::Unix(_) => ConnectionType::Local,
-            };
-
-            let client_addr = match peer_addr.as_ref() {
-                Address::Tcp(socket_addr) => Some(&socket_addr.ip()),
-                _ => None,
+            let (connection_type, client_addr) = match peer_addr.as_ref() {
+                Address::Tcp(socket_addr) => (ConnectionType::Host, Some(&socket_addr.ip())),
+                Address::Unix(_) => (ConnectionType::Local, None),
             };
+            tracing::debug!(
+                "receive connection: type={:?}, client_addr={:?}",
+                connection_type,
+                client_addr
+            );
 
             let hba_entry_opt = self.env.frontend_config().hba_config.find_matching_entry(
                 &connection_type,

src/utils/pgwire/src/pg_server.rs

@@ -285,6 +285,35 @@ pub async fn pg_serve(
     let listener = Listener::bind(addr).await?;
     tracing::info!(addr, "server started");
 
+    // FIXME: only for debugging, remove it after hba config feature is complete.
+    // check and create another unix socket if addr is binding locally with a port
+    let extra_unix_listener = match &listener {
+        Listener::Tcp(_) => {
+            let host_and_port = addr.split(':').collect::<Vec<_>>();
+            let port = if host_and_port.len() == 2 {
+                host_and_port[1].parse::<i16>().ok()
+            } else {
+                None
+            };
+            if let Some(port) = port {
+                let local_unix_addr = format!("unix:/tmp/.s.PGSQL.{}", port);
+                match Listener::bind(&local_unix_addr).await {
+                    Ok(l) => {
+                        tracing::info!(local_unix_addr, "also bind a local unix socket");
+                        Some(l)
+                    }
+                    Err(e) => {
+                        tracing::warn!(error = %e.as_report(), local_unix_addr, "failed to bind a local unix socket");
+                        None
+                    }
+                }
+            } else {
+                None
+            }
+        }
+        Listener::Unix(_) => None,
+    };
+
     let acceptor_runtime = BackgroundShutdownRuntime::from({
         let mut builder = tokio::runtime::Builder::new_multi_thread();
         builder.worker_threads(1);
@@ -300,17 +329,20 @@ pub async fn pg_serve(
     #[cfg(madsim)]
     let worker_runtime = tokio::runtime::Builder::new_multi_thread().build().unwrap();
     let session_mgr_clone = session_mgr.clone();
+    let worker_runtime_clone = worker_runtime.clone();
+    let context_clone = context.clone();
+    let tcp_keepalive_clone = tcp_keepalive.clone();
     let f = async move {
         loop {
-            let conn_ret = listener.accept(&tcp_keepalive).await;
+            let conn_ret = listener.accept(&tcp_keepalive_clone).await;
             match conn_ret {
                 Ok((stream, peer_addr)) => {
                     tracing::info!(%peer_addr, "accept connection");
-                    worker_runtime.spawn(handle_connection(
+                    worker_runtime_clone.spawn(handle_connection(
                         stream,
                         session_mgr_clone.clone(),
                         Arc::new(peer_addr),
-                        context.clone(),
+                        context_clone.clone(),
                     ));
                 }
 
@@ -322,6 +354,31 @@ pub async fn pg_serve(
     };
     acceptor_runtime.spawn(f);
 
+    if let Some(local_unix_listener) = extra_unix_listener {
+        let session_mgr_clone = session_mgr.clone();
+        let f = async move {
+            loop {
+                let conn_ret = local_unix_listener.accept(&tcp_keepalive).await;
+                match conn_ret {
+                    Ok((stream, peer_addr)) => {
+                        tracing::info!(%peer_addr, "accept connection from local unix socket");
+                        worker_runtime.spawn(handle_connection(
+                            stream,
+                            session_mgr_clone.clone(),
+                            Arc::new(peer_addr),
+                            context.clone(),
+                        ));
+                    }
+
+                    Err(e) => {
+                        tracing::error!(error = %e.as_report(), "failed to accept connection from local unix socket",);
+                    }
+                }
+            }
+        };
+        acceptor_runtime.spawn(f);
+    }
+
     // Wait for the shutdown signal.
     shutdown.cancelled().await;