From d46a36fbf43d1540776275315596e64c8bfe60cd Mon Sep 17 00:00:00 2001
From: hev <r@hev.cc>
Date: Wed, 4 Dec 2024 21:18:01 +0800
Subject: [PATCH] HevSocks5SessionUDP: Fix timeout for bidirectional link.

---
 src/hev-socks5-session-udp.c | 37 ++++++++++++++++++++++++++++--------
 src/hev-socks5-session-udp.h |  8 ++++++++
 2 files changed, 37 insertions(+), 8 deletions(-)

diff --git a/src/hev-socks5-session-udp.c b/src/hev-socks5-session-udp.c
index 1b5fed8..c4d95cc 100644
--- a/src/hev-socks5-session-udp.c
+++ b/src/hev-socks5-session-udp.c
@@ -54,8 +54,12 @@ hev_socks5_session_udp_fwd_f (HevSocks5SessionUDP *self)
             break;
 
         res = task_io_yielder (HEV_TASK_WAITIO, self);
-        if (res < 0)
+        if (res < 0) {
+            self->alive &= ~HEV_SOCKS5_SESSION_UDP_ALIVE_F;
+            if (self->alive && hev_socks5_get_timeout (HEV_SOCKS5 (self)))
+                return 0;
             return -1;
+        }
     }
 
     frame = container_of (node, HevSocks5UDPFrame, node);
@@ -64,17 +68,23 @@ hev_socks5_session_udp_fwd_f (HevSocks5SessionUDP *self)
 
     udp = HEV_SOCKS5_UDP (self);
     res = hev_socks5_udp_sendto (udp, buf->payload, buf->len, addr);
+    hev_list_del (&self->frame_list, node);
+    hev_free (frame);
+    pbuf_free (buf);
+    self->frames--;
     if (res <= 0) {
+        if (res < -1) {
+            self->alive &= ~HEV_SOCKS5_SESSION_UDP_ALIVE_F;
+            if (self->alive && hev_socks5_get_timeout (HEV_SOCKS5 (self)))
+                return 0;
+        }
         LOG_D ("%p socks5 session udp fwd f send", self);
         res = -1;
     }
 
-    hev_list_del (&self->frame_list, node);
-    hev_free (frame);
-    pbuf_free (buf);
-    self->frames--;
+    self->alive |= HEV_SOCKS5_SESSION_UDP_ALIVE_F;
 
-    return res;
+    return 0;
 }
 
 static int
@@ -109,6 +119,11 @@ hev_socks5_session_udp_fwd_b (HevSocks5SessionUDP *self)
 
     res = hev_socks5_udp_recvfrom (udp, buf->payload, buf->len, saddr);
     if (res <= 0) {
+        if (res < -1) {
+            self->alive &= ~HEV_SOCKS5_SESSION_UDP_ALIVE_B;
+            if (self->alive && hev_socks5_get_timeout (HEV_SOCKS5 (self)))
+                return 0;
+        }
         LOG_D ("%p socks5 session udp fwd b recv", self);
         pbuf_free (buf);
         return -1;
@@ -134,6 +149,8 @@ hev_socks5_session_udp_fwd_b (HevSocks5SessionUDP *self)
         return -1;
     }
 
+    self->alive |= HEV_SOCKS5_SESSION_UDP_ALIVE_B;
+
     return 0;
 }
 
@@ -220,10 +237,10 @@ static void
 splice_task_entry (void *data)
 {
     HevTask *task = hev_task_self ();
-    HevSocks5UDP *self = data;
+    HevSocks5SessionUDP *self = data;
     int fd;
 
-    fd = hev_task_io_dup (hev_socks5_udp_get_fd (self));
+    fd = hev_task_io_dup (hev_socks5_udp_get_fd (HEV_SOCKS5_UDP (self)));
     if (fd < 0)
         return;
 
@@ -235,6 +252,7 @@ splice_task_entry (void *data)
             break;
     }
 
+    self->alive &= ~HEV_SOCKS5_SESSION_UDP_ALIVE_B;
     hev_task_del_fd (task, fd);
     close (fd);
 }
@@ -249,6 +267,8 @@ hev_socks5_session_udp_splice (HevSocks5Session *base)
 
     LOG_D ("%p socks5 session udp splice", self);
 
+    self->alive = HEV_SOCKS5_SESSION_UDP_ALIVE_F |
+                  HEV_SOCKS5_SESSION_UDP_ALIVE_B;
     fd = hev_socks5_udp_get_fd (HEV_SOCKS5_UDP (self));
     if (hev_task_mod_fd (task, fd, POLLOUT) < 0)
         hev_task_add_fd (task, fd, POLLOUT);
@@ -263,6 +283,7 @@ hev_socks5_session_udp_splice (HevSocks5Session *base)
             break;
     }
 
+    self->alive &= ~HEV_SOCKS5_SESSION_UDP_ALIVE_F;
     hev_task_join (task);
     hev_task_unref (task);
 }
diff --git a/src/hev-socks5-session-udp.h b/src/hev-socks5-session-udp.h
index 3bce45f..4c9b596 100644
--- a/src/hev-socks5-session-udp.h
+++ b/src/hev-socks5-session-udp.h
@@ -18,15 +18,23 @@
 #define HEV_SOCKS5_SESSION_UDP_CLASS(p) ((HevSocks5SessionUDPClass *)p)
 #define HEV_SOCKS5_SESSION_UDP_TYPE (hev_socks5_session_udp_class ())
 
+typedef enum _HevSocks5SessionUDPAlive HevSocks5SessionUDPAlive;
 typedef struct _HevSocks5SessionUDP HevSocks5SessionUDP;
 typedef struct _HevSocks5SessionUDPClass HevSocks5SessionUDPClass;
 
+enum _HevSocks5SessionUDPAlive
+{
+    HEV_SOCKS5_SESSION_UDP_ALIVE_F = (1 << 0),
+    HEV_SOCKS5_SESSION_UDP_ALIVE_B = (1 << 1),
+};
+
 struct _HevSocks5SessionUDP
 {
     HevSocks5ClientUDP base;
 
     HevSocks5SessionData data;
 
+    HevSocks5SessionUDPAlive alive;
     HevList frame_list;
     struct udp_pcb *pcb;
     HevTaskMutex *mutex;