network-partition-key.py (6455B)
1 import mimetypes 2 import os 3 4 from wptserve.utils import isomorphic_decode, isomorphic_encode 5 6 # Test server that tracks the last partition_id was used with each connection for each uuid, and 7 # lets consumers query if multiple different partition_ids have been been used for any socket. 8 # 9 # Server assumes that ports aren't reused, so a client address and a server port uniquely identify 10 # a connection. If that constraint is ever violated, the test will be flaky. No sockets being 11 # closed for the duration of the test is sufficient to ensure that, though even if sockets are 12 # closed, the OS should generally prefer to use new ports for new connections, if any are 13 # available. 14 def main(request, response): 15 response.headers.set(b"Cache-Control", b"no-store") 16 dispatch = request.GET.first(b"dispatch", None) 17 uuid = request.GET.first(b"uuid", None) 18 partition_id = request.GET.first(b"partition_id", None) 19 20 if not uuid or not dispatch or not partition_id: 21 return simple_response(request, response, 404, b"Not found", b"Invalid query parameters") 22 23 # Unless nocheck_partition is true, check partition_id against server_state, and update server_state. 24 stash = request.server.stash 25 test_failed = False 26 request_count = 0; 27 connection_count = 0; 28 if request.GET.first(b"nocheck_partition", None) != b"True": 29 # Need to grab the lock to access the Stash, since requests are made in parallel. 30 with stash.lock: 31 # Don't use server hostname here, since H2 allows multiple hosts to reuse a connection. 32 # Server IP is not currently available, unfortunately. 33 address_key = isomorphic_encode(str(request.client_address) + u"|" + str(request.url_parts.port)) 34 server_state = stash.take(uuid) or {b"test_failed": False, 35 b"request_count": 0, b"connection_count": 0} 36 request_count = server_state[b"request_count"] 37 request_count += 1 38 server_state[b"request_count"] = request_count 39 if address_key in server_state: 40 if server_state[address_key] != partition_id: 41 server_state[b"test_failed"] = True 42 else: 43 connection_count = server_state[b"connection_count"] 44 connection_count += 1 45 server_state[b"connection_count"] = connection_count 46 server_state[address_key] = partition_id 47 test_failed = server_state[b"test_failed"] 48 stash.put(uuid, server_state) 49 50 origin = request.headers.get(b"Origin") 51 if origin: 52 response.headers.set(b"Access-Control-Allow-Origin", origin) 53 response.headers.set(b"Access-Control-Allow-Credentials", b"true") 54 55 if request.method == u"OPTIONS": 56 return handle_preflight(request, response) 57 58 if dispatch == b"fetch_file": 59 return handle_fetch_file(request, response, partition_id, uuid) 60 61 if dispatch == b"check_partition": 62 status = request.GET.first(b"status", 200) 63 if test_failed: 64 return simple_response(request, response, status, b"OK", b"Multiple partition IDs used on a socket") 65 body = b"ok" 66 if request.GET.first(b"addcounter", False): 67 body += (". Request was sent " + str(request_count) + " times. " + 68 str(connection_count) + " connections were created.").encode('utf-8') 69 return simple_response(request, response, status, b"OK", body) 70 71 if dispatch == b"clean_up": 72 stash.take(uuid) 73 if test_failed: 74 return simple_response(request, response, 200, b"OK", b"Test failed, but cleanup completed.") 75 return simple_response(request, response, 200, b"OK", b"cleanup complete") 76 77 return simple_response(request, response, 404, b"Not Found", b"Unrecognized dispatch parameter: " + dispatch) 78 79 def handle_preflight(request, response): 80 response.status = (200, b"OK") 81 response.headers.set(b"Access-Control-Allow-Methods", b"GET") 82 response.headers.set(b"Access-Control-Allow-Headers", b"header-to-force-cors") 83 response.headers.set(b"Access-Control-Max-Age", b"86400") 84 return b"Preflight request" 85 86 def simple_response(request, response, status_code, status_message, body, content_type=b"text/plain"): 87 response.status = (status_code, status_message) 88 response.headers.set(b"Content-Type", content_type) 89 return body 90 91 def handle_fetch_file(request, response, partition_id, uuid): 92 subresource_origin = request.GET.first(b"subresource_origin", None) 93 rel_path = request.GET.first(b"path", None) 94 95 # This needs to be passed on to subresources so they all have access to it. 96 include_credentials = request.GET.first(b"include_credentials", None) 97 if not subresource_origin or not rel_path or not include_credentials: 98 return simple_response(request, response, 404, b"Not found", b"Invalid query parameters") 99 100 cur_path = os.path.realpath(isomorphic_decode(__file__)) 101 base_path = os.path.abspath(os.path.join(os.path.dirname(cur_path), os.pardir, os.pardir, os.pardir)) 102 path = os.path.abspath(os.path.join(base_path, isomorphic_decode(rel_path))) 103 104 # Basic security check. 105 if not path.startswith(base_path): 106 return simple_response(request, response, 404, b"Not found", b"Invalid path") 107 108 sandbox = request.GET.first(b"sandbox", None) 109 if sandbox == b"true": 110 response.headers.set(b"Content-Security-Policy", b"sandbox allow-scripts") 111 112 file = open(path, mode="rb") 113 body = file.read() 114 file.close() 115 116 subresource_path = b"/" + isomorphic_encode(os.path.relpath(isomorphic_decode(__file__), base_path)).replace(b'\\', b'/') 117 subresource_params = b"?partition_id=" + partition_id + b"&uuid=" + uuid + b"&subresource_origin=" + subresource_origin + b"&include_credentials=" + include_credentials 118 body = body.replace(b"SUBRESOURCE_PREFIX:", subresource_origin + subresource_path + subresource_params) 119 120 other_origin = request.GET.first(b"other_origin", None) 121 if other_origin: 122 body = body.replace(b"OTHER_PREFIX:", other_origin + subresource_path + subresource_params) 123 124 mimetypes.init() 125 mimetype_pair = mimetypes.guess_type(path) 126 mimetype = mimetype_pair[0] 127 128 if mimetype == None or mimetype_pair[1] != None: 129 return simple_response(request, response, 500, b"Server Error", b"Unknown MIME type") 130 return simple_response(request, response, 200, b"OK", body, mimetype)