switch to use common quic listen code

bytedance · Nov 20, 2023 · bacbccf · bacbccf
1 parent 692a8ba
commit bacbccf
Show file tree

Hide file tree

Showing 22 changed files with 276 additions and 583 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/g3proxy/src/serve/dummy_close/server.rs b/g3proxy/src/serve/dummy_close/server.rs
@@ -24,7 +24,9 @@ use tokio::sync::broadcast;
 use tokio_openssl::SslStream;
 use tokio_rustls::server::TlsStream;
 
-use g3_daemon::listen::{AcceptTcpServer, ArcAcceptTcpServer, ListenStats};
+use g3_daemon::listen::{
+    AcceptQuicServer, AcceptTcpServer, ArcAcceptQuicServer, ArcAcceptTcpServer, ListenStats,
+};
 use g3_daemon::server::{BaseServer, ClientConnectionInfo, ServerReloadCommand};
 use g3_types::metrics::MetricsName;
 
@@ -151,6 +153,15 @@ impl AcceptTcpServer for DummyCloseServer {
     }
 }
 
+#[async_trait]
+impl AcceptQuicServer for DummyCloseServer {
+    async fn run_quic_task(&self, _connection: Connection, _cc_info: ClientConnectionInfo) {}
+
+    fn get_reloaded(&self) -> ArcAcceptQuicServer {
+        crate::serve::get_or_insert_default(self.config.name())
+    }
+}
+
 #[async_trait]
 impl Server for DummyCloseServer {
     fn escaper(&self) -> &MetricsName {
@@ -187,6 +198,4 @@ impl Server for DummyCloseServer {
         _cc_info: ClientConnectionInfo,
     ) {
     }
-
-    async fn run_quic_task(&self, _connection: Connection, _cc_info: ClientConnectionInfo) {}
 }
diff --git a/g3proxy/src/serve/http_proxy/server.rs b/g3proxy/src/serve/http_proxy/server.rs
@@ -30,7 +30,10 @@ use tokio::sync::{broadcast, mpsc};
 use tokio_openssl::SslStream;
 use tokio_rustls::{server::TlsStream, TlsAcceptor};
 
-use g3_daemon::listen::{AcceptTcpServer, ArcAcceptTcpServer, ListenStats, ListenTcpRuntime};
+use g3_daemon::listen::{
+    AcceptQuicServer, AcceptTcpServer, ArcAcceptQuicServer, ArcAcceptTcpServer, ListenStats,
+    ListenTcpRuntime,
+};
 use g3_daemon::server::{BaseServer, ClientConnectionInfo, ServerReloadCommand};
 use g3_types::acl::{AclAction, AclNetworkRule};
 use g3_types::acl_set::AclDstHostRuleSet;
@@ -385,6 +388,38 @@ impl AcceptTcpServer for HttpProxyServer {
     }
 }
 
