From a611a72c48567dd19e897fde48a4813fbe254f71 Mon Sep 17 00:00:00 2001
From: jess <jess@else-if.org>
Date: Thu, 2 Apr 2026 16:10:59 -0700
Subject: [PATCH] clean up UDP sessions when WiFi station disconnects

---
 main/wifi_transport.c | 40 ++++++++++++++++++++++++++++++++++++----
 1 file changed, 36 insertions(+), 4 deletions(-)

diff --git a/main/wifi_transport.c b/main/wifi_transport.c
index e633569..d221578 100644
--- a/main/wifi_transport.c
+++ b/main/wifi_transport.c
@@ -6,6 +6,7 @@
 #include "freertos/task.h"
 #include "freertos/queue.h"
 #include "esp_wifi.h"
+#include "esp_wifi_ap_get_sta_list.h"
 #include "esp_netif.h"
 #include "esp_event.h"
 #include "lwip/sockets.h"
@@ -21,10 +22,12 @@
 #define CLIENT_TIMEOUT_MS  30000
 
 static int udp_sock = -1;
+static esp_netif_t *ap_netif;
 
 static struct {
     struct sockaddr_in addr;
     TickType_t last_seen;
+    uint8_t mac[6];
     bool active;
 } clients[MAX_UDP_CLIENTS];
 
@@ -46,6 +49,20 @@ static void client_touch(const struct sockaddr_in *addr)
         clients[client_count].addr = *addr;
         clients[client_count].last_seen = now;
         clients[client_count].active = true;
+        memset(clients[client_count].mac, 0, 6);
+
+        wifi_sta_list_t sta_list;
+        wifi_sta_mac_ip_list_t ip_list;
+        if (esp_wifi_ap_get_sta_list(&sta_list) == ESP_OK &&
+            esp_wifi_ap_get_sta_list_with_ip(&sta_list, &ip_list) == ESP_OK) {
+            for (int j = 0; j < ip_list.num; j++) {
+                if (ip_list.sta[j].ip.addr == addr->sin_addr.s_addr) {
+                    memcpy(clients[client_count].mac, ip_list.sta[j].mac, 6);
+                    break;
+                }
+            }
+        }
+
         client_count++;
         printf("UDP: client added (%d/%d)\n", client_count, MAX_UDP_CLIENTS);
     }
@@ -66,6 +83,18 @@ static void clients_expire(void)
     }
 }
 
+static void client_remove_by_mac(const uint8_t *mac)
+{
+    for (int i = 0; i < client_count; ) {
+        if (memcmp(clients[i].mac, mac, 6) == 0) {
+            clients[i] = clients[--client_count];
+            printf("UDP: client removed by MAC (%d/%d)\n", client_count, MAX_UDP_CLIENTS);
+        } else {
+            i++;
+        }
+    }
+}
+
 static void parse_udp_sysex(const uint8_t *data, uint16_t len)
 {
     if (len < 3 || data[0] != 0xF0 || data[1] != 0x7D)
@@ -219,12 +248,15 @@ int wifi_get_client_count(void)
 static void wifi_event_handler(void *arg, esp_event_base_t base,
                                int32_t id, void *data)
 {
-    (void)arg; (void)data;
+    (void)arg;
     if (base == WIFI_EVENT) {
-        if (id == WIFI_EVENT_AP_STACONNECTED)
+        if (id == WIFI_EVENT_AP_STACONNECTED) {
             printf("WiFi: station connected\n");
-        else if (id == WIFI_EVENT_AP_STADISCONNECTED)
+        } else if (id == WIFI_EVENT_AP_STADISCONNECTED) {
+            wifi_event_ap_stadisconnected_t *evt = data;
             printf("WiFi: station disconnected\n");
+            client_remove_by_mac(evt->mac);
+        }
     }
 }
 
@@ -250,7 +282,7 @@ static void sta_event_handler(void *arg, esp_event_base_t base,
 
 static int wifi_ap_init(void)
 {
-    esp_netif_create_default_wifi_ap();
+    ap_netif = esp_netif_create_default_wifi_ap();
     esp_netif_create_default_wifi_sta();
 
     wifi_init_config_t wifi_cfg = WIFI_INIT_CONFIG_DEFAULT();