From d2ee6379dbded41121cf1b66592427510b1bd148 Mon Sep 17 00:00:00 2001
From: iou1name <iou1name@national.shitposting.agency>
Date: Fri, 15 Nov 2019 23:27:45 -0500
Subject: [PATCH] websocket connection more resilient

---
 aberrant.py          |  8 ++---
 static/aberrant.js   | 83 ++++++++++++++++++++++----------------------
 templates/index.html |  1 -
 3 files changed, 43 insertions(+), 49 deletions(-)

diff --git a/aberrant.py b/aberrant.py
index e9aed2e..637afd2 100644
--- a/aberrant.py
+++ b/aberrant.py
@@ -28,16 +28,11 @@ async def index(request):
 	tracker_stats = rtorrent.get_stats()
 	return render_template("index.html", request, locals())
 
-@routes.get(config.prefix + "/get_active_torrents", name='active-torrents')
-async def get_active_torrents(request):
-	"""Returns all active torrents formatted as JSON."""
-	data = [vars(t) for t in rtorrent.get_active()]
-	return web.json_response(data)
 
 @routes.get(config.prefix + '/ws', name='ws')
 async def websocket_handler(request):
 	"""The websocket endpoint."""
-	ws = web.WebSocketResponse()
+	ws = web.WebSocketResponse(heartbeat=30)
 	ws_ready = ws.can_prepare(request)
 	if not ws_ready.ok:
 		return web.Response(text="Cannot start websocket.")
@@ -60,6 +55,7 @@ async def websocket_handler(request):
 	await ws.close()
 	return ws
 
+
 app.router.add_routes(routes)
 
 if __name__ == "__main__":
diff --git a/static/aberrant.js b/static/aberrant.js
index 377e882..ab008f5 100644
--- a/static/aberrant.js
+++ b/static/aberrant.js
@@ -1,56 +1,55 @@
+var socket;
+
 function load() {
-	//let intervalID = window.setInterval(get_active_torrents_ajax, 20000);
+	socket = init_websocket();
 	let intervalID = window.setInterval(update, 5000);
 }
 
-function get_active_torrents_ajax() {
-	let httpRequest;
-	httpRequest = new XMLHttpRequest();
-	httpRequest.onreadystatechange = function() {
-		if (httpRequest.readyState !== XMLHttpRequest.DONE) { return; }
-		if (httpRequest.status !== 200) { return; }
-		torrents = JSON.parse(httpRequest.responseText);
-		let html_str = '';
-		for (let i = 0; i < torrents.length; i++) {
-			html_str += '<tr>'
-			html_str += '<td class="name">' + torrents[i].name + '</td>';
-			html_str += '<td class="total_size_str">' + torrents[i].total_size_str + '</td>';
-			html_str += '<td class="state">' + torrents[i].state + '</td>';
-			html_str += '<td class="downrate_str">' + torrents[i].downrate_str + '</td>';
-			html_str += '<td class="down_percent">' + torrents[i].down_percent + '%</td>';
-			html_str += '<td class="eta_str">' + torrents[i].eta_str + '</td>';
-			html_str += '<td class="uprate_str">' + torrents[i].uprate_str + '</td>';
-			html_str += '<td class="tracker">' + torrents[i].tracker + '</td>';
-			html_str += '<td class="rtorrent_id">' + torrents[i].rtorrent_id + '</td>';
-			html_str += '</tr>';
+function init_websocket() {
+	let socket = new WebSocket('wss://' + window.location.hostname + ws_uri);
+	socket._send = socket.send;
+	socket.send = function(event_title, data) {
+		data = JSON.stringify({event: event_title, data: data});
+		if (socket.readyState == 0) {
+			console.log("Socket is still opening!");
+			return;
 		}
-		document.getElementById('torrents').children[1].innerHTML = html_str;
-	};
-	httpRequest.open('GET', get_torrents_uri, true);
-	httpRequest.send();
+		socket._send(data);
+	}
+	socket.onmessage = onmessage;
+	socket.onclose = onclose;
+	socket.onerror = onerror;
+	socket.events = {};
+	socket.events['active_torrents'] = active_torrents_recv;
+	socket.events['tracker_stats'] = tracker_stats_recv;
+	return socket;
 }
 
-var socket = new WebSocket('wss://' + window.location.hostname + ws_uri);
-socket.oldSend = socket.send;
-socket.send = function(event_title, data) {
-		data = JSON.stringify({event: event_title, data: data});
-		socket.oldSend.apply(this, [data]);
+function onmessage (e) {
+	let data = JSON.parse(e.data);
+	let event = data.event;
+	data = data.data;
+	if (socket.events[event] === undefined) {
+		console.log("Unknown socket event: " + event);
+		return;
+	}
+	socket.events[event](data);
 }
-socket.events = {};
-socket.onmessage = function(e) {
-		let data = JSON.parse(e.data);
-		let event = data.event;
-		data = data.data;
-		if (socket.events[event] === undefined) { return; }
-		socket.events[event](data);
+
+function onclose(e) {
+	if (e.wasClean) { return; } // no need to reconnect
+	console.log(e);
+	console.log('Websocket lost connection to server. Re-trying...');
+	socket = init_websocket();
 }
-socket.onclose = function(e) {
-	console.log('WebSocket lost connection to server. Re-trying...');
-	// TODO: reconnect
+
+function onerror(e) {
+	console.log("Websocket error!")
+	console.log(e);
 }
 
 /* Websocket receive */
-socket.events['active_torrents'] = function(data) {
+function active_torrents_recv(data) {
 	let table = document.querySelector('#torrents tbody');
 	while (table.firstChild) {
 		table.removeChild(table.firstChild);
@@ -64,7 +63,7 @@ socket.events['active_torrents'] = function(data) {
 		table.appendChild(node);
 	});
 }
-socket.events['tracker_stats'] = function(data) {
+function tracker_stats(data) {
 	let table = document.querySelector('#tracker_stats tbody');
 	while (table.firstChild) {
 		table.removeChild(table.firstChild);
diff --git a/templates/index.html b/templates/index.html
index ccae2e0..42e0137 100644
--- a/templates/index.html
+++ b/templates/index.html
@@ -4,7 +4,6 @@
 	<title>Aberrant</title>
 	<link rel="stylesheet" type="text/css" href="/static/aberrant.css">
 	<script>
-		const get_torrents_uri = "{{ request.app.router['active-torrents'].url_for() }}";
 		const ws_uri = "{{ request.app.router['ws'].url_for() }}";
 	</script>
 	<script type="text/javascript" src="/static/aberrant.js"></script>