fix(query): defer MySQL session creation until the handshake completes (#18956)

* fix(query): build mysql tls config

* z

* z

* z

* z

* z

Author: everpcpc

Cargo.lock

@@ -410,7 +410,7 @@ openraft = { version = "0.10.0", features = [
     "serde",
     "tracing-log",
 ] }
-opensrv-mysql = { git = "https://github.com/databendlabs/opensrv.git", tag = "v0.9.0", features = ["tls"] }
+opensrv-mysql = { git = "https://github.com/databendlabs/opensrv.git", tag = "v0.10.0", features = ["tls"] }
 orc-rust = "0.6.0"
 ordered-float = { version = "5.0.0", default-features = false }
 ordq = "0.2.0"

Cargo.toml

@@ -16,8 +16,6 @@ mod mysql_federated;
 mod mysql_handler;
 mod mysql_interactive_worker;
 mod mysql_session;
-#[allow(clippy::unused_io_amount)]
-mod reject_connection;
 mod tls;
 mod writers;
 

src/query/service/src/servers/mysql/mod.rs

@@ -21,24 +21,18 @@ use databend_common_base::base::tokio::net::TcpStream;
 use databend_common_base::base::tokio::task::JoinHandle;
 use databend_common_base::runtime::Runtime;
 use databend_common_base::runtime::TrySpawn;
-use databend_common_catalog::session_type::SessionType;
 use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
 use futures::future::AbortHandle;
 use futures::future::AbortRegistration;
 use futures::future::Abortable;
 use futures::StreamExt;
 use log::error;
-use log::info;
-use log::warn;
-use opensrv_mysql::*;
 use rustls::ServerConfig;
-use socket2::SockRef;
 use socket2::TcpKeepalive;
 use tokio_stream::wrappers::TcpListenerStream;
 
 use crate::servers::mysql::mysql_session::MySQLConnection;
-use crate::servers::mysql::reject_connection::RejectConnection;
 use crate::servers::mysql::tls::MySQLTlsConfig;
 use crate::servers::server::ListeningStream;
 use crate::servers::server::Server;
@@ -110,50 +104,13 @@ impl MySQLHandler {
         tls: Option<Arc<ServerConfig>>,
     ) {
         executor.spawn(async move {
-            match session_manager.create_session(SessionType::MySQL).await {
-                Err(error) => {
-                    warn!("create session failed, {:?}", error);
-                    Self::reject_session(socket, error).await
-                }
-                Ok(session) => {
-                    info!("MySQL connection coming: {:?}", socket.peer_addr());
-
-                    match session_manager.register_session(session) {
-                        Ok(session) => {
-                            // TcpStream must implement AsFd for socket2 0.5, wait https://github.com/tokio-rs/tokio/pull/5514
-                            if let Err(e) = SockRef::from(&socket).set_tcp_keepalive(&keepalive) {
-                                warn!("failed to set socket option keepalive {}", e);
-                            }
-
-                            if let Err(error) = MySQLConnection::run_on_stream(session, socket, tls)
-                            {
-                                error!("Unexpected error occurred during query: {:?}", error);
-                            };
-                        }
-                        Err(error) => {
-                            warn!("fail to register session, {:?}", error);
-                            Self::reject_session(socket, error).await
-                        }
-                    }
-                }
+            if let Err(error) =
+                MySQLConnection::run_on_stream(session_manager, socket, keepalive, tls).await
+            {
+                error!("Unexpected error occurred during query: {:?}", error);
             }
         });
     }
-
-    #[async_backtrace::framed]
-    async fn reject_session(stream: TcpStream, error: ErrorCode) {
-        let (kind, message) = match error.code() {
-            41 => (ErrorKind::ER_TOO_MANY_USER_CONNECTIONS, error.message()),
-            _ => (ErrorKind::ER_INTERNAL_ERROR, error.message()),
-        };
-
-        if let Err(error) = RejectConnection::reject_mysql_connection(stream, kind, message).await {
-            error!(
-                "Unexpected error occurred during reject connection: {:?}",
-                error
-            );
-        }
-    }
 }
 
 #[async_trait::async_trait]

src/query/service/src/servers/mysql/mysql_handler.rs

@@ -48,7 +48,6 @@ use opensrv_mysql::QueryResultWriter;
 use opensrv_mysql::StatementMetaWriter;
 use rand::thread_rng;
 use rand::Rng as _;
-use rand::RngCore;
 use uuid::Uuid;
 
 use crate::auth::CredentialType;
@@ -519,23 +518,12 @@ impl InteractiveWorker {
         session: Arc<Session>,
         version: BuildInfoRef,
         client_addr: String,
+        salt: [u8; 20],
     ) -> InteractiveWorker {
-        let mut bs = vec![0u8; 20];
-        let mut rng = rand::thread_rng();
-        rng.fill_bytes(bs.as_mut());
-
-        let mut scramble: [u8; 20] = [0; 20];
-        for i in 0..20 {
-            scramble[i] = bs[i] & 0x7fu8;
-            if scramble[i] == b'\0' || scramble[i] == b'$' {
-                scramble[i] += 1;
-            }
-        }
-
         InteractiveWorker {
             version: format!("{MYSQL_VERSION}-{}", version.commit_detail),
             base: InteractiveWorkerBase { session, version },
-            salt: scramble,
+            salt,
             client_addr,
             keep_alive_task_started: false,
         }