From db50d96503b62365d34a2a929746256a2233638e Mon Sep 17 00:00:00 2001
From: ngthhu <hungnt.code@proton.me>
Date: Fri, 3 May 2024 09:05:03 +0700
Subject: [PATCH] refactor: [#682] handle UDP Tracker timeouts

---
 src/console/clients/udp/checker.rs           |  2 ++
 src/shared/bit_torrent/tracker/udp/client.rs | 11 ++++++-----
 2 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/src/console/clients/udp/checker.rs b/src/console/clients/udp/checker.rs
index 9b2a9011..23690891 100644
--- a/src/console/clients/udp/checker.rs
+++ b/src/console/clients/udp/checker.rs
@@ -20,6 +20,8 @@ pub enum ClientError {
     NotConnected,
     #[error("Unexpected response while connecting the the remote server.")]
     UnexpectedConnectionResponse,
+    #[error("Connection timeout.")]
+    ConnectionTimeout,
 }
 
 /// A UDP Tracker client to make test requests (checks).
diff --git a/src/shared/bit_torrent/tracker/udp/client.rs b/src/shared/bit_torrent/tracker/udp/client.rs
index 9af9571b..a124e181 100644
--- a/src/shared/bit_torrent/tracker/udp/client.rs
+++ b/src/shared/bit_torrent/tracker/udp/client.rs
@@ -4,12 +4,13 @@ use std::net::SocketAddr;
 use std::sync::Arc;
 use std::time::Duration;
 
-use anyhow::{anyhow, Context, Result};
+use anyhow::{anyhow, bail, Context, Result};
 use aquatic_udp_protocol::{ConnectRequest, Request, Response, TransactionId};
 use log::debug;
 use tokio::net::UdpSocket;
 use tokio::time;
 
+use crate::console::clients::udp::checker::ClientError;
 use crate::shared::bit_torrent::tracker::udp::{source_address, MAX_PACKET_SIZE};
 
 /// Default timeout for sending and receiving packets. And waiting for sockets
@@ -78,7 +79,7 @@ impl UdpClient {
                     Err(e) => return Err(anyhow!("IO error waiting for the socket to become readable: {e:?}")),
                 };
             }
-            Err(e) => return Err(anyhow!("Timeout waiting for the socket to become readable: {e:?}")),
+            Err(_) => bail!(ClientError::ConnectionTimeout),
         };
 
         match time::timeout(self.timeout, self.socket.send(bytes)).await {
@@ -86,7 +87,7 @@ impl UdpClient {
                 Ok(size) => Ok(size),
                 Err(e) => Err(anyhow!("IO error during send: {e:?}")),
             },
-            Err(e) => Err(anyhow!("Send operation timed out: {e:?}")),
+            Err(_) => bail!(ClientError::ConnectionTimeout),
         }
     }
 
@@ -109,7 +110,7 @@ impl UdpClient {
                     Err(e) => return Err(anyhow!("IO error waiting for the socket to become readable: {e:?}")),
                 };
             }
-            Err(e) => return Err(anyhow!("Timeout waiting for the socket to become readable: {e:?}")),
+            Err(_) => bail!(ClientError::ConnectionTimeout),
         };
 
         let size_result = match time::timeout(self.timeout, self.socket.recv(bytes)).await {
@@ -117,7 +118,7 @@ impl UdpClient {
                 Ok(size) => Ok(size),
                 Err(e) => Err(anyhow!("IO error during send: {e:?}")),
             },
-            Err(e) => Err(anyhow!("Receive operation timed out: {e:?}")),
+            Err(_) => bail!(ClientError::ConnectionTimeout),
         };
 
         if size_result.is_ok() {