🔨 refactor the listening server

.github/workflows/release.yml

@@ -74,6 +74,8 @@ jobs:
           target: ${{ matrix.target }}
           override: true
 
+      - uses: extractions/setup-just@v1
+
       - name: cargo install patch-crate
         uses: baptiste0928/cargo-install@v2
         with:

.github/workflows/test.yml

@@ -36,7 +36,9 @@ jobs:
           toolchain: stable
           target: ${{ matrix.target }}
           override: true
-
+
+      - uses: extractions/setup-just@v1
+
       - name: cargo install patch-crate
         uses: baptiste0928/cargo-install@v2
         with:
@@ -76,6 +78,8 @@ jobs:
           target: ${{ matrix.target }}
           override: true
 
+      - uses: extractions/setup-just@v1
+
       - name: cargo install patch-crate
         uses: baptiste0928/cargo-install@v2
         with:
@@ -155,6 +159,8 @@ jobs:
           target: ${{ matrix.target }}
           override: true
 
+      - uses: extractions/setup-just@v1
+
       - name: cargo install patch-crate
         uses: baptiste0928/cargo-install@v2
         with:

Cargo.toml

@@ -41,6 +41,8 @@ dns-over-https-rustls = [
 
 dnssec = ["hickory-proto/dnssec"]
 
+legacy_dns_server = ["dep:hickory-server"]
+
 experimental = ["experimental-trie", "experimental-phf"]
 
 experimental-trie = []
@@ -58,7 +60,7 @@ hickory-proto = { path = "./target/patch/hickory-proto-0.24.0" }
 # rustls-native-certs = { git = "https://github.com/mokeyish/rustls-native-certs.git" }
 hostname = { git = "https://github.com/mokeyish/hostname.git", branch = "dev" }
 # enum_dispatch = { git = "https://gitlab.com/mokeyish/enum_dispatch.git", branch = "master"}
-axum = { git = "https://github.com/tokio-rs/axum.git", rev = "8854e66" }
+# axum = { git = "https://github.com/tokio-rs/axum.git", rev = "8854e66" }
 # boomphf ={ path = "./target/patch/boomphf-0.6.0"}
 
 [dependencies]
@@ -73,8 +75,8 @@ enum_dispatch = "0.3.12"
 
 
 # api
-axum = { version = "0.6.16" }
-axum-server = { version = "0.5.1", features = ["tls-rustls"] }
+axum = { version = "0.7.2" }
+axum-server = { version = "0.6.0", features = ["tls-rustls"] }
 
 # serde
 serde = { version = "1.0", features = ["derive"]}
@@ -94,6 +96,7 @@ tokio = { version = "1.28", features = [
   "parking_lot",
 ] }
 tokio-rustls = "0.24.0"
+tokio-util = "0.7.10"
 socket2 = { version = "0.5", features = ["all"] }
 reqwest = { version = "0.11", default-features = false, features = [
   "blocking",
@@ -115,7 +118,7 @@ hickory-resolver = { version = "0.24", features = [
   "serde-config",
   "system-config",
 ] }
-hickory-server = { version = "0.24", features = ["resolver"] }
+hickory-server = { version = "0.24", features = ["resolver"], optional = true }
 
 # ssl
 webpki-roots = "0.25.2"

justfile

@@ -0,0 +1,24 @@
+
+
+build: init
+  cargo build -r
+
+# Run tests
+test: init
+  cargo test
+
+# Run clippy
+clippy: init
+  cargo clippy --fix --all 
+
+# Check the format
+fmt: init
+  cargo fmt --all
+
+apply-patch:
+  cargo patch-crate
+
+# Initialize all tools needed
+init:
+  @cargo install patch-crate -q
+

src/api/mod.rs

@@ -1,10 +1,11 @@
 use std::{io, sync::Arc};
 
 use axum::{routing::get, Json, Router};
-use axum_server::tls_rustls::RustlsConfig;
+use axum_server::{tls_rustls::RustlsConfig, Handle};
 use rustls::{Certificate, PrivateKey};
 use serde::{Deserialize, Serialize};
 use tokio::net::TcpListener;
+use tokio_util::sync::CancellationToken;
 
 mod address;
 mod audit;
@@ -16,24 +17,26 @@ mod nameserver;
 mod serve_dns;
 mod settings;
 
-use crate::{app::App, dns_server::DnsServerHandler};
+use crate::{app::App, server::DnsHandle};
 
 type StatefulRouter = Router<Arc<ServeState>>;
 
 pub struct ServeState {
     app: Arc<App>,
-    dns_handler: DnsServerHandler,
+    dns_handle: DnsHandle,
 }
 
-pub async fn register_https(
+pub async fn serve(
     app: Arc<App>,
-    dns_handler: DnsServerHandler,
+    dns_handle: DnsHandle,
     tcp_listener: TcpListener,
     certificate: Vec<Certificate>,
     certificate_key: PrivateKey,
-    handle: axum_server::Handle,
-) -> io::Result<()> {
-    let state = Arc::new(ServeState { app, dns_handler });
+) -> io::Result<CancellationToken> {
+    let token = CancellationToken::new();
+    let cancellation_token = token.clone();
+
+    let state = Arc::new(ServeState { app, dns_handle });
 
     let app = Router::new()
         .merge(serve_dns::routes())
@@ -43,15 +46,33 @@ pub async fn register_https(
     let certificate = certificate.into_iter().map(|c| c.0).collect::<Vec<_>>();
     let certificate_key = certificate_key.0;
 
-    axum_server::from_tcp_rustls(
-        tcp_listener.into_std()?,
-        RustlsConfig::from_der(certificate, certificate_key).await?,
-    )
-    .handle(handle)
-    .serve(app.into_make_service())
-    .await?;
-
-    Ok(())
+    let tcp_listener = tcp_listener.into_std()?;
+    let rustls_config = RustlsConfig::from_der(certificate, certificate_key).await?;
+
+    tokio::spawn(async move {
+        use crate::log;
+        let shutdown_handle = Handle::new();
+
+        tokio::select! {
+            result = axum_server::from_tcp_rustls(
+                tcp_listener,
+                rustls_config,
+            )
+            .handle(shutdown_handle.clone())
+            .serve(app.into_make_service()) => match result {
+                Ok(()) => (),
+                Err(e) => {
+                    log::debug!("error receiving quic connection: {e}");
+                }
+            },
+            _ = cancellation_token.cancelled() => {
+                // A graceful shutdown was initiated. Break out of the loop.
+                shutdown_handle.graceful_shutdown(Some(std::time::Duration::from_secs(5)))
+            },
+        };
+    });
+
+    Ok(token)
 }
 
 fn api_routes() -> StatefulRouter {

src/api/serve_dns.rs

@@ -8,7 +8,7 @@ use axum::{
 };
 
 use super::{ServeState, StatefulRouter};
-use crate::libdns::{proto::xfer::SerialMessage, server::server::Protocol};
+use crate::{dns::SerialMessage, libdns::Protocol};
 
 pub fn routes() -> StatefulRouter {
     Router::new().route("/dns-query", any(serve_dns))
@@ -24,15 +24,10 @@ async fn serve_dns(State(state): State<Arc<ServeState>>, req: Request) -> Bytes
     println!("{}", s.join("\n"));
 
     if let Ok(bytes) = Bytes::from_request(req, &state).await {
-        state
-            .dns_handler
-            .handle(
-                SerialMessage::new(bytes.into(), "0.0.0.0:0".parse().unwrap()),
-                Protocol::Https,
-            )
-            .await
-            .into_parts()
-            .0
+        let req_msg =
+            SerialMessage::binary(bytes.into(), "0.0.0.0:0".parse().unwrap(), Protocol::Https);
+        let res_msg = state.dns_handle.send(req_msg).await;
+        res_msg.message
     } else {
         Default::default()
     }