99
1010import asyncio
1111
12- from .handlers import NamedLocalProxyHandler , SuperviseAndProxyHandler
1312from tornado import web
1413
14+ from .handlers import NamedLocalProxyHandler , SuperviseAndProxyHandler
15+
1516
1617class RawSocketProtocol (asyncio .Protocol ):
1718 """
1819 A protocol handler for the proxied stream connection.
1920 Sends any received blocks directly as websocket messages.
2021 """
22+
2123 def __init__ (self , handler ):
2224 self .handler = handler
2325
@@ -30,14 +32,18 @@ def data_received(self, data):
3032
3133 def connection_lost (self , exc ):
3234 "Close the websocket connection."
33- self .handler .log .info (f"Raw websocket { self .handler .name } connection lost: { exc } " )
35+ self .handler .log .info (
36+ f"Raw websocket { self .handler .name } connection lost: { exc } "
37+ )
3438 self .handler .close ()
3539
40+
3641class RawSocketHandler (NamedLocalProxyHandler ):
3742 """
3843 HTTP handler that proxies websocket connections into a backend stream.
3944 All other HTTP requests return 405.
4045 """
46+
4147 def _create_ws_connection (self , proto : asyncio .BaseProtocol ):
4248 "Create the appropriate backend asyncio connection"
4349 loop = asyncio .get_running_loop ()
@@ -46,17 +52,21 @@ def _create_ws_connection(self, proto: asyncio.BaseProtocol):
4652 return loop .create_unix_connection (proto , self .unix_socket )
4753 else :
4854 self .log .info (f"RawSocket { self .name } connecting to port { self .port } " )
49- return loop .create_connection (proto , ' localhost' , self .port )
55+ return loop .create_connection (proto , " localhost" , self .port )
5056
5157 async def proxy (self , port , path ):
52- raise web .HTTPError (405 , "this raw_socket_proxy backend only supports websocket connections" )
58+ raise web .HTTPError (
59+ 405 , "this raw_socket_proxy backend only supports websocket connections"
60+ )
5361
5462 async def proxy_open (self , host , port , proxied_path = "" ):
5563 """
5664 Open the backend connection. host and port are ignored (as they are in
5765 the parent for unix sockets) since they are always passed known values.
5866 """
59- transp , proto = await self ._create_ws_connection (lambda : RawSocketProtocol (self ))
67+ transp , proto = await self ._create_ws_connection (
68+ lambda : RawSocketProtocol (self )
69+ )
6070 self .ws_transp = transp
6171 self .ws_proto = proto
6272 self ._record_activity ()
@@ -66,8 +76,10 @@ def on_message(self, message):
6676 "Send websocket messages as stream writes, encoding if necessary."
6777 self ._record_activity ()
6878 if isinstance (message , str ):
69- message = message .encode ('utf-8' )
70- self .ws_transp .write (message ) # buffered non-blocking. should block (needs new enough tornado)
79+ message = message .encode ("utf-8" )
80+ self .ws_transp .write (
81+ message
82+ ) # buffered non-blocking. should block (needs new enough tornado)
7183
7284 def on_ping (self , message ):
7385 "No-op"
@@ -79,6 +91,7 @@ def on_close(self):
7991 if hasattr (self , "ws_transp" ):
8092 self .ws_transp .close ()
8193
94+
8295class SuperviseAndRawSocketHandler (SuperviseAndProxyHandler , RawSocketHandler ):
8396 async def _http_ready_func (self , p ):
8497 # not really HTTP here, just try an empty connection
0 commit comments