Skip to content

Commit

Permalink
[pre-commit.ci] auto fixes from pre-commit.com hooks
Browse files Browse the repository at this point in the history
for more information, see https://pre-commit.ci
  • Loading branch information
pre-commit-ci[bot] committed Aug 5, 2024
1 parent acf3d2c commit 6313168
Show file tree
Hide file tree
Showing 8 changed files with 48 additions and 35 deletions.
1 change: 1 addition & 0 deletions contrib/theia/jupyter_theia_proxy/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
See https://jupyter-server-proxy.readthedocs.io/en/latest/server-process.html
for more information.
"""

import os
import shutil

Expand Down
18 changes: 7 additions & 11 deletions jupyter_server_proxy/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -54,7 +54,11 @@ def _make_proxy_handler(sp: ServerProcess):
Create an appropriate handler with given parameters
"""
if sp.command:
cls = SuperviseAndRawSocketHandler if sp.raw_socket_proxy else SuperviseAndProxyHandler
cls = (
SuperviseAndRawSocketHandler
if sp.raw_socket_proxy
else SuperviseAndProxyHandler
)
args = dict(state={})
elif not (sp.port or isinstance(sp.unix_socket, str)):
warn(
Expand Down Expand Up @@ -122,13 +126,7 @@ def make_handlers(base_url, server_processes):
handler = _make_proxy_handler(sp)
if not handler:
continue
handlers.append(
(
ujoin(base_url, sp.name, r"(.*)"),
handler,
handler.kwargs
)
)
handlers.append((ujoin(base_url, sp.name, r"(.*)"), handler, handler.kwargs))
handlers.append((ujoin(base_url, sp.name), AddSlashHandler))
return handlers

Expand Down Expand Up @@ -159,9 +157,7 @@ def make_server_process(name, server_process_config, serverproxy_config):
"rewrite_response",
tuple(),
),
update_last_activity=server_process_config.get(
"update_last_activity", True
),
update_last_activity=server_process_config.get("update_last_activity", True),
raw_socket_proxy=server_process_config.get("raw_socket_proxy", False),
)

Expand Down
27 changes: 20 additions & 7 deletions jupyter_server_proxy/rawsocket.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,15 +9,17 @@

import asyncio

from .handlers import NamedLocalProxyHandler, SuperviseAndProxyHandler
from tornado import web

from .handlers import NamedLocalProxyHandler, SuperviseAndProxyHandler


class RawSocketProtocol(asyncio.Protocol):
"""
A protocol handler for the proxied stream connection.
Sends any received blocks directly as websocket messages.
"""

def __init__(self, handler):
self.handler = handler

Expand All @@ -30,14 +32,18 @@ def data_received(self, data):

def connection_lost(self, exc):
"Close the websocket connection."
self.handler.log.info(f"Raw websocket {self.handler.name} connection lost: {exc}")
self.handler.log.info(
f"Raw websocket {self.handler.name} connection lost: {exc}"
)
self.handler.close()


class RawSocketHandler(NamedLocalProxyHandler):
"""
HTTP handler that proxies websocket connections into a backend stream.
All other HTTP requests return 405.
"""

def _create_ws_connection(self, proto: asyncio.BaseProtocol):
"Create the appropriate backend asyncio connection"
loop = asyncio.get_running_loop()
Expand All @@ -46,17 +52,21 @@ def _create_ws_connection(self, proto: asyncio.BaseProtocol):
return loop.create_unix_connection(proto, self.unix_socket)
else:
self.log.info(f"RawSocket {self.name} connecting to port {self.port}")
return loop.create_connection(proto, 'localhost', self.port)
return loop.create_connection(proto, "localhost", self.port)

async def proxy(self, port, path):
raise web.HTTPError(405, "this raw_socket_proxy backend only supports websocket connections")
raise web.HTTPError(
405, "this raw_socket_proxy backend only supports websocket connections"
)

async def proxy_open(self, host, port, proxied_path=""):
"""
Open the backend connection. host and port are ignored (as they are in
the parent for unix sockets) since they are always passed known values.
"""
transp, proto = await self._create_ws_connection(lambda: RawSocketProtocol(self))
transp, proto = await self._create_ws_connection(
lambda: RawSocketProtocol(self)
)
self.ws_transp = transp
self.ws_proto = proto
self._record_activity()
Expand All @@ -66,8 +76,10 @@ def on_message(self, message):
"Send websocket messages as stream writes, encoding if necessary."
self._record_activity()
if isinstance(message, str):
message = message.encode('utf-8')
self.ws_transp.write(message) # buffered non-blocking. should block (needs new enough tornado)
message = message.encode("utf-8")
self.ws_transp.write(
message
) # buffered non-blocking. should block (needs new enough tornado)

def on_ping(self, message):
"No-op"
Expand All @@ -79,6 +91,7 @@ def on_close(self):
if hasattr(self, "ws_transp"):
self.ws_transp.close()


class SuperviseAndRawSocketHandler(SuperviseAndProxyHandler, RawSocketHandler):
async def _http_ready_func(self, p):
# not really HTTP here, just try an empty connection
Expand Down
1 change: 1 addition & 0 deletions tests/conftest.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
"""Reusable test fixtures for ``jupyter_server_proxy``."""

import os
import shutil
import socket
Expand Down
1 change: 1 addition & 0 deletions tests/resources/httpinfo.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
Simple webserver to respond with an echo of the sent request. It can listen to
either a tcp port or a unix socket.
"""

import argparse
import socket
import sys
Expand Down
10 changes: 5 additions & 5 deletions tests/resources/jupyter_server_config.py
Original file line number Diff line number Diff line change
Expand Up @@ -42,10 +42,10 @@ def cats_only(response, path):
response.code = 403
response.body = b"dogs not allowed"


def my_env():
return {
"MYVAR": "String with escaped {{var}}"
}
return {"MYVAR": "String with escaped {{var}}"}


c.ServerProxy.servers = {
"python-http": {
Expand Down Expand Up @@ -129,12 +129,12 @@ def my_env():
"python-proxyto54321-no-command": {"port": 54321},
"python-rawsocket-tcp": {
"command": [sys.executable, "./tests/resources/rawsocket.py", "{port}"],
"raw_socket_proxy": True
"raw_socket_proxy": True,
},
"python-rawsocket-unix": {
"command": [sys.executable, "./tests/resources/rawsocket.py", "{unix_socket}"],
"unix_socket": True,
"raw_socket_proxy": True
"raw_socket_proxy": True,
},
}

Expand Down
3 changes: 1 addition & 2 deletions tests/resources/rawsocket.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,5 @@
#!/usr/bin/env python

import os
import socket
import sys

Expand All @@ -11,7 +10,7 @@
try:
port = int(where)
family = socket.AF_INET
addr = ('localhost', port)
addr = ("localhost", port)
except ValueError:
family = socket.AF_UNIX
addr = where
Expand Down
22 changes: 12 additions & 10 deletions tests/test_proxies.py
Original file line number Diff line number Diff line change
Expand Up @@ -471,18 +471,20 @@ def test_callable_environment_formatting(
assert r.code == 200


@pytest.mark.parametrize("rawsocket_type", [
"tcp",
pytest.param(
"unix",
marks=pytest.mark.skipif(
sys.platform == "win32", reason="Unix socket not supported on Windows"
@pytest.mark.parametrize(
"rawsocket_type",
[
"tcp",
pytest.param(
"unix",
marks=pytest.mark.skipif(
sys.platform == "win32", reason="Unix socket not supported on Windows"
),
),
),
])
],
)
async def test_server_proxy_rawsocket(
rawsocket_type: str,
a_server_port_and_token: Tuple[int, str]
rawsocket_type: str, a_server_port_and_token: Tuple[int, str]
) -> None:
PORT, TOKEN = a_server_port_and_token
url = f"ws://{LOCALHOST}:{PORT}/python-rawsocket-{rawsocket_type}/?token={TOKEN}"
Expand Down

0 comments on commit 6313168

Please sign in to comment.