+#[async_trait]
+impl AcceptQuicServer for HttpProxyServer {
+    async fn run_quic_task(&self, connection: Connection, cc_info: ClientConnectionInfo) {
+        let client_addr = cc_info.client_addr();
+        self.server_stats.add_conn(client_addr);
+        if self.drop_early(client_addr) {
+            return;
+        }
+
+        loop {
+            // TODO update ctx and quit gracefully
+            match connection.accept_bi().await {
+                Ok((send_stream, recv_stream)) => {
+                    self.spawn_quic_stream_task(send_stream, recv_stream, cc_info.clone())
+                }
+                Err(e) => {
+                    debug!(
+                        "{} - {} quic connection error: {e:?}",
+                        cc_info.sock_local_addr(),
+                        cc_info.sock_peer_addr()
+                    );
+                    break;
+                }
+            }
+        }
+    }
+
+    fn get_reloaded(&self) -> ArcAcceptQuicServer {
+        crate::serve::get_or_insert_default(self.config.name())
+    }
+}
+
 #[async_trait]
 impl Server for HttpProxyServer {
     fn escaper(&self) -> &MetricsName {
@@ -435,29 +470,4 @@ impl Server for HttpProxyServer {
 
         self.spawn_stream_task(stream, cc_info).await;
     }
-
-    async fn run_quic_task(&self, connection: Connection, cc_info: ClientConnectionInfo) {
-        let client_addr = cc_info.client_addr();
-        self.server_stats.add_conn(client_addr);
-        if self.drop_early(client_addr) {
-            return;
-        }
-
-        loop {
-            // TODO update ctx and quit gracefully
-            match connection.accept_bi().await {
-                Ok((send_stream, recv_stream)) => {
-                    self.spawn_quic_stream_task(send_stream, recv_stream, cc_info.clone())
-                }
-                Err(e) => {
-                    debug!(
-                        "{} - {} quic connection error: {e:?}",
-                        cc_info.sock_local_addr(),
-                        cc_info.sock_peer_addr()
-                    );
-                    break;
-                }
-            }
-        }
-    }
 }
diff --git a/g3proxy/src/serve/http_rproxy/server.rs b/g3proxy/src/serve/http_rproxy/server.rs
@@ -31,7 +31,10 @@ use tokio_openssl::SslStream;
 use tokio_rustls::server::TlsStream;
 use tokio_rustls::LazyConfigAcceptor;
 
-use g3_daemon::listen::{AcceptTcpServer, ArcAcceptTcpServer, ListenStats, ListenTcpRuntime};
+use g3_daemon::listen::{
+    AcceptQuicServer, AcceptTcpServer, ArcAcceptQuicServer, ArcAcceptTcpServer, ListenStats,
+    ListenTcpRuntime,
+};
 use g3_daemon::server::{BaseServer, ClientConnectionInfo, ServerReloadCommand};
 use g3_types::acl::{AclAction, AclNetworkRule};
 use g3_types::metrics::MetricsName;
@@ -402,6 +405,15 @@ impl AcceptTcpServer for HttpRProxyServer {
     }
 }
 
+#[async_trait]
+impl AcceptQuicServer for HttpRProxyServer {
+    async fn run_quic_task(&self, _connection: Connection, _cc_info: ClientConnectionInfo) {}
+
+    fn get_reloaded(&self) -> ArcAcceptQuicServer {
+        crate::serve::get_or_insert_default(self.config.name())
+    }
+}
+
 #[async_trait]
 impl Server for HttpRProxyServer {
     fn escaper(&self) -> &MetricsName {
@@ -452,6 +464,4 @@ impl Server for HttpRProxyServer {
 
         self.spawn_stream_task(stream, cc_info).await;
     }
-
-    async fn run_quic_task(&self, _connection: Connection, _cc_info: ClientConnectionInfo) {}
 }
diff --git a/g3proxy/src/serve/intelli_proxy/server.rs b/g3proxy/src/serve/intelli_proxy/server.rs
@@ -26,7 +26,10 @@ use tokio::sync::broadcast;
 use tokio_openssl::SslStream;
 use tokio_rustls::server::TlsStream;
 
-use g3_daemon::listen::{AcceptTcpServer, ArcAcceptTcpServer, ListenStats, ListenTcpRuntime};
+use g3_daemon::listen::{
+    AcceptQuicServer, AcceptTcpServer, ArcAcceptQuicServer, ArcAcceptTcpServer, ListenStats,
+    ListenTcpRuntime,
+};
 use g3_daemon::server::{BaseServer, ClientConnectionInfo, ServerReloadCommand};
 use g3_io_ext::haproxy::{ProxyProtocolV1Reader, ProxyProtocolV2Reader};
 use g3_types::acl::{AclAction, AclNetworkRule};
@@ -266,6 +269,15 @@ impl AcceptTcpServer for IntelliProxy {
     }
 }
 
+#[async_trait]
+impl AcceptQuicServer for IntelliProxy {
+    async fn run_quic_task(&self, _connection: Connection, _cc_info: ClientConnectionInfo) {}
+
+    fn get_reloaded(&self) -> ArcAcceptQuicServer {
+        crate::serve::get_or_insert_default(self.config.name())
+    }
+}
+
 #[async_trait]
 impl Server for IntelliProxy {
     fn escaper(&self) -> &MetricsName {
@@ -302,6 +314,4 @@ impl Server for IntelliProxy {
         _cc_info: ClientConnectionInfo,
     ) {
     }
-
-    async fn run_quic_task(&self, _connection: Connection, _cc_info: ClientConnectionInfo) {}
 }
diff --git a/g3proxy/src/serve/mod.rs b/g3proxy/src/serve/mod.rs
@@ -17,13 +17,12 @@
 use std::sync::Arc;
 
 use async_trait::async_trait;
-use quinn::Connection;
 use tokio::net::TcpStream;
 use tokio::sync::broadcast;
 use tokio_openssl::SslStream;
 use tokio_rustls::server::TlsStream;
 
-use g3_daemon::listen::{AcceptTcpServer, ListenStats};
+use g3_daemon::listen::{AcceptQuicServer, AcceptTcpServer, ListenStats};
 use g3_daemon::server::{BaseServer, ClientConnectionInfo, ServerQuitPolicy, ServerReloadCommand};
 use g3_types::metrics::MetricsName;
 
@@ -35,9 +34,6 @@ pub(crate) use registry::{foreach_online as foreach_server, get_names, get_or_in
 mod idle_check;
 pub(crate) use idle_check::ServerIdleChecker;
 
-mod runtime;
-use runtime::ListenQuicRuntime;
-
 mod dummy_close;
 mod intelli_proxy;
 mod native_tls_port;
@@ -90,7 +86,9 @@ pub(crate) trait ServerInternal {
 }
 
 #[async_trait]
-pub(crate) trait Server: ServerInternal + BaseServer + AcceptTcpServer {
+pub(crate) trait Server:
+    ServerInternal + BaseServer + AcceptTcpServer + AcceptQuicServer
+{
     fn escaper(&self) -> &MetricsName;
     fn user_group(&self) -> &MetricsName;
     fn auditor(&self) -> &MetricsName;
@@ -106,8 +104,6 @@ pub(crate) trait Server: ServerInternal + BaseServer + AcceptTcpServer {
     async fn run_rustls_task(&self, stream: TlsStream<TcpStream>, cc_info: ClientConnectionInfo);
 
     async fn run_openssl_task(&self, stream: SslStream<TcpStream>, cc_info: ClientConnectionInfo);
-
-    async fn run_quic_task(&self, connection: Connection, cc_info: ClientConnectionInfo);
 }
 
 pub(crate) type ArcServer = Arc<dyn Server + Send + Sync>;

diff --git a/g3proxy/src/serve/native_tls_port/mod.rs b/g3proxy/src/serve/native_tls_port/mod.rs
@@ -29,7 +29,10 @@ use tokio::sync::broadcast;
 use tokio_openssl::SslStream;
 use tokio_rustls::server::TlsStream;
 
-use g3_daemon::listen::{AcceptTcpServer, ArcAcceptTcpServer, ListenStats, ListenTcpRuntime};
+use g3_daemon::listen::{
+    AcceptQuicServer, AcceptTcpServer, ArcAcceptQuicServer, ArcAcceptTcpServer, ListenStats,
+    ListenTcpRuntime,
+};
 use g3_daemon::server::{BaseServer, ClientConnectionInfo, ServerReloadCommand};
 use g3_io_ext::haproxy::{ProxyProtocolV1Reader, ProxyProtocolV2Reader};
 use g3_types::acl::{AclAction, AclNetworkRule};
@@ -282,6 +285,15 @@ impl AcceptTcpServer for NativeTlsPort {
     }
 }
 
+#[async_trait]
+impl AcceptQuicServer for NativeTlsPort {
+    async fn run_quic_task(&self, _connection: Connection, _cc_info: ClientConnectionInfo) {}
+
+    fn get_reloaded(&self) -> ArcAcceptQuicServer {
+        crate::serve::get_or_insert_default(self.config.name())
+    }
+}
+
 #[async_trait]
 impl Server for NativeTlsPort {
     fn escaper(&self) -> &MetricsName {
@@ -318,6 +330,4 @@ impl Server for NativeTlsPort {
         _cc_info: ClientConnectionInfo,
     ) {
     }
-
-    async fn run_quic_task(&self, _connection: Connection, _cc_info: ClientConnectionInfo) {}
 }
diff --git a/g3proxy/src/serve/plain_quic_port/mod.rs b/g3proxy/src/serve/plain_quic_port/mod.rs
@@ -26,15 +26,18 @@ use tokio::sync::{broadcast, watch};
 use tokio_openssl::SslStream;
 use tokio_rustls::server::TlsStream;
 
-use g3_daemon::listen::{AcceptTcpServer, ArcAcceptTcpServer, ListenStats, ListenQuicConf};
+use g3_daemon::listen::{
+    AcceptQuicServer, AcceptTcpServer, ArcAcceptQuicServer, ArcAcceptTcpServer, ListenQuicConf,
+    ListenQuicRuntime, ListenStats,
+};
 use g3_daemon::server::{BaseServer, ClientConnectionInfo, ServerReloadCommand};
 use g3_types::acl::AclNetworkRule;
 use g3_types::metrics::MetricsName;
 use g3_types::net::UdpListenConfig;
 
 use crate::config::server::plain_quic_port::{PlainQuicPortConfig, PlainQuicPortUpdateFlags};
 use crate::config::server::{AnyServerConfig, ServerConfig};
-use crate::serve::{ArcServer, ListenQuicRuntime, Server, ServerInternal, ServerQuitPolicy};
+use crate::serve::{ArcServer, Server, ServerInternal, ServerQuitPolicy};
 
 #[derive(Clone)]
 struct PlainQuicPortAuxConfig {
@@ -229,7 +232,11 @@ impl ServerInternal for PlainQuicPort {
 
     fn _start_runtime(&self, server: &ArcServer) -> anyhow::Result<()> {
         let config = self.config.load();
-        let runtime = ListenQuicRuntime::new(server, config.as_ref(), config.listen.clone());
+        let runtime = ListenQuicRuntime::new(
+            server.clone(),
+            server.get_listen_stats(),
+            config.listen.clone(),
+        );
         runtime.run_all_instances(
             config.listen_in_worker,
             &self.quinn_config,
@@ -269,6 +276,18 @@ impl AcceptTcpServer for PlainQuicPort {
     }
 }
 
+#[async_trait]
+impl AcceptQuicServer for PlainQuicPort {
+    async fn run_quic_task(&self, connection: Connection, cc_info: ClientConnectionInfo) {
+        let next_server = self.next_server.load().as_ref().clone();
+        next_server.run_quic_task(connection, cc_info).await
+    }
+
+    fn get_reloaded(&self) -> ArcAcceptQuicServer {
+        crate::serve::get_or_insert_default(&self.name)
+    }
+}
+
 #[async_trait]
 impl Server for PlainQuicPort {
     fn escaper(&self) -> &MetricsName {
@@ -305,9 +324,4 @@ impl Server for PlainQuicPort {
         _cc_info: ClientConnectionInfo,
     ) {
     }
-
-    async fn run_quic_task(&self, connection: Connection, cc_info: ClientConnectionInfo) {
-        let next_server = self.next_server.load().as_ref().clone();
-        next_server.run_quic_task(connection, cc_info).await;
-    }
 }
diff --git a/g3proxy/src/serve/plain_tcp_port/mod.rs b/g3proxy/src/serve/plain_tcp_port/mod.rs
@@ -26,7 +26,10 @@ use tokio::sync::broadcast;
 use tokio_openssl::SslStream;
 use tokio_rustls::server::TlsStream;
 
-use g3_daemon::listen::{AcceptTcpServer, ArcAcceptTcpServer, ListenStats, ListenTcpRuntime};
+use g3_daemon::listen::{
+    AcceptQuicServer, AcceptTcpServer, ArcAcceptQuicServer, ArcAcceptTcpServer, ListenStats,
+    ListenTcpRuntime,
+};
 use g3_daemon::server::{BaseServer, ClientConnectionInfo, ServerReloadCommand};
 use g3_io_ext::haproxy::{ProxyProtocolV1Reader, ProxyProtocolV2Reader};
 use g3_types::acl::{AclAction, AclNetworkRule};
@@ -232,6 +235,15 @@ impl AcceptTcpServer for PlainTcpPort {
     }
 }
 
+#[async_trait]
+impl AcceptQuicServer for PlainTcpPort {
+    async fn run_quic_task(&self, _connection: Connection, _cc_info: ClientConnectionInfo) {}
+
+    fn get_reloaded(&self) -> ArcAcceptQuicServer {
+        crate::serve::get_or_insert_default(self.config.name())
+    }
+}
+
 #[async_trait]
 impl Server for PlainTcpPort {
     fn escaper(&self) -> &MetricsName {
@@ -268,6 +280,4 @@ impl Server for PlainTcpPort {
         _cc_info: ClientConnectionInfo,
     ) {
     }
-
-    async fn run_quic_task(&self, _connection: Connection, _cc_info: ClientConnectionInfo) {}
 }