Skip to content

Commit

Permalink
Split out basic WebSocket server template
Browse files Browse the repository at this point in the history
  • Loading branch information
@ossman
ossman committed Feb 1, 2017
1 parent 8a69762 commit e47591f
Show file tree
Hide file tree
Showing 7 changed files with 899 additions and 827 deletions.
6 changes: 3 additions & 3 deletions tests/echo.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,9 +12,9 @@

import os, sys, select, optparse, logging
sys.path.insert(0,os.path.join(os.path.dirname(__file__), ".."))
from websockify.websocketserver import WebSocketServer, WebSocketRequestHandler
from websockify.websockifyserver import WebSockifyServer, WebSockifyRequestHandler

class WebSocketEcho(WebSocketRequestHandler):
class WebSocketEcho(WebSockifyRequestHandler):
"""
WebSockets server that echos back whatever is received from the
client. """
Expand Down Expand Up @@ -71,6 +71,6 @@ def new_websocket_client(self):
logging.basicConfig(level=logging.INFO)

opts.web = "."
server = WebSocketServer(WebSocketEcho, **opts.__dict__)
server = WebSockifyServer(WebSocketEcho, **opts.__dict__)
server.start_server()

8 changes: 4 additions & 4 deletions tests/load.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,20 +8,20 @@

import sys, os, select, random, time, optparse, logging
sys.path.insert(0,os.path.join(os.path.dirname(__file__), ".."))
from websockify.websocketserver import WebSocketServer, WebSocketRequestHandler
from websockify.websockifyserver import WebSockifyServer, WebSockifyRequestHandler

class WebSocketLoadServer(WebSocketServer):
class WebSocketLoadServer(WebSockifyServer):

recv_cnt = 0
send_cnt = 0

def __init__(self, *args, **kwargs):
self.delay = kwargs.pop('delay')

WebSocketServer.__init__(self, *args, **kwargs)
WebSockifyServer.__init__(self, *args, **kwargs)


class WebSocketLoad(WebSocketRequestHandler):
class WebSocketLoad(WebSockifyRequestHandler):

max_packet_size = 10000

Expand Down
4 changes: 2 additions & 2 deletions tests/test_websocketproxy.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@

from mox3 import stubout

