server: free udp hash when exit

brevent · Jun 11, 2024 · 6c44d78 · 6c44d78
1 parent 3c5b3f7
commit 6c44d78
Show file tree

Hide file tree

Showing 4 changed files with 50 additions and 24 deletions.
diff --git a/common.c b/common.c
@@ -43,9 +43,7 @@ static int is_udp(struct bufferevent *bev) {
 }
 
 static void bev_context_udp_free(void *context) {
-#ifdef WSS_PROXY_CLIENT
     lh_bev_context_udp_delete(((struct bev_context_udp *) context)->hash, context);
-#endif
     free(context);
 }
 
@@ -54,6 +52,26 @@ const struct bev_context const_bev_context_udp = {
         bev_context_udp_free,
 };
 
+static void free_udp(bev_context_udp *udp) {
+    struct bufferevent *raw, *wev, *tev;
+
+    raw = udp->bev;
+    wev = raw->cbarg;
+    tev = wev ? bufferevent_get_underlying(wev) : NULL;
+    LOGD("free udp for peer %d, raw: %p, wev: %p, tev: %p", get_peer_port(raw), raw, wev, tev);
+    if (tev) {
+        close_wss(tev, close_reason_eof, BEV_EVENT_EOF);
+    } else {
+        raw->errorcb(raw, BEV_EVENT_EOF, get_cbarg(raw));
+    }
+}
+
+void free_all_udp(LHASH_OF(bev_context_udp) *hash) {
+    lh_bev_context_udp_set_down_load(hash, 0);
+    lh_bev_context_udp_doall(hash, free_udp);
+    lh_bev_context_udp_free(hash);
+}
+
 uint16_t get_peer_port(struct bufferevent *bev) {
     evutil_socket_t sock;
     ev_socklen_t socklen;

diff --git a/common.h b/common.h
@@ -81,13 +81,11 @@ struct udp_frame {
 
 typedef struct bev_context_udp bev_context_udp;
 
-#ifdef WSS_PROXY_CLIENT
 #ifdef DEFINE_LHASH_OF_EX
 DEFINE_LHASH_OF_EX(bev_context_udp);
 #else
 DEFINE_LHASH_OF(bev_context_udp);
 #endif
-#endif
 
 struct bev_context {
     const char *name;
@@ -99,15 +97,20 @@ struct bev_context_udp {
     ev_socklen_t socklen;
     struct sockaddr *sockaddr;
     struct bufferevent *bev;
+    LHASH_OF(bev_context_udp) *hash;
 #ifdef WSS_PROXY_CLIENT
     union {
         struct sockaddr_in sin;
         struct sockaddr_in6 sin6;
     } sockaddr_storage;
-    LHASH_OF(bev_context_udp) *hash;
+#endif
+#ifdef WSS_PROXY_SERVER
+    evutil_socket_t sock;
 #endif
 };
 
+void free_all_udp(LHASH_OF(bev_context_udp) *hash);
+
 static inline void bufferevent_set_context(struct bufferevent *bev, void *context) {
     if (context == NULL || (!bev->wm_read.low && !bev->wm_write.low)) {
         bev->wm_read.low = bev->wm_write.low = (size_t) context;

diff --git a/wss-proxy-client.c b/wss-proxy-client.c
@@ -625,20 +625,6 @@ static void udp_read_cb_server(evutil_socket_t sock, short event, void *ctx) {
     }
 }
 
-static void free_udp(bev_context_udp *udp) {
-    struct bufferevent *raw, *wev, *tev;
-
-    raw = udp->bev;
-    wev = raw->cbarg;
-    tev = wev ? bufferevent_get_underlying(wev) : NULL;
-    LOGD("free udp for peer %d, raw: %p, wev: %p, tev: %p", get_peer_port(raw), raw, wev, tev);
-    if (tev) {
-        close_wss(tev, close_reason_eof, BEV_EVENT_EOF);
-    } else {
-        raw->errorcb(raw, BEV_EVENT_EOF, get_cbarg(raw));
-    }
-}
-
 static void server_context_free(const struct server_context *server_context) {
     if (server_context->listener) {
         evconnlistener_free(server_context->listener);
@@ -647,9 +633,7 @@ static void server_context_free(const struct server_context *server_context) {
         evutil_closesocket(server_context->udp_sock);
     }
     if (server_context->udp_context.hash) {
-        lh_bev_context_udp_set_down_load(server_context->udp_context.hash, 0);
-        lh_bev_context_udp_doall(server_context->udp_context.hash, free_udp);
-        lh_bev_context_udp_free(server_context->udp_context.hash);
+        free_all_udp(server_context->udp_context.hash);
     }
     if (server_context->udp_event) {
         event_free(server_context->udp_event);

diff --git a/wss-proxy-server.c b/wss-proxy-server.c
@@ -17,8 +17,17 @@ struct raw_server_info {
     struct sockaddr_storage sockaddr;
     int udp_port;
     struct sockaddr_storage udp_sockaddr;
+    LHASH_OF(bev_context_udp) *hash;
 };
 
+static unsigned long bev_context_udp_hash(const bev_context_udp *a) {
+    return a->sock;
+}
+
+static int bev_context_udp_cmp(const bev_context_udp *a, const bev_context_udp *b) {
+    return a->sock - b->sock;
+}
+
 static int init_ws_info(const char **addr, int *port) {
     int mux;
     char *end;
@@ -168,8 +177,6 @@ static struct bufferevent *init_udp_client(struct event_base *base, struct raw_s
         LOGE("cannot calloc for udp socket");
         goto error;
     }
-    bev_context_udp->socklen = raw_server_info->socklen;
-    bev_context_udp->sockaddr = (struct sockaddr *) &(raw_server_info->udp_sockaddr);
     raw = bufferevent_socket_new(base, -1, BEV_OPT_CLOSE_ON_FREE);
     if (!raw) {
         LOGE("cannot create udp bufferevent");
@@ -180,8 +187,14 @@ static struct bufferevent *init_udp_client(struct event_base *base, struct raw_s
     event_assign(&(raw->ev_write), base, sock, EV_WRITE | EV_PERSIST, bev_context_udp_writecb, raw);
     event_assign(&(raw->ev_read), base, sock, EV_READ | EV_PERSIST, udp_read_cb_client, raw);
     event_add(&(raw->ev_read), &one_minute);
+    bev_context_udp->socklen = raw_server_info->socklen;
+    bev_context_udp->sockaddr = (struct sockaddr *) &(raw_server_info->udp_sockaddr);
+    bev_context_udp->bev = raw;
+    bev_context_udp->hash = raw_server_info->hash;
+    bev_context_udp->sock = sock;
     bev_context_udp->bev_context = &const_bev_context_udp;
     bufferevent_set_context(raw, bev_context_udp);
+    lh_bev_context_udp_insert(raw_server_info->hash, bev_context_udp);
     return raw;
 error:
     if (sock > 0) {
@@ -318,6 +331,11 @@ int main() {
     if (raw_server_info.udp_port > 0) {
         memcpy(&(raw_server_info.udp_sockaddr), &(raw_server_info.sockaddr), raw_server_info.socklen);
         set_port(&(raw_server_info.udp_sockaddr), raw_server_info.udp_port);
+        raw_server_info.hash = lh_bev_context_udp_new(bev_context_udp_hash, bev_context_udp_cmp);
+        if (!raw_server_info.hash) {
+            LOGE("cannot create lhash for udp");
+            goto error;
+        }
     }
 
     if (init_ws_info(&addr, &port)) {
@@ -369,6 +387,9 @@ int main() {
 
     code = 0;
 error:
+    if (raw_server_info.hash) {
+        free_all_udp(raw_server_info.hash);
+    }
     if (event_parent) {
         event_free(event_parent);
     }