Add tokio feature flag for async runtime

rust-bitcoin · Oct 30, 2024 · b04d54c · b04d54c
1 parent d9135e0
commit b04d54c
Show file tree

Hide file tree

Showing 6 changed files with 19 additions and 30 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/protocol/Cargo.toml b/protocol/Cargo.toml
@@ -11,11 +11,13 @@ rust-version = "1.63.0"
 [features]
 default = ["std"]
 async = ["std", "futures/std"]
+tokio = ["std", "tokio/io-util"]
 std = ["alloc", "bitcoin/std", "rand/std", "rand/std_rng"]
 alloc = []
 
 [dependencies]
 futures = { version = "0.3.30", default-features = false, optional = true }
+tokio = { version = "1.37.0", default-features = false, optional = true }
 rand = { version = "0.8.0", default-features = false }
 bitcoin = { version = "0.32.4", default-features = false }
 

diff --git a/protocol/README.md b/protocol/README.md
@@ -12,7 +12,8 @@ The lower-level `Handshake` and `PacketHandler` types can be directly used by ap
 
 * `alloc` -- Expose memory allocation dependent features.
 * `std`   -- Includes the `alloc` memory allocation feature as well as extra standard library dependencies for I/O and random number generators.
-* `async` -- High level wrappers for asynchronous read and write runtimes.
+* `async` -- High level wrappers for asynchronous read and write runtimes using agnostic futures-rs traits.
+* `tokio` -- Same wrappers as `async`, but using the popular tokio runtime's specific traits instead of futures-rs.
 
 ## Handshake
 

diff --git a/protocol/src/lib.rs b/protocol/src/lib.rs
@@ -47,8 +47,13 @@ use bitcoin::{
     },
 };
 use fschacha20poly1305::{FSChaCha20, FSChaCha20Poly1305};
+// Default to the futures-rs traits, but can overwrite with more specific
+// tokio implemenations for easier caller integration.
 #[cfg(feature = "async")]
 use futures::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
+#[cfg(feature = "tokio")]
+use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
+
 use hkdf::Hkdf;
 use rand::Rng;
 
@@ -1064,7 +1069,7 @@ impl fmt::Display for ProtocolError {
 }
 
 /// A protocol session with handshake and send/receive packet management.
-#[cfg(feature = "async")]
+#[cfg(any(feature = "async", feature = "tokio"))]
 pub struct AsyncProtocol<R, W>
 where
     R: AsyncRead + Unpin + Send,
@@ -1074,7 +1079,7 @@ where
     writer: AsyncProtocolWriter<W>,
 }
 
-#[cfg(feature = "async")]
+#[cfg(any(feature = "async", feature = "tokio"))]
 impl<R, W> AsyncProtocol<R, W>
 where
     R: AsyncRead + Unpin + Send,
@@ -1213,7 +1218,7 @@ where
 }
 
 /// State machine of an asynchronous packet read.
-#[cfg(feature = "async")]
+#[cfg(any(feature = "async", feature = "tokio"))]
 #[derive(Debug)]
 enum DecryptState {
     ReadingLength {
@@ -1226,7 +1231,7 @@ enum DecryptState {
     },
 }
 
-#[cfg(feature = "async")]
+#[cfg(any(feature = "async", feature = "tokio"))]
 impl Default for DecryptState {
     fn default() -> Self {
         DecryptState::ReadingLength {
@@ -1237,7 +1242,7 @@ impl Default for DecryptState {
 }
 
 /// Manages an async buffer to automatically decrypt contents of received packets.
-#[cfg(feature = "async")]
+#[cfg(any(feature = "async", feature = "tokio"))]
 pub struct AsyncProtocolReader<R>
 where
     R: AsyncRead + Unpin + Send,
@@ -1247,7 +1252,7 @@ where
     state: DecryptState,
 }
 
-#[cfg(feature = "async")]
+#[cfg(any(feature = "async", feature = "tokio"))]
 impl<R> AsyncProtocolReader<R>
 where
     R: AsyncRead + Unpin + Send,
@@ -1298,7 +1303,7 @@ where
 }
 
 /// Manages an async buffer to automatically encrypt and send contents in packets.
-#[cfg(feature = "async")]
+#[cfg(any(feature = "async", feature = "tokio"))]
 pub struct AsyncProtocolWriter<W>
 where
     W: AsyncWrite + Unpin + Send,
@@ -1307,7 +1312,7 @@ where
     packet_writer: PacketWriter,
 }
 
-#[cfg(feature = "async")]
+#[cfg(any(feature = "async", feature = "tokio"))]
 impl<W> AsyncProtocolWriter<W>
 where
     W: AsyncWrite + Unpin + Send,

diff --git a/proxy/Cargo.toml b/proxy/Cargo.toml
@@ -14,9 +14,8 @@ spec = "config_spec.toml"
 [dependencies]
 bitcoin = { version = "0.32.0" }
 tokio = { version = "1.37.0", features = ["full"] }
-tokio-util = { version = "0.7.12", features = ["compat"] }
 hex = { package = "hex-conservative", version = "0.2.0" }
-bip324 = { path = "../protocol", features = ["async"] }
+bip324 = { path = "../protocol", features = ["tokio"] }
 configure_me = "0.4.0"
 log = "0.4"
 env_logger = "0.10"

diff --git a/proxy/src/bin/proxy.rs b/proxy/src/bin/proxy.rs
@@ -13,7 +13,6 @@ use tokio::{
     net::{TcpListener, TcpStream},
     select,
 };
-use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt};
 
 configure_me::include_config!();
 
@@ -73,9 +72,6 @@ async fn v2_proxy(
 
     info!("Initiating handshake.");
     let (remote_reader, remote_writer) = remote.into_split();
-    // Convert to futures-compatible types.
-    let remote_reader = remote_reader.compat();
-    let remote_writer = remote_writer.compat_write();
 
     let protocol = match AsyncProtocol::new(
         network,