pus-handler.py (3158B)
1 """ 2 This script handles requests for the "main" and "signal" resources. 3 The "main" resource serves an HTML page that either initiates or acts as the 4 target for a prerendering speculation rule. The "signal" resource is used to 5 coordinate between the initiator and the prerendered page, ensuring the 6 prerendering is complete before proceeding. It uses query parameters to 7 determine its behavior and coordinate the process. 8 """ 9 10 import os 11 import copy 12 import time 13 from urllib.parse import parse_qs 14 15 def encode_query(query_dict): 16 encoded_pairs = [f"{key}={value[0]}" if value[0] !="" else key for key, value in query_dict.items()] 17 return '&'.join(encoded_pairs) 18 19 def calculate_signal_path(url_parts): 20 url_parts = copy.deepcopy(url_parts) 21 query_dict = parse_qs(url_parts.query, keep_blank_values=True) 22 query_dict['type'] = ['signal'] 23 return url_parts._replace(query=encode_query(query_dict)).geturl() 24 25 def calculate_prerendering_path(url_parts): 26 query_dict = parse_qs(url_parts.query) 27 query_dict['isprerendering'] = [True] 28 return url_parts._replace(query=encode_query(query_dict)).geturl() 29 30 def main_resource_handler(request, response): 31 is_prerendering = b"isprerendering" in request.GET 32 uid = request.GET.get(b"uid") 33 signal_path = calculate_signal_path(request.url_parts) 34 content = '' 35 if not is_prerendering: 36 # First request, which is the initiator 37 template_path = os.path.join(os.path.dirname(__file__), "pus-initiator-page-template.html") 38 prerendering_path = calculate_prerendering_path(request.url_parts) 39 with open(template_path, "r") as f: 40 content = f.read().replace("{{signal_url}}", signal_path).replace( 41 "{{prerendering_url}}", prerendering_path) 42 else: 43 template_path = os.path.join(os.path.dirname( 44 __file__), "pus-page-template.html") 45 with open(template_path, "r") as f: 46 content = f.read().replace("{{signal_path}}", signal_path) 47 response.headers.set(b"Content-Type", b"text/html") 48 response.status = 200 49 response.content = content.encode('utf-8') 50 51 def signal_handler(request, response): 52 is_prerendering = b"isprerendering" in request.GET 53 uid = request.GET.get(b"uid") 54 if is_prerendering: 55 with request.server.stash.lock: 56 request.server.stash.put(uid, "ok") 57 else: 58 # This will hang until the gate is released 59 while True: 60 with request.server.stash.lock: 61 if request.server.stash.take(uid) is None: 62 time.sleep(0.1) 63 else: 64 break 65 response.headers.set(b"Content-Type", b"text/javascript") 66 response.status = 200 67 response.content = "console.error('orz')".encode('utf-8') 68 69 def main(request, response): 70 resource_router = { 71 b"main":main_resource_handler, 72 b"signal": signal_handler 73 } 74 resource_type = request.GET.get(b"type") 75 resource_handler = resource_router.get(resource_type) 76 if resource_handler is not None: 77 return resource_handler(request, response) 78 response.status = 400 # Bad Request 79 response.content = b"Invalid resource type" 80 return