set_beacon.py (3235B)
1 """An HTTP request handler for WPT that handles /set_beacon.py requests.""" 2 3 _BEACON_ID_KEY = b"uuid" 4 _BEACON_DATA_PATH = "beacon_data" 5 _BEACON_FORM_PAYLOAD_KEY = b"payload" 6 _BEACON_BODY_PAYLOAD_KEY = "payload=" 7 _BEACON_EXPECT_ORIGIN_KEY = b"expectOrigin" 8 _BEACON_EXPECT_PREFLIGHT_KEY = b"expectPreflight" 9 _BEACON_EXPECT_CREDS_KEY = b"expectCredentials" 10 11 12 def main(request, response): 13 """Stores the given beacon's data keyed by uuid in the server. 14 15 For GET request, this handler assumes no data. 16 For POST request, this handler extracts data from request body: 17 - Content-Type=multipart/form-data: data keyed by 'payload'. 18 - the entire request body. 19 20 Multiple data can be added for the same uuid. 21 22 The data is stored as UTF-8 format. 23 """ 24 if _BEACON_ID_KEY not in request.GET: 25 response.status = 400 26 return "Must provide a UUID to store beacon data" 27 uuid = request.GET.first(_BEACON_ID_KEY) 28 29 expected_origin = request.GET.get(_BEACON_EXPECT_ORIGIN_KEY) 30 if b"origin" in request.headers: 31 origin = request.headers.get(b"origin") 32 if expected_origin: 33 assert origin == expected_origin, f"expected {expected_origin}, got {origin}" 34 response.headers.set(b"Access-Control-Allow-Origin", origin) 35 else: 36 assert expected_origin is None, f"expected None, got {expected_origin}" 37 38 # Handles preflight request first. 39 if request.method == u"OPTIONS": 40 assert request.GET.get( 41 _BEACON_EXPECT_PREFLIGHT_KEY) == b"true", "Preflight not expected." 42 43 # preflight must not have cookies. 44 assert b"Cookie" not in request.headers 45 46 requested_headers = request.headers.get( 47 b"Access-Control-Request-Headers") 48 assert b"content-type" in requested_headers, f"expected content-type, got {requested_headers}" 49 response.headers.set(b"Access-Control-Allow-Headers", b"content-type") 50 51 requested_method = request.headers.get(b"Access-Control-Request-Method") 52 assert requested_method == b"POST", f"expected POST, got {requested_method}" 53 response.headers.set(b"Access-Control-Allow-Methods", b"POST") 54 55 return response 56 57 expect_creds = request.GET.get(_BEACON_EXPECT_CREDS_KEY) == b"true" 58 if expect_creds: 59 assert b"Cookie" in request.headers 60 else: 61 assert b"Cookie" not in request.headers 62 63 data = None 64 if request.method == u"POST": 65 if b"multipart/form-data" in request.headers.get(b"Content-Type", b""): 66 if _BEACON_FORM_PAYLOAD_KEY in request.POST: 67 data = request.POST.first(_BEACON_FORM_PAYLOAD_KEY).decode( 68 'utf-8') 69 elif request.body: 70 data = request.body.decode('utf-8') 71 if data.startswith(_BEACON_BODY_PAYLOAD_KEY): 72 data = data.split(_BEACON_BODY_PAYLOAD_KEY)[1] 73 74 with request.server.stash.lock: 75 saved_data = request.server.stash.take(key=uuid, path=_BEACON_DATA_PATH) 76 if not saved_data: 77 saved_data = [data] 78 else: 79 saved_data.append(data) 80 request.server.stash.put( 81 key=uuid, value=saved_data, path=_BEACON_DATA_PATH) 82 83 response.status = 200