from websockify import websocketserver
from websockify import websockifyserver
from websockify import websocketproxy
from websockify import token_plugins
from websockify import auth_plugins
Expand Down Expand Up @@ -75,7 +75,7 @@ def setUp(self):
FakeSocket(''), "127.0.0.1", FakeServer())
self.handler.path = "https://localhost:6080/websockify?token=blah"
self.handler.headers = None
self.stubs.Set(websocketserver.WebSocketServer, 'socket',
self.stubs.Set(websockifyserver.WebSockifyServer, 'socket',
staticmethod(lambda *args, **kwargs: None))

def tearDown(self):
Expand Down
60 changes: 30 additions & 30 deletions tests/test_websocketserver.py → tests/test_websockifyserver.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@
# License for the specific language governing permissions and limitations
# under the License.

""" Unit tests for websocketserver """
""" Unit tests for websockifyserver """
import errno
import os
import logging
Expand All @@ -28,12 +28,12 @@
import unittest
import socket
import signal
from websockify import websocketserver
from websockify import websockifyserver

try:
from SimpleHTTPServer import SimpleHTTPRequestHandler
from BaseHTTPServer import BaseHTTPRequestHandler
except ImportError:
from http.server import SimpleHTTPRequestHandler
from http.server import BaseHTTPRequestHandler

try:
from StringIO import StringIO
Expand Down Expand Up @@ -70,64 +70,64 @@ def makefile(self, mode='r', buffsize=None):
return StringIO(self._data.decode('latin_1'))


class WebSocketRequestHandlerTestCase(unittest.TestCase):
class WebSockifyRequestHandlerTestCase(unittest.TestCase):
def setUp(self):
super(WebSocketRequestHandlerTestCase, self).setUp()
super(WebSockifyRequestHandlerTestCase, self).setUp()
self.stubs = stubout.StubOutForTesting()
self.tmpdir = tempfile.mkdtemp('-websockify-tests')
# Mock this out cause it screws tests up
self.stubs.Set(os, 'chdir', lambda *args, **kwargs: None)
self.stubs.Set(SimpleHTTPRequestHandler, 'send_response',
self.stubs.Set(BaseHTTPRequestHandler, 'send_response',
lambda *args, **kwargs: None)

def tearDown(self):
"""Called automatically after each test."""
self.stubs.UnsetAll()
os.rmdir(self.tmpdir)
super(WebSocketRequestHandlerTestCase, self).tearDown()
super(WebSockifyRequestHandlerTestCase, self).tearDown()

def _get_server(self, handler_class=websocketserver.WebSocketRequestHandler,
def _get_server(self, handler_class=websockifyserver.WebSockifyRequestHandler,
**kwargs):
web = kwargs.pop('web', self.tmpdir)
return websocketserver.WebSocketServer(
return websockifyserver.WebSockifyServer(
handler_class, listen_host='localhost',
listen_port=80, key=self.tmpdir, web=web,
record=self.tmpdir, daemon=False, ssl_only=0, idle_timeout=1,
**kwargs)

def test_normal_get_with_only_upgrade_returns_error(self):
server = self._get_server(web=None)
handler = websocketserver.WebSocketRequestHandler(
handler = websockifyserver.WebSockifyRequestHandler(
FakeSocket('GET /tmp.txt HTTP/1.1'), '127.0.0.1', server)

def fake_send_response(self, code, message=None):
self.last_code = code

self.stubs.Set(SimpleHTTPRequestHandler, 'send_response',
self.stubs.Set(BaseHTTPRequestHandler, 'send_response',
fake_send_response)

handler.do_GET()
self.assertEqual(handler.last_code, 405)

def test_list_dir_with_file_only_returns_error(self):
server = self._get_server(file_only=True)
handler = websocketserver.WebSocketRequestHandler(
handler = websockifyserver.WebSockifyRequestHandler(
FakeSocket('GET / HTTP/1.1'), '127.0.0.1', server)

def fake_send_response(self, code, message=None):
self.last_code = code

self.stubs.Set(SimpleHTTPRequestHandler, 'send_response',
self.stubs.Set(BaseHTTPRequestHandler, 'send_response',
fake_send_response)

handler.path = '/'
handler.do_GET()
self.assertEqual(handler.last_code, 404)


class WebSocketServerTestCase(unittest.TestCase):
class WebSockifyServerTestCase(unittest.TestCase):
def setUp(self):
super(WebSocketServerTestCase, self).setUp()
super(WebSockifyServerTestCase, self).setUp()
self.stubs = stubout.StubOutForTesting()
self.tmpdir = tempfile.mkdtemp('-websockify-tests')
# Mock this out cause it screws tests up
Expand All @@ -137,11 +137,11 @@ def tearDown(self):
"""Called automatically after each test."""
self.stubs.UnsetAll()
os.rmdir(self.tmpdir)
super(WebSocketServerTestCase, self).tearDown()
super(WebSockifyServerTestCase, self).tearDown()

def _get_server(self, handler_class=websocketserver.WebSocketRequestHandler,
def _get_server(self, handler_class=websockifyserver.WebSockifyRequestHandler,
**kwargs):
return websocketserver.WebSocketServer(
return websockifyserver.WebSockifyServer(
handler_class, listen_host='localhost',
listen_port=80, key=self.tmpdir, web=self.tmpdir,
record=self.tmpdir, **kwargs)
Expand Down Expand Up @@ -174,7 +174,7 @@ def fake_select(rlist, wlist, xlist, timeout=None):

self.stubs.Set(select, 'select', fake_select)
self.assertRaises(
websocketserver.WebSocketServer.EClose, server.do_handshake,
websockifyserver.WebSockifyServer.EClose, server.do_handshake,
FakeSocket(), '127.0.0.1')

def test_empty_handshake_fails(self):
Expand All @@ -187,7 +187,7 @@ def fake_select(rlist, wlist, xlist, timeout=None):

self.stubs.Set(select, 'select', fake_select)
self.assertRaises(
websocketserver.WebSocketServer.EClose, server.do_handshake,
websockifyserver.WebSockifyServer.EClose, server.do_handshake,
sock, '127.0.0.1')

def test_handshake_policy_request(self):
Expand All @@ -204,7 +204,7 @@ def fake_select(rlist, wlist, xlist, timeout=None):

self.stubs.Set(select, 'select', fake_select)
self.assertRaises(
websocketserver.WebSocketServer.EClose, server.do_handshake,
websockifyserver.WebSockifyServer.EClose, server.do_handshake,
sock, '127.0.0.1')

def test_do_handshake_no_ssl(self):
Expand Down Expand Up @@ -247,7 +247,7 @@ def fake_select(rlist, wlist, xlist, timeout=None):

self.stubs.Set(select, 'select', fake_select)
self.assertRaises(
websocketserver.WebSocketServer.EClose, server.do_handshake,
websockifyserver.WebSockifyServer.EClose, server.do_handshake,
sock, '127.0.0.1')

def test_do_handshake_ssl_error_eof_raises_close_error(self):
Expand All @@ -264,7 +264,7 @@ def fake_wrap_socket(*args, **kwargs):
self.stubs.Set(select, 'select', fake_select)
self.stubs.Set(ssl, 'wrap_socket', fake_wrap_socket)
self.assertRaises(
websocketserver.WebSocketServer.EClose, server.do_handshake,
websockifyserver.WebSockifyServer.EClose, server.do_handshake,
sock, '127.0.0.1')

def test_fallback_sigchld_handler(self):
Expand All @@ -278,9 +278,9 @@ def test_start_server_error(self):
def fake_select(rlist, wlist, xlist, timeout=None):
raise Exception("fake error")

self.stubs.Set(websocketserver.WebSocketServer, 'socket',
self.stubs.Set(websockifyserver.WebSockifyServer, 'socket',
lambda *args, **kwargs: sock)
self.stubs.Set(websocketserver.WebSocketServer, 'daemonize',
self.stubs.Set(websockifyserver.WebSockifyServer, 'daemonize',
lambda *args, **kwargs: None)
self.stubs.Set(select, 'select', fake_select)
server.start_server()
Expand All @@ -292,9 +292,9 @@ def test_start_server_keyboardinterrupt(self):
def fake_select(rlist, wlist, xlist, timeout=None):
raise KeyboardInterrupt

self.stubs.Set(websocketserver.WebSocketServer, 'socket',
self.stubs.Set(websockifyserver.WebSockifyServer, 'socket',
lambda *args, **kwargs: sock)
self.stubs.Set(websocketserver.WebSocketServer, 'daemonize',
self.stubs.Set(websockifyserver.WebSockifyServer, 'daemonize',
lambda *args, **kwargs: None)
self.stubs.Set(select, 'select', fake_select)
server.start_server()
Expand All @@ -306,9 +306,9 @@ def test_start_server_systemexit(self):
def fake_select(rlist, wlist, xlist, timeout=None):
sys.exit()

self.stubs.Set(websocketserver.WebSocketServer, 'socket',
self.stubs.Set(websockifyserver.WebSockifyServer, 'socket',
lambda *args, **kwargs: sock)
self.stubs.Set(websocketserver.WebSocketServer, 'daemonize',
self.stubs.Set(websockifyserver.WebSockifyServer, 'daemonize',
lambda *args, **kwargs: None)
self.stubs.Set(select, 'select', fake_select)
server.start_server()
Expand Down
12 changes: 6 additions & 6 deletions websockify/websocketproxy.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,15 +17,15 @@
try: from http.server import HTTPServer
except: from BaseHTTPServer import HTTPServer
import select
from websockify import websocketserver
from websockify import websockifyserver
from websockify import auth_plugins as auth
try:
from urllib.parse import parse_qs, urlparse
except:
from cgi import parse_qs
from urlparse import urlparse

class ProxyRequestHandler(websocketserver.WebSocketRequestHandler):
class ProxyRequestHandler(websockifyserver.WebSockifyRequestHandler):

buffer_size = 65536

Expand Down Expand Up @@ -88,7 +88,7 @@ def new_websocket_client(self):
msg += " (using SSL)"
self.log_message(msg)

tsock = websocketserver.WebSocketServer.socket(self.server.target_host,
tsock = websockifyserver.WebSockifyServer.socket(self.server.target_host,
self.server.target_port,
connect=True,
use_ssl=self.server.ssl_target,
Expand Down Expand Up @@ -221,7 +221,7 @@ def do_proxy(self, target):
cqueue.append(buf)
self.print_traffic("{")

class WebSocketProxy(websocketserver.WebSocketServer):
class WebSocketProxy(websockifyserver.WebSockifyServer):
"""
Proxy traffic to and from a WebSockets client to a normal TCP
socket server target.
Expand Down Expand Up @@ -274,7 +274,7 @@ def __init__(self, RequestHandlerClass=ProxyRequestHandler, *args, **kwargs):
"REBIND_OLD_PORT": str(kwargs['listen_port']),
"REBIND_NEW_PORT": str(self.target_port)})

websocketserver.WebSocketServer.__init__(self, RequestHandlerClass, *args, **kwargs)
websockifyserver.WebSockifyServer.__init__(self, RequestHandlerClass, *args, **kwargs)

def run_wrap_cmd(self):
self.msg("Starting '%s'", " ".join(self.wrap_cmd))
Expand Down Expand Up @@ -469,7 +469,7 @@ def websockify_init():
if len(args) > 2:
parser.error("Too many arguments")

if not websocketserver.ssl and opts.ssl_target:
if not websockifyserver.ssl and opts.ssl_target:
parser.error("SSL target requested and Python SSL module not loaded.");

if opts.ssl_only and not os.path.exists(opts.cert):
Expand Down
Loading

0 comments on commit e47591f

Please sign in to comment.