From 94bbc987c6af33091539fad2f2bff1aa6903cd37 Mon Sep 17 00:00:00 2001
From: Frank Ploss <frank@fqxp.de>
Date: Mon, 28 Dec 2015 17:39:08 +0100
Subject: [PATCH] Add auto-retry functionality

---
 extension.js                | 13 +++++--------
 tor_control_client.js       | 38 ++++++++++++++++++++++++++++++++++---
 ui/tor_button.js            |  2 +-
 ui/tor_disconnected_menu.js |  6 ------
 4 files changed, 41 insertions(+), 18 deletions(-)

diff --git a/extension.js b/extension.js
index 64af7b6..73bd2fa 100644
--- a/extension.js
+++ b/extension.js
@@ -38,19 +38,16 @@ function init(extensionMeta) {
 }
 
 function enable() {
-    torControlClient = new TorControlClient(TOR_CONTROL_HOST, TOR_CONTROL_PORT);
+    torControlClient = new TorControlClient(TOR_CONTROL_HOST, TOR_CONTROL_PORT, true);
     torButton = new TorButton(torControlClient);
     Main.panel.addToStatusArea(torButton.Name, torButton);
     torControlClient.openConnection();
 }
 
 function disable() {
-    if (torButton !== null) {
+    if (torControlClient !== null)
+        torControlClient.destroy();
+
+    if (torButton !== null)
         torButton.destroy();
-        torButton = null;
-    }
-    if (torControlClient !== null) {
-        torControlClient.closeConnection();
-        torControlClient = null;
-    }
 }
diff --git a/tor_control_client.js b/tor_control_client.js
index 4ef5abf..8749732 100644
--- a/tor_control_client.js
+++ b/tor_control_client.js
@@ -20,6 +20,7 @@ along with gnome-shell-extension-tor.  If not, see <http://www.gnu.org/licenses/
 'use strict';
 
 const Gio = imports.gi.Gio;
+const GLib = imports.gi.GLib;
 const Lang = imports.lang;
 const Signals = imports.signals;
 
@@ -44,10 +45,16 @@ const TorControlClient = new Lang.Class({
     Name: 'TorControlClient',
 
 
-    _init: function(host, port) {
+    _init: function(host, port, autoRetry) {
         this._host = host;
         this._port = port;
-        this._fail_reason = null;
+        this._autoRetry = autoRetry;
+        this._autoRetryTimerId = null;
+    },
+
+    destroy: function() {
+        this.stopAutoRetry();
+        this.closeConnection();
     },
 
     openConnection: function() {
@@ -57,11 +64,13 @@ const TorControlClient = new Lang.Class({
             this._ensureProtocolCompatibility();
             this._authenticate();
             this.emit('changed-connection-state', 'ready');
+            this.stopAutoRetry();
         } catch (e if e instanceof TorConnectionError) {
             this.closeConnection(e.message);
+            this.startAutoRetry();
         } catch (e if e instanceof TorProtocolError) {
-            //this.emit('protocol-error', 'Error while connecting to Tor control port', e.message);
             this.closeConnection(e.message);
+            this.startAutoRetry();
         }
     },
 
@@ -75,6 +84,28 @@ const TorControlClient = new Lang.Class({
         this.emit('changed-connection-state', 'closed', reason);
     },
 
+    startAutoRetry: function() {
+        if (!this._autoRetry)
+            return;
+
+        if (this._autoRetryTimerId !== null)
+            return;
+
+        this._autoRetryTimerId = GLib.timeout_add(GLib.PRIORITY_DEFAULT, 1000, Lang.bind(this, function() {
+            //log('Retrying to open connection ...');
+            this.openConnection();
+            return this._connection === null || !this._connection.is_connected();
+        }));
+    },
+
+    stopAutoRetry: function() {
+        if (this._autoRetryTimerId === null)
+            return;
+
+        GLib.source_remove(this._autoRetryTimerId);
+        this._autoRetryTimerId = null;
+    },
+
     switchIdentity: function() {
         var reply = this._runCommand('SIGNAL NEWNYM');
 
@@ -173,6 +204,7 @@ const TorControlClient = new Lang.Class({
             if (line === null) {
                 let reason = 'Lost connection to Tor server';
                 this.closeConnection(reason);
+                this.startAutoRetry();
                 return {replyLines: [reason]};
             }
 
diff --git a/ui/tor_button.js b/ui/tor_button.js
index 1ea1cbe..8c109c6 100644
--- a/ui/tor_button.js
+++ b/ui/tor_button.js
@@ -76,7 +76,7 @@ const TorButton = new Lang.Class({
     },
 
     _showConnectedMenu: function() {
-        if (this._menu instanceof TorConnectedIcon) {
+        if (this._menu instanceof TorPopupMenu) {
             return;
         }
 
diff --git a/ui/tor_disconnected_menu.js b/ui/tor_disconnected_menu.js
index 8b7ee7e..eb6311a 100644
--- a/ui/tor_disconnected_menu.js
+++ b/ui/tor_disconnected_menu.js
@@ -49,11 +49,5 @@ const TorDisconnectedMenu = new Lang.Class({
         this._msgLabel = new St.Label();
         errorMessageMenuItem.actor.add_actor(this._msgLabel);
         this.addMenuItem(errorMessageMenuItem);
-
-        this.addAction('Reconnect', Lang.bind(this, this._reconnect));
     },
-
-    _reconnect: function() {
-        this._torControlClient.openConnection();
-    }
 });