From 5d8f5ccfab04936b714d27dd2dbfc3d0f8171c25 Mon Sep 17 00:00:00 2001 From: Joshua Ashton Date: Sun, 29 Aug 2021 08:49:30 +0100 Subject: [PATCH] Support for UDP socket forwarding Allow forwarding WS->UDP. Useful for experiments and hobby projects. --- websockify/websocketproxy.py | 11 +++++--- websockify/websockifyserver.py | 48 ++++++++++++++++++++-------------- 2 files changed, 36 insertions(+), 23 deletions(-) diff --git a/websockify/websocketproxy.py b/websockify/websocketproxy.py index 09d7882..274ee98 100644 --- a/websockify/websocketproxy.py +++ b/websockify/websocketproxy.py @@ -86,6 +86,7 @@ Traffic Legend: """ Called after a new WebSocket connection has been established. """ + use_tcp = False # Checking for a token is done in validate_connection() # Connect to the target @@ -106,15 +107,17 @@ Traffic Legend: self.server.target_port, connect=True, use_ssl=self.server.ssl_target, - unix_socket=self.server.unix_target) + unix_socket=self.server.unix_target, + use_tcp=use_tcp) except Exception as e: self.log_message("Failed to connect to %s:%s: %s", self.server.target_host, self.server.target_port, e) raise self.CClose(1011, "Failed to connect to downstream server") - self.request.setsockopt(socket.SOL_TCP, socket.TCP_NODELAY, 1) - if not self.server.wrap_cmd and not self.server.unix_target: - tsock.setsockopt(socket.SOL_TCP, socket.TCP_NODELAY, 1) + if use_tcp: + self.request.setsockopt(socket.SOL_TCP, socket.TCP_NODELAY, 1) + if not self.server.wrap_cmd and not self.server.unix_target: + tsock.setsockopt(socket.SOL_TCP, socket.TCP_NODELAY, 1) self.print_traffic(self.traffic_legend) diff --git a/websockify/websockifyserver.py b/websockify/websockifyserver.py index 0199e42..4d47e60 100644 --- a/websockify/websockifyserver.py +++ b/websockify/websockifyserver.py @@ -422,7 +422,8 @@ class WebSockifyServer(): @staticmethod def socket(host, port=None, connect=False, prefer_ipv6=False, unix_socket=None, use_ssl=False, tcp_keepalive=True, - tcp_keepcnt=None, tcp_keepidle=None, tcp_keepintvl=None): + tcp_keepcnt=None, tcp_keepidle=None, tcp_keepintvl=None, + use_tcp=True): """ Resolve a host (and optional port) to an IPv4 or IPv6 address. Create a socket. Bind to it if listen is set, otherwise connect to it. Return the socket. @@ -440,8 +441,12 @@ class WebSockifyServer(): flags = flags | socket.AI_PASSIVE if not unix_socket: - addrs = socket.getaddrinfo(host, port, 0, socket.SOCK_STREAM, - socket.IPPROTO_TCP, flags) + if use_tcp: + addrs = socket.getaddrinfo(host, port, 0, socket.SOCK_STREAM, + socket.IPPROTO_TCP, flags) + else: + addrs = socket.getaddrinfo(host, port, 0, socket.SOCK_DGRAM, + socket.IPPROTO_UDP, flags) if not addrs: raise Exception("Could not resolve host '%s'" % host) addrs.sort(key=lambda x: x[0]) @@ -449,29 +454,35 @@ class WebSockifyServer(): addrs.reverse() sock = socket.socket(addrs[0][0], addrs[0][1]) - if tcp_keepalive: - sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1) - if tcp_keepcnt: - sock.setsockopt(socket.SOL_TCP, socket.TCP_KEEPCNT, - tcp_keepcnt) - if tcp_keepidle: - sock.setsockopt(socket.SOL_TCP, socket.TCP_KEEPIDLE, - tcp_keepidle) - if tcp_keepintvl: - sock.setsockopt(socket.SOL_TCP, socket.TCP_KEEPINTVL, - tcp_keepintvl) + if use_tcp: + if tcp_keepalive: + sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1) + if tcp_keepcnt: + sock.setsockopt(socket.SOL_TCP, socket.TCP_KEEPCNT, + tcp_keepcnt) + if tcp_keepidle: + sock.setsockopt(socket.SOL_TCP, socket.TCP_KEEPIDLE, + tcp_keepidle) + if tcp_keepintvl: + sock.setsockopt(socket.SOL_TCP, socket.TCP_KEEPINTVL, + tcp_keepintvl) if connect: sock.connect(addrs[0][4]) if use_ssl: sock = ssl.wrap_socket(sock) else: - sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + if not use_tcp: + sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) sock.bind(addrs[0][4]) - sock.listen(100) + if use_tcp: + sock.listen(100) else: - sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) - sock.connect(unix_socket) + if use_tcp: + sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + sock.connect(unix_socket) + else: + sock = socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM) return sock @@ -697,7 +708,6 @@ class WebSockifyServer(): is a WebSockets client then call new_websocket_client() method (which must be overridden) for each new client connection. """ - if self.listen_fd != None: lsock = socket.fromfd(self.listen_fd, socket.AF_INET, socket.SOCK_STREAM) else: