diff --git a/src/server_udp_relay.c b/src/server_udp_relay.c index 3ca2ce5..ad201fb 100644 --- a/src/server_udp_relay.c +++ b/src/server_udp_relay.c @@ -4,6 +4,7 @@ #include #define UDP_RELAY_BUF_SIZE (64U * 1024U) +#define UDP_RELAY_ROUTE_TIMEOUT_MS 30000U struct udp_relay { int downstream_fd; @@ -15,6 +16,7 @@ struct udp_relay { struct sockaddr_storage client_addr; socklen_t client_addr_len; int has_client; + uint32_t client_last_seen_ms; struct udp_relay_route *routes; pthread_mutex_t lock; pthread_mutex_t log_mu; @@ -35,8 +37,17 @@ typedef struct udp_relay_route { uint32_t conv; struct sockaddr_storage client_addr; socklen_t client_addr_len; + uint32_t last_seen_ms; } udp_relay_route_t; +static uint32_t udp_relay_now_ms(void) { + return omni_now_millis32(); +} + +static uint32_t udp_relay_elapsed_ms(uint32_t now_ms, uint32_t then_ms) { + return now_ms - then_ms; +} + static void udp_relay_parse_kcp_summary(const uint8_t *packet, size_t len, int *has_conv, uint32_t *conv, size_t *segment_count) { size_t offset = 0; size_t count = 0; @@ -144,22 +155,61 @@ static void udp_relay_record_client(udp_relay_t *relay, const struct sockaddr_st memcpy(&relay->client_addr, addr, sizeof(*addr)); relay->client_addr_len = addr_len; relay->has_client = 1; + relay->client_last_seen_ms = udp_relay_now_ms(); pthread_mutex_unlock(&relay->lock); } +static void udp_relay_prune_routes_locked(udp_relay_t *relay, uint32_t now_ms) { + udp_relay_route_t *prev = NULL; + udp_relay_route_t *route; + + if (relay == NULL) { + return; + } + + route = relay->routes; + while (route != NULL) { + udp_relay_route_t *next = route->next; + + if (udp_relay_elapsed_ms(now_ms, route->last_seen_ms) >= UDP_RELAY_ROUTE_TIMEOUT_MS) { + if (prev == NULL) { + relay->routes = next; + } else { + prev->next = next; + } + free(route); + route = next; + continue; + } + + prev = route; + route = next; + } + + if (relay->has_client && udp_relay_elapsed_ms(now_ms, relay->client_last_seen_ms) >= UDP_RELAY_ROUTE_TIMEOUT_MS) { + relay->has_client = 0; + relay->client_addr_len = 0; + memset(&relay->client_addr, 0, sizeof(relay->client_addr)); + } +} + static int udp_relay_record_route(udp_relay_t *relay, uint32_t conv, const struct sockaddr_storage *addr, socklen_t addr_len) { udp_relay_route_t *route; + uint32_t now_ms; if (relay == NULL || addr == NULL || addr_len == 0) { errno = EINVAL; return -1; } + now_ms = udp_relay_now_ms(); pthread_mutex_lock(&relay->lock); + udp_relay_prune_routes_locked(relay, now_ms); for (route = relay->routes; route != NULL; route = route->next) { if (route->conv == conv) { memcpy(&route->client_addr, addr, sizeof(*addr)); route->client_addr_len = addr_len; + route->last_seen_ms = now_ms; pthread_mutex_unlock(&relay->lock); return 0; } @@ -173,6 +223,7 @@ static int udp_relay_record_route(udp_relay_t *relay, uint32_t conv, const struc route->conv = conv; memcpy(&route->client_addr, addr, sizeof(*addr)); route->client_addr_len = addr_len; + route->last_seen_ms = now_ms; route->next = relay->routes; relay->routes = route; pthread_mutex_unlock(&relay->lock); @@ -181,8 +232,11 @@ static int udp_relay_record_route(udp_relay_t *relay, uint32_t conv, const struc static int udp_relay_copy_client(udp_relay_t *relay, struct sockaddr_storage *addr, socklen_t *addr_len) { int has_client; + uint32_t now_ms; + now_ms = udp_relay_now_ms(); pthread_mutex_lock(&relay->lock); + udp_relay_prune_routes_locked(relay, now_ms); has_client = relay->has_client; if (has_client) { memcpy(addr, &relay->client_addr, sizeof(*addr)); @@ -194,8 +248,11 @@ static int udp_relay_copy_client(udp_relay_t *relay, struct sockaddr_storage *ad static int udp_relay_copy_route(udp_relay_t *relay, uint32_t conv, struct sockaddr_storage *addr, socklen_t *addr_len) { udp_relay_route_t *route; + uint32_t now_ms; + now_ms = udp_relay_now_ms(); pthread_mutex_lock(&relay->lock); + udp_relay_prune_routes_locked(relay, now_ms); for (route = relay->routes; route != NULL; route = route->next) { if (route->conv == conv) { memcpy(addr, &route->client_addr, sizeof(*addr)); @@ -305,8 +362,11 @@ static void *udp_relay_forward_upstream_to_downstream(void *arg) { udp_relay_print_packet(relay, "relay_upstream_rx", relay->upstream_local_addr, &relay->upstream_addr, relay->upstream_addr_len, buffer, (size_t) n); udp_relay_parse_kcp_summary(buffer, (size_t) n, &has_conv, &conv, NULL); - if ((has_conv && !udp_relay_copy_route(relay, conv, &client_addr, &client_addr_len)) && - !udp_relay_copy_client(relay, &client_addr, &client_addr_len)) { + if (has_conv && !udp_relay_copy_route(relay, conv, &client_addr, &client_addr_len)) { + udp_relay_print_packet(relay, "relay_upstream_drop_unknown_conv", relay->upstream_local_addr, &relay->upstream_addr, relay->upstream_addr_len, buffer, (size_t) n); + continue; + } + if (!has_conv && !udp_relay_copy_client(relay, &client_addr, &client_addr_len)) { udp_relay_print_packet(relay, "relay_upstream_drop_no_client", relay->upstream_local_addr, &relay->upstream_addr, relay->upstream_addr_len, buffer, (size_t) n); continue; }