server.py

"""
Tiny relay server with LAN auto-discovery.

Run:
    python server.py
"""

import socket
import threading
import struct
import pickle
import time

RELAY_HOST = "0.0.0.0"
RELAY_PORT = 9000
BROADCAST_PORT = 54545
BROADCAST_TOKEN = b"VIDEO_RELAY"
BROADCAST_INTERVAL = 1.0


def recv_exact(sock, size: int):
    """Receive exactly size bytes or return None."""
    data = b""
    while len(data) < size:
        chunk = sock.recv(size - len(data))
        if not chunk:
            return None
        data += chunk
    return data


def start_beacon(relay_port: int):
    def _beacon():
        sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
        msg = BROADCAST_TOKEN + b" " + str(relay_port).encode()
        while True:
            try:
                sock.sendto(msg, ("255.255.255.255", BROADCAST_PORT))
            except Exception:
                pass
            time.sleep(BROADCAST_INTERVAL)

    t = threading.Thread(target=_beacon, daemon=True)
    t.start()


class RelayServer:
    def __init__(self, host=RELAY_HOST, port=RELAY_PORT):
        self.host = host
        self.port = port
        self.running = False
        self.server_socket = None
        self.clients = []
        self.lock = threading.Lock()

    def broadcast(self, sender, payload: bytes):
        header = struct.pack("!I", len(payload))
        with self.lock:
            targets = [c for c in self.clients if c is not sender]
        for c in targets:
            try:
                c.sendall(header + payload)
            except Exception:
                pass

    def handle(self, sock, addr):
        try:
            while self.running:
                raw = recv_exact(sock, 4)
                if not raw:
                    break
                size = struct.unpack("!I", raw)[0]
                data = recv_exact(sock, size)
                if not data:
                    break
                self.broadcast(sock, data)
        finally:
            with self.lock:
                if sock in self.clients:
                    self.clients.remove(sock)
            try:
                sock.close()
            except:
                pass

    def start(self):
        self.server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self.server_socket.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
        self.server_socket.bind((self.host, self.port))
        self.server_socket.listen(10)
        self.running = True
        print(f"Relay listening on {self.host}:{self.port}")
        start_beacon(self.port)
        print(f"Broadcasting beacon on UDP 255.255.255.255:{BROADCAST_PORT}")
        try:
            while True:
                client, addr = self.server_socket.accept()
                client.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
                with self.lock:
                    self.clients.append(client)
                threading.Thread(target=self.handle, args=(client, addr), daemon=True).start()
        except KeyboardInterrupt:
            pass
        finally:
            self.stop()

    def stop(self):
        self.running = False
        if self.server_socket:
            try:
                self.server_socket.close()
            except:
                pass
        with self.lock:
            sockets = list(self.clients)
            self.clients.clear()
        for s in sockets:
            try:
                s.close()
            except:
                pass
        print("Relay stopped.")


if __name__ == "__main__":
    RelayServer().start()