From aaaf660e9d1c380697e6c30046b30bcbd893e03f Mon Sep 17 00:00:00 2001 From: pi Date: Mon, 6 Apr 2026 21:28:00 +1200 Subject: [PATCH] feat(driver): discover controller via udp and rediscover on reconnect Made-with: Cursor --- src/hello.py | 162 ++++++++++++++++++++++++++++++++++++++++++++++++ src/main.py | 27 ++++++-- src/settings.py | 1 - 3 files changed, 184 insertions(+), 6 deletions(-) create mode 100644 src/hello.py diff --git a/src/hello.py b/src/hello.py new file mode 100644 index 0000000..1c6e16e --- /dev/null +++ b/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) diff --git a/src/main.py b/src/main.py index 9cd8568..e183a5e 100644 --- a/src/main.py +++ b/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() diff --git a/src/settings.py b/src/settings.py index c74c8a2..f28d6bc 100644 --- a/src/settings.py +++ b/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: