From 9d17fd51c204af8aea0082237e72093095084567 Mon Sep 17 00:00:00 2001
From: kanpov <karpovanton729@gmail.com>
Date: Sun, 9 Feb 2025 21:47:18 +0300
Subject: [PATCH] Implement optional support for axum 0.8's Listener trait

Doing this is required due to how the orphan rule works in Rust

Signed-off-by: kanpov <karpovanton729@gmail.com>
---
 Cargo.toml          |  6 +++++-
 src/axum_support.rs | 45 +++++++++++++++++++++++++++++++++++++++++++++
 src/lib.rs          |  1 +
 3 files changed, 51 insertions(+), 1 deletion(-)
 create mode 100644 src/axum_support.rs

diff --git a/Cargo.toml b/Cargo.toml
index da7ab7c..20fb5e2 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "tokio-vsock"
-version = "0.7.0"
+version = "0.7.1"
 authors = ["fsyncd", "rust-vsock"]
 description = "Asynchronous Virtio socket support for Rust"
 repository = "https://github.com/rust-vsock/tokio-vsock"
@@ -25,6 +25,10 @@ tonic09 = { package = "tonic", version = "0.9", optional = true }
 tonic010 = { package = "tonic", version = "0.10", optional = true }
 tonic011 = { package = "tonic", version = "0.11", optional = true }
 tonic012 = { package = "tonic", version = "0.12", optional = true }
+axum08 = { package = "axum", version = "0.8", optional = true, default-features = false, features = [
+    "tokio",
+    "http1",
+] }
 
 [dev-dependencies]
 sha2 = "0.10.6"
diff --git a/src/axum_support.rs b/src/axum_support.rs
new file mode 100644
index 0000000..b7034d0
--- /dev/null
+++ b/src/axum_support.rs
@@ -0,0 +1,45 @@
+#[cfg(feature = "axum08")]
+#[cfg_attr(docsrs, doc(cfg(feature = "axum08")))]
+impl axum08::serve::Listener for crate::VsockListener {
+    type Io = crate::VsockStream;
+
+    type Addr = vsock::VsockAddr;
+
+    async fn accept(&mut self) -> (Self::Io, Self::Addr) {
+        loop {
+            match std::future::poll_fn(|cx| self.poll_accept(cx)).await {
+                Ok(tuple) => return tuple,
+                Err(err) => handle_accept_error(err).await,
+            }
+        }
+    }
+
+    fn local_addr(&self) -> std::io::Result<Self::Addr> {
+        self.local_addr()
+    }
+}
+
+#[cfg(feature = "axum08")]
+async fn handle_accept_error(err: std::io::Error) {
+    if matches!(
+        err.kind(),
+        std::io::ErrorKind::ConnectionRefused
+            | std::io::ErrorKind::ConnectionAborted
+            | std::io::ErrorKind::ConnectionReset
+    ) {
+        return;
+    }
+
+    // [From `hyper::Server` in 0.14](https://github.com/hyperium/hyper/blob/v0.14.27/src/server/tcp.rs#L186)
+    //
+    // > A possible scenario is that the process has hit the max open files
+    // > allowed, and so trying to accept a new connection will fail with
+    // > `EMFILE`. In some cases, it's preferable to just wait for some time, if
+    // > the application will likely close some files (or connections), and try
+    // > to accept the connection again. If this option is `true`, the error
+    // > will be logged at the `error` level, since it is still a big deal,
+    // > and then the listener will sleep for 1 second.
+    //
+    // hyper allowed customizing this but axum does not.
+    tokio::time::sleep(std::time::Duration::from_secs(1)).await;
+}
diff --git a/src/lib.rs b/src/lib.rs
index 13f5782..d8bf598 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -16,6 +16,7 @@
 
 #![cfg_attr(docsrs, feature(doc_cfg))]
 
+mod axum_support;
 mod listener;
 mod split;
 mod stream;