4
4
Some original inspiration from https://github.com/senko/tornado-proxy
5
5
"""
6
6
7
- import os , json , re
7
+ import os
8
+ import re
8
9
import socket
9
10
from asyncio import Lock
10
11
from copy import copy
@@ -307,7 +308,7 @@ def _build_proxy_request(self, host, port, proxied_path, body, **extra_opts):
307
308
decompress_response = False ,
308
309
headers = headers ,
309
310
** self .proxy_request_options (),
310
- ** extra_opts ,
311
+ ** extra_opts ,
311
312
)
312
313
return req
313
314
@@ -376,12 +377,12 @@ async def proxy(self, host, port, proxied_path):
376
377
else :
377
378
client = httpclient .AsyncHTTPClient (force_instance = True )
378
379
# check if the request is stream request
379
- accept_header = self .request .headers .get (' Accept' )
380
- if accept_header == ' text/event-stream' :
380
+ accept_header = self .request .headers .get (" Accept" )
381
+ if accept_header == " text/event-stream" :
381
382
return await self ._proxy_progressive (host , port , proxied_path , body , client )
382
383
else :
383
384
return await self ._proxy_buffered (host , port , proxied_path , body , client )
384
-
385
+
385
386
async def _proxy_progressive (self , host , port , proxied_path , body , client ):
386
387
# Proxy in progressive flush mode, whenever chunks are received. Potentially slower but get results quicker for voila
387
388
# Set up handlers so we can progressively flush result
@@ -390,15 +391,19 @@ async def _proxy_progressive(self, host, port, proxied_path, body, client):
390
391
391
392
def dump_headers (headers_raw ):
392
393
for line in headers_raw :
393
- r = re .match (' ^([a-zA-Z0-9\-_]+)\s*\:\ s*([^\r \n ]+)[\r \n ]*$' , line )
394
+ r = re .match (" ^([a-zA-Z0-9\\ -_]+)\\ s*\\ : \\ s*([^\r \n ]+)[\r \n ]*$" , line )
394
395
if r :
395
- k ,v = r .groups ([1 ,2 ])
396
- if k not in ('Content-Length' , 'Transfer-Encoding' ,
397
- 'Content-Encoding' , 'Connection' ):
396
+ k , v = r .groups ([1 , 2 ])
397
+ if k not in (
398
+ "Content-Length" ,
399
+ "Transfer-Encoding" ,
400
+ "Content-Encoding" ,
401
+ "Connection" ,
402
+ ):
398
403
# some header appear multiple times, eg 'Set-Cookie'
399
- self .set_header (k ,v )
404
+ self .set_header (k , v )
400
405
else :
401
- r = re .match (' ^HTTP[^\s]* ([0-9]+)' , line )
406
+ r = re .match (r" ^HTTP[^\s]* ([0-9]+)" , line )
402
407
if r :
403
408
status_code = r .group (1 )
404
409
self .set_status (int (status_code ))
@@ -414,20 +419,27 @@ def streaming_callback(chunk):
414
419
# record activity at start and end of requests
415
420
self ._record_activity ()
416
421
# Do this here, not in header_callback so we can be sure headers are out of the way first
417
- dump_headers (headers_raw ) # array will be empty if this was already called before
422
+ dump_headers (
423
+ headers_raw
424
+ ) # array will be empty if this was already called before
418
425
self .write (chunk )
419
426
self .flush ()
420
427
421
428
# Now make the request
422
429
423
- req = self ._build_proxy_request (host , port , proxied_path , body ,
424
- streaming_callback = streaming_callback ,
425
- header_callback = header_callback )
426
-
430
+ req = self ._build_proxy_request (
431
+ host ,
432
+ port ,
433
+ proxied_path ,
434
+ body ,
435
+ streaming_callback = streaming_callback ,
436
+ header_callback = header_callback ,
437
+ )
438
+
427
439
# no timeout for stream api
428
440
req .request_timeout = 7200
429
441
req .connect_timeout = 600
430
-
442
+
431
443
try :
432
444
response = await client .fetch (req , raise_error = False )
433
445
except httpclient .HTTPError as err :
@@ -444,15 +456,16 @@ def streaming_callback(chunk):
444
456
self .set_status (500 )
445
457
self .write (str (response .error ))
446
458
else :
447
- self .set_status (response .code , response .reason ) # Should already have been set
459
+ self .set_status (
460
+ response .code , response .reason
461
+ ) # Should already have been set
448
462
449
- dump_headers (headers_raw ) # Should already have been emptied
463
+ dump_headers (headers_raw ) # Should already have been emptied
450
464
451
- if response .body : # Likewise, should already be chunked out and flushed
465
+ if response .body : # Likewise, should already be chunked out and flushed
452
466
self .write (response .body )
453
467
454
468
async def _proxy_buffered (self , host , port , proxied_path , body , client ):
455
-
456
469
req = self ._build_proxy_request (host , port , proxied_path , body )
457
470
458
471
self .log .debug (f"Proxying request to { req .url } " )
@@ -535,7 +548,6 @@ def rewrite_pe(rewritable_response: RewritableResponse):
535
548
if rewritten_response .body :
536
549
self .write (rewritten_response .body )
537
550
538
-
539
551
async def proxy_open (self , host , port , proxied_path = "" ):
540
552
"""
541
553
Called when a client opens a websocket connection.
0 commit comments