From 8dae1126d5f5699c20c2624a30bd493a71dc9d0a Mon Sep 17 00:00:00 2001
From: Jacob Kiers <code@kiers.eu>
Date: Wed, 4 Oct 2023 20:50:40 +0200
Subject: [PATCH] Deduplicate copy method

Signed-off-by: Jacob Kiers <code@kiers.eu>
---
 src/servers/mod.rs          | 16 ++++++++++++++++
 src/servers/protocol/kcp.rs | 16 +---------------
 src/servers/protocol/tcp.rs | 18 ++----------------
 3 files changed, 19 insertions(+), 31 deletions(-)

diff --git a/src/servers/mod.rs b/src/servers/mod.rs
index 0161e68..ec6d32f 100644
--- a/src/servers/mod.rs
+++ b/src/servers/mod.rs
@@ -2,6 +2,8 @@ use log::{error, info};
 use std::collections::{HashMap, HashSet};
 use std::net::SocketAddr;
 use std::sync::Arc;
+use tokio::io;
+use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt};
 use tokio::task::JoinHandle;
 
 mod protocol;
@@ -210,3 +212,17 @@ mod tests {
         // conn.shutdown().await.unwrap();
     }
 }
+
+async fn copy<'a, R, W>(reader: &'a mut R, writer: &'a mut W) -> io::Result<u64>
+where
+    R: AsyncRead + Unpin + ?Sized,
+    W: AsyncWrite + Unpin + ?Sized,
+{
+    match io::copy(reader, writer).await {
+        Ok(u64) => {
+            let _ = writer.shutdown().await;
+            Ok(u64)
+        }
+        Err(_) => Ok(0),
+    }
+}
diff --git a/src/servers/protocol/kcp.rs b/src/servers/protocol/kcp.rs
index 5cc4698..eaba266 100644
--- a/src/servers/protocol/kcp.rs
+++ b/src/servers/protocol/kcp.rs
@@ -1,6 +1,6 @@
 use crate::config::Upstream;
 use crate::plugins::kcp::{KcpConfig, KcpListener, KcpStream};
-use crate::servers::Proxy;
+use crate::servers::{copy, Proxy};
 use futures::future::try_join;
 use log::{debug, error, warn};
 use std::net::SocketAddr;
@@ -96,17 +96,3 @@ async fn process(
     };
     Ok(())
 }
-
-async fn copy<'a, R, W>(reader: &'a mut R, writer: &'a mut W) -> io::Result<u64>
-where
-    R: AsyncRead + Unpin + ?Sized,
-    W: AsyncWrite + Unpin + ?Sized,
-{
-    match io::copy(reader, writer).await {
-        Ok(u64) => {
-            let _ = writer.shutdown().await;
-            Ok(u64)
-        }
-        Err(_) => Ok(0),
-    }
-}
diff --git a/src/servers/protocol/tcp.rs b/src/servers/protocol/tcp.rs
index f4565f6..e193a4a 100644
--- a/src/servers/protocol/tcp.rs
+++ b/src/servers/protocol/tcp.rs
@@ -1,11 +1,11 @@
 use crate::config::Upstream;
 use crate::servers::protocol::tls::get_sni;
-use crate::servers::Proxy;
+use crate::servers::{copy, Proxy};
 use futures::future::try_join;
 use log::{debug, error, info, warn};
 use std::sync::Arc;
 use tokio::io;
-use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt};
+use tokio::io::AsyncWriteExt;
 use tokio::net::{TcpListener, TcpStream};
 
 pub(crate) async fn proxy(config: Arc<Proxy>) -> Result<(), Box<dyn std::error::Error>> {
@@ -120,17 +120,3 @@ async fn process(
     };
     Ok(())
 }
-
-async fn copy<'a, R, W>(reader: &'a mut R, writer: &'a mut W) -> io::Result<u64>
-where
-    R: AsyncRead + Unpin + ?Sized,
-    W: AsyncWrite + Unpin + ?Sized,
-{
-    match io::copy(reader, writer).await {
-        Ok(u64) => {
-            let _ = writer.shutdown().await;
-            Ok(u64)
-        }
-        Err(_) => Ok(0),
-    }
-}