request whitelist fixes + error handling unfuckery

ext/bg.js

@@ -15,7 +15,9 @@ var proxy_map = {};
 
 var web_response_whitelist = ["main_frame", "xmlhttprequest"];
 
-const log_debug = false;
+var webreq_completed = new Map();
+
+const log_debug = true;
 
 browser.browserAction.setBadgeBackgroundColor({
 	color: [0, 0, 0, 0]
@@ -473,6 +475,7 @@ function attach_ws_events(ws){
 	ws.addEventListener("close", async function(e){
 		
 		console.log("ws: offline");
+		webreq_completed = new Map();
 		await set_status(STATUS_OFFLINE);
 	});
 }
@@ -530,21 +533,52 @@ async function buff2b64(uint8Array) {
 	});
 }
 
+// store incoming data
 browser.webRequest.onBeforeRequest.addListener(
-	async function(details){
+	function(details){
+		
+		if(!web_response_whitelist.includes(details.type)){
+			
+			return;
+		}
 		
 		if(
 			web_response_whitelist.length === 0 ||
 			global_ws === null
-		){ return; }
+		){
+			
+			return;
+		}
 		
-		const filter = browser.webRequest.filterResponseData(
-			details.requestId
+		var resolve_func = null;
+		var reject_func = null;
+		
+		webreq_completed.set(
+			details.requestId,
+			new Promise(
+				function(resolve){
+					
+					resolve_func = resolve;
+				}
+			)
+			
 		);
 		
+		var filter = null;
+		
+		try{
+			
+			filter = browser.webRequest.filterResponseData(details.requestId);
+		}catch(err){
+			
+			resolve_func({success: false});
+			webreq_completed.delete(details.requestId);
+			return;
+		}	
+		
 		var chunks = [];
 		
-		filter.ondata = async function(event){
+		filter.ondata = function(event){
 			
 			chunks.push(event.data);
 			
@@ -552,7 +586,7 @@ browser.webRequest.onBeforeRequest.addListener(
 			filter.write(event.data);
 		};
 		
-		filter.onstop = async function(){
+		filter.onstop = function(){
 			
 			// we got the full response data
 			var len = 0;
@@ -570,33 +604,103 @@ browser.webRequest.onBeforeRequest.addListener(
 				offset += c.byteLength;
 			}
 			
-			var b64 = await buff2b64(merged);
-			
-			send_event(
-				global_ws,
-				{
-					action: "web_response",
-					data: {
-						id: details.tabId,
-						url: details.url,
-						status: details.statusCode,
-						origin: details.originUrl,
-						type: details.type,
-						method: details.method,
-						container: details.cookieStoreId,
-						url: details.url,
-						body: b64
-					}
-				}
-			);
-			
+			resolve_func({success: true, data: merged});
 			filter.close();
 		};
+		
+		filter.onerror = function(){
+			
+			resolve_func({success: false});
+			webreq_completed.delete(details.requestId);
+		}
 	},
-	{ urls: ["<all_urls>"], types: web_response_whitelist },
+	{ urls: ["<all_urls>"] },
 	["blocking"]
 );
 
+// detect page complete and send payloads
+browser.webRequest.onCompleted.addListener(
+	async function(details){
+		
+		var pending_req = null;
+		
+		pending_req = await webreq_completed.get(details.requestId);
+		
+		if(
+			// not in capture group
+			!pending_req ||
+			// some fuckass error occured
+			pending_req.success === false
+		){
+			
+			return;
+		}
+		
+		var b64 = await buff2b64(pending_req.data);
+		
+		send_event(
+			global_ws,
+			{
+				action: "web_response",
+				data: {
+					id: details.tabId,
+					url: details.url,
+					status: details.statusCode,
+					origin: details.originUrl,
+					type: details.type,
+					method: details.method,
+					container: details.cookieStoreId,
+					body: b64
+				}
+			}
+		);
+		
+		// clean up
+		webreq_completed.delete(details.requestId);
+	},
+	{ urls: ["<all_urls>"] }
+);
+
+// shit out an error when the request fails
+browser.webRequest.onErrorOccurred.addListener(
+	async function(page){
+		
+		if(connected === false){
+			
+			// disconnect ws handler should clean old payloads
+			return;
+		}
+		
+		var pending_req = webreq_completed.get(page.requestId);
+
+		if(
+			// we're not watching that req
+			!pending_req ||
+			// some bullshit happened
+			pending_req.success === false
+		){
+			
+			webreq_completed.delete(details.requestId);
+			return;
+		}
+		
+		send_event(
+			global_ws,
+			{
+				"action": "dom_load_fail",
+				"data": {
+					id: page.tabId,
+					url: page.url,
+					error: page.error
+				}
+			}
+		);
+		
+		webreq_completed.delete(page.requestId);
+	},
+	{ urls: ["<all_urls>"] }
+);
+
 browser.proxy.onRequest.addListener(function(request){
 	
 	const proxy_config = proxy_map[request.cookieStoreId];
@@ -664,35 +768,6 @@ browser.webRequest.onHeadersReceived.addListener(function(page){
 	{ urls: ["<all_urls>"] }
 );
 
-browser.webNavigation.onErrorOccurred.addListener(function(page){
-	
-	if(
-		connected === false ||
-		// must be an error in parent frame
-		page.frameId !== 0 ||
-		// ignore HSTS error. See:
-		// https://bugzilla.mozilla.org/show_bug.cgi?id=1801326
-		page.error == "Error code 2152398850" ||
-		// page-load abort
-		page.error == "NS_BINDING_ABORTED"
-	){
-		
-		return;
-	}
-	
-	send_event(
-		global_ws,
-		{
-			"action": "dom_load_fail",
-			"data": {
-				id: page.tabId,
-				url: page.url,
-				error: page.error
-			}
-		}
-	);
-});
-
 (async function(){
 	await ws_init();
 })();