11import asyncio
2+ import hashlib
3+ import json
24import logging
5+ from pathlib import Path
6+ from typing import Optional
37
48import httpx
59import websockets
6- from fastapi import APIRouter , HTTPException , Request , WebSocket , WebSocketDisconnect
7- from fastapi .responses import Response
10+ from fastapi import APIRouter , BackgroundTasks , HTTPException , Request , WebSocket , WebSocketDisconnect
11+ from fastapi .responses import Response , StreamingResponse
12+ from starlette .background import BackgroundTask
813
914logger = logging .getLogger (__name__ )
1015router = APIRouter ()
11-
16+ cache_dir = Path ( "./cache" )
1217
1318# HTTP 转发
14- @router .api_route ("/http/{target_url:path}" , methods = ["GET" , "POST" , "PUT" , "DELETE" , "PATCH" , "OPTIONS" ])
19+ @router .api_route ("/proxy/ http/{target_url:path}" , methods = ["GET" , "POST" , "PUT" , "DELETE" , "PATCH" , "OPTIONS" ])
1520async def proxy_http (request : Request , target_url : str ):
1621 logger .info (f"HTTP target_url: { target_url } )
1722
1823 async with httpx .AsyncClient (timeout = httpx .Timeout (30.0 )) as client :
19- body = await request .body ()
24+ body = await request .body () if request .method in {"POST" , "PUT" , "PATCH" , "DELETE" } else None
25+ headers = {k : v for k , v in request .headers .items () if k .lower () not in "host" , "x-target-url" }}
26+ headers ['accept-encoding' ] = '' # disable gzip
2027 resp = await client .request (
2128 request .method ,
2229 target_url ,
2330 content = body ,
24- headers = { k : v for k , v in request . headers . items () if k . lower () != "host" and k . lower () != "x-target-url" }
31+ headers = headers ,
2532 )
2633 return Response (content = resp .content , status_code = resp .status_code , headers = dict (resp .headers ))
2734
35+
36+ @router .get ("/" )
37+ @router .get ("/android/{path:path}" )
38+ @router .get ("/ios/{path:path}" )
39+ @router .get ("/demo/{path:path}" )
40+ @router .get ("/harmony/{path:path}" )
41+ async def proxy_html (request : Request ):
42+ target_url = "https://uiauto.dev/"
43+ cache = HTTPCache (cache_dir , target_url , key = 'homepage' )
44+ response = await cache .proxy_request (request , update_cache = True )
45+ return response
46+ # update
47+
48+ @router .get ("/assets/{path:path}" )
49+ @router .get ('/favicon.ico' )
50+ async def proxy_assets (request : Request , path : str = "" ):
51+ target_url = f"https://uiauto.dev{ request .url .path }
52+ cache = HTTPCache (cache_dir , target_url )
53+ return await cache .proxy_request (request )
54+
55+
56+ class HTTPCache :
57+ def __init__ (self , cache_dir : Path , target_url : str , key : Optional [str ] = None ):
58+ self .cache_dir = cache_dir
59+ self .target_url = target_url
60+ self .key = key or hashlib .md5 (target_url .encode ()).hexdigest ()
61+ self .file_body = self .cache_dir / 'http' / (self .key + ".body" )
62+ self .file_headers = self .file_body .with_suffix (".headers" )
63+
64+ async def proxy_request (self , request : Request , update_cache : bool = False ):
65+ response = await self .get_cached_response (request )
66+ if not response :
67+ response = await self .proxy_and_save_response (request )
68+ return response
69+ if update_cache :
70+ # async update cache in background
71+ asyncio .create_task (self .update_cache (request ))
72+ return response
73+
74+ async def get_cached_response (self , request : Request ):
75+ if request .method == 'GET' and self .file_body .exists ():
76+ logger .info (f"Cache hit: { self .file_body } )
77+ headers = {}
78+ if self .file_headers .exists ():
79+ with self .file_headers .open ('rb' ) as f :
80+ headers = json .load (f )
81+ body_fd = self .file_body .open ("rb" )
82+ return StreamingResponse (
83+ content = body_fd ,
84+ status_code = 200 ,
85+ headers = headers ,
86+ background = BackgroundTask (body_fd .close )
87+ )
88+ return None
89+
90+ async def update_cache (self , request : Request ):
91+ try :
92+ await self .proxy_and_save_response (request )
93+ except Exception as e :
94+ logger .error ("Update cache failed" )
95+
96+ async def proxy_and_save_response (self , request : Request ) -> Response :
97+ logger .debug (f"Proxying request... { request .url .path } )
98+ response = await proxy_http (request , self .target_url )
99+ # save response to cache
100+ if request .method == "GET" and response .status_code == 200 and self .cache_dir .exists ():
101+ self .file_body .parent .mkdir (parents = True , exist_ok = True )
102+ with self .file_body .open ("wb" ) as f :
103+ f .write (response .body )
104+ with self .file_headers .open ("w" , encoding = "utf-8" ) as f :
105+ headers = response .headers
106+ headers ['cache-status' ] = 'HIT'
107+ json .dump (dict (headers ), f , indent = 2 , ensure_ascii = False )
108+ return response
109+
110+
28111# WebSocket 转发
29- @router .websocket ("/ws/{target_url:path}" )
112+ @router .websocket ("/proxy/ ws/{target_url:path}" )
30113async def proxy_ws (websocket : WebSocket , target_url : str ):
31114 await websocket .accept ()
32115 logger .info (f"WebSocket target_url: { target_url } )
@@ -54,4 +137,44 @@ async def from_server():
54137 pass
55138 except Exception as e :
56139 logger .error (f"WS Error: { e } )
57- await websocket .close ()
140+ await websocket .close ()
141+
142+ # ref: https://stackoverflow.com/questions/74555102/how-to-forward-fastapi-requests-to-another-server
143+ def make_reverse_proxy (base_url : str , strip_prefix : str = "" ):
144+ async def _reverse_proxy (request : Request ):
145+ client = httpx .AsyncClient (base_url = base_url )
146+ client .timeout = httpx .Timeout (30.0 , read = 300.0 )
147+ path = request .url .path
148+ if strip_prefix and path .startswith (strip_prefix ):
149+ path = path [len (strip_prefix ):]
150+ target_url = httpx .URL (
151+ path = path , query = request .url .query .encode ("utf-8" )
152+ )
153+ exclude_headers = [b"host" , b"connection" , b"accept-encoding" ]
154+ headers = [(k , v ) for k , v in request .headers .raw if k not in exclude_headers ]
155+ headers .append ((b'accept-encoding' , b'' ))
156+
157+ req = client .build_request (
158+ request .method , target_url , headers = headers , content = request .stream ()
159+ )
160+ r = await client .send (req , stream = True )#, follow_redirects=True)
161+
162+ response_headers = {
163+ k : v for k , v in r .headers .items ()
164+ if k .lower () not in "transfer-encoding" , "connection" , "content-length" }
165+ }
166+ async def gen_content ():
167+ async for chunk in r .aiter_bytes (chunk_size = 40960 ):
168+ yield chunk
169+
170+ async def aclose ():
171+ await client .aclose ()
172+
173+ return StreamingResponse (
174+ content = gen_content (),
175+ status_code = r .status_code ,
176+ headers = response_headers ,
177+ background = BackgroundTask (aclose ),
178+ )
179+
180+ return _reverse_proxy
0 commit comments