feat(driver): discover controller via udp and rediscover on reconnect

Made-with: Cursor

src/hello.py

@@ -0,0 +1,162 @@
+"""LED hello payload and UDP broadcast discovery (controller IP via echo on port 8766).
+
+Wi-Fi must already be connected; this module does not use Settings or call connect().
+"""
+
+import json
+import socket
+import ubinascii
+
+import network
+
+# Match led-controller/tests/udp_server.py
+DISCOVERY_UDP_PORT = 8766
+DEFAULT_RECV_TIMEOUT_S = 3
+
+
+def pack_hello_dict(sta, device_name=""):
+    """Same fields as main HTTP/ESP-NOW hello."""
+    mac = sta.config("mac")
+    return {
+        "v": "1",
+        "device_name": device_name,
+        "mac": ubinascii.hexlify(mac).decode().lower(),
+        "type": "led",
+    }
+
+
+def pack_hello_bytes(sta, device_name=""):
+    return json.dumps(pack_hello_dict(sta, device_name)).encode("utf-8")
+
+
+def pack_hello_line(sta, device_name=""):
+    """JSON hello + newline (HTTP/UDP discovery payloads)."""
+    return pack_hello_bytes(sta, device_name) + b"\n"
+
+
+def ipv4_broadcast(ip, netmask):
+    """Directed broadcast (e.g. 192.168.1.0/24 -> 192.168.1.255)."""
+    ia = [int(x) for x in ip.split(".")]
+    im = [int(x) for x in netmask.split(".")]
+    if len(ia) != 4 or len(im) != 4:
+        return None
+    return ".".join(str(ia[i] | (255 - im[i])) for i in range(4))
+
+
+def udp_discovery_targets(ip, mask):
+    """(directed_broadcast, port) then limited broadcast."""
+    out = [("255.255.255.255", DISCOVERY_UDP_PORT)]
+    b = ipv4_broadcast(ip, mask)
+    if b:
+        out.insert(0, (b, DISCOVERY_UDP_PORT))
+    return out
+
+
+def broadcast_hello_udp(
+    sta,
+    device_name="",
+    *,
+    wait_reply=True,
+    recv_timeout_s=DEFAULT_RECV_TIMEOUT_S,
+    wdt=None,
+):
+    """
+    Send pack_hello_line via directed then 255.255.255.255 on DISCOVERY_UDP_PORT.
+    STA must already be connected with a valid IPv4 (caller brings up Wi-Fi).
+
+    If wait_reply, wait for first UDP echo. Returns controller IP string or None.
+    """
+    ip, mask, _gw, _dns = sta.ifconfig()
+    msg = pack_hello_line(sta, device_name)
+    print("hello:", msg)
+
+    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+    try:
+        sock.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
+    except (AttributeError, OSError) as e:
+        print("SO_BROADCAST not set:", e)
+    try:
+        sock.bind((ip, 0))
+    except (AttributeError, OSError, TypeError) as e:
+        try:
+            sock.bind(("0.0.0.0", 0))
+        except (AttributeError, OSError) as e2:
+            print("bind skipped:", e, e2)
+    if wait_reply:
+        try:
+            sock.settimeout(recv_timeout_s)
+        except (AttributeError, OSError):
+            pass
+
+    discovered = None
+    for dest_ip, dest_port in udp_discovery_targets(ip, mask):
+        if wdt is not None:
+            wdt.feed()
+        label = "%s:%s" % (dest_ip, dest_port)
+        target = (dest_ip, dest_port)
+        try:
+            sock.sendto(msg, target)
+            print("sent hello ->", target)
+        except OSError as e:
+            print("sendto failed:", e)
+            continue
+        if not wait_reply:
+            continue
+        if wdt is not None:
+            wdt.feed()
+        try:
+            data, addr = sock.recvfrom(2048)
+            print("reply from", addr, ":", data)
+            remote_ip = addr[0]
+            if data != msg:
+                print("(warning: reply payload differs from hello; still using source IP.)")
+            discovered = remote_ip
+            print("Discovered controller at", remote_ip)
+            break
+        except OSError as e:
+            print("recv (no reply):", e, "via", label)
+            if dest_ip == "255.255.255.255":
+                print(
+                    "(hint: many APs drop Wi-Fi client broadcast; try wired server or AP without client isolation.)"
+                )
+
+    sock.close()
+    return discovered
+
+
+def discover_controller_udp(device_name="", wdt=None):
+    """
+    Broadcast hello; return controller IP from first UDP echo, or None.
+    STA must already be connected.
+
+    device_name: logical name in the JSON (caller supplies, e.g. from Settings elsewhere).
+    wdt: optional WDT to feed during waits.
+    """
+    sta = network.WLAN(network.STA_IF)
+    if not sta.isconnected():
+        print("hello: STA not connected — connect Wi-Fi before discovery.")
+        raise SystemExit(1)
+
+    ip, mask, _g, _d = sta.ifconfig()
+    if ip == "0.0.0.0":
+        print("hello: STA has no IP address.")
+        raise SystemExit(1)
+
+    print("STA IP:", ip, "mask:", mask)
+
+    discovered = broadcast_hello_udp(
+        sta,
+        device_name,
+        wait_reply=True,
+        wdt=wdt,
+    )
+    if discovered:
+        print("discover done; controller =", repr(discovered))
+    else:
+        print("discover done; controller not found")
+    return discovered
+
+
+if __name__ == "__main__":
+    if not discover_controller_udp():
+        raise SystemExit(1)

src/main.py

@@ -10,6 +10,7 @@ import time
 import select
 import socket
 import ubinascii
+from hello import discover_controller_udp
 
 BROADCAST_MAC = b"\xff\xff\xff\xff\xff\xff"
 CONTROLLER_TCP_PORT = 8765
@@ -122,14 +123,13 @@ sta_if.active(True)
 sta_if.config(pm=network.WLAN.PM_NONE)
 
 mac = sta_if.config("mac")
-hello = {
+hello_payload = {
     "v": "1",
     "device_name": settings.get("name", ""),
     "mac": ubinascii.hexlify(mac).decode().lower(),
     "type": "led",
 }
-hello_bytes = json.dumps(hello).encode("utf-8")
-hello_line = hello_bytes + b"\n"
+hello_bytes = json.dumps(hello_payload).encode("utf-8")
 
 if settings["transport_type"] == "espnow":
     sta_if.disconnect()
@@ -152,6 +152,22 @@ elif settings["transport_type"] == "wifi":
     while not sta_if.isconnected():
         time.sleep(1)
     print(f"WiFi connected {sta_if.ifconfig()[0]}")
+    controller_ip = discover_controller_udp(
+        device_name=settings.get("name", ""),
+        wdt=wdt,
+    )
+    if not controller_ip:
+        raise SystemExit("No controller IP discovered for Wi-Fi transport")
+
+    def pick_controller_ip(current):
+        ip = discover_controller_udp(
+            device_name=settings.get("name", ""),
+            wdt=wdt,
+        )
+        if ip and ip != current:
+            print("Controller IP updated to", ip)
+        return ip if ip else current
+
     reconnect_ms = 1000
     next_connect_at = 0
     client = None
@@ -165,8 +181,7 @@ elif settings["transport_type"] == "wifi":
             c = None
             try:
                 c = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-                c.connect((settings["server_ip"], CONTROLLER_TCP_PORT))
-                c.send(hello_line)
+                c.connect((controller_ip, CONTROLLER_TCP_PORT))
                 c.setblocking(False)
                 p = select.poll()
                 p.register(c, select.POLLIN)
@@ -180,6 +195,7 @@ elif settings["transport_type"] == "wifi":
                         c.close()
                     except Exception:
                         pass
+                controller_ip = pick_controller_ip(controller_ip)
                 next_connect_at = utime.ticks_add(now, reconnect_ms)
 
         if client is not None and poller is not None:
@@ -221,6 +237,7 @@ elif settings["transport_type"] == "wifi":
                 client = None
                 poller = None
                 buf = b""
+                controller_ip = pick_controller_ip(controller_ip)
                 next_connect_at = utime.ticks_add(now, reconnect_ms)
 
         presets.tick()

src/settings.py

@@ -26,7 +26,6 @@ class Settings(dict):
         # Wi-Fi + TCP to Pi: leave ssid empty for ESP-NOW-only.
         self["ssid"] = ""
         self["password"] = ""
-        self["server_ip"] = ""
 
     def save(self):
         try: