[refactor] re-add retries for connecting to management

2 files changed, 37 insertions, 17 deletions

diff --git a/src/leap/bitmask/vpn/process.py b/src/leap/bitmask/vpn/process.py
index 48d3f1a..4ef9fd9 100644
--- a/src/leap/bitmask/vpn/process.py
+++ b/src/leap/bitmask/vpn/process.py
@@ -89,6 +89,7 @@ class _VPNProcess(protocol.ProcessProtocol):
         self._restartfun = restartfun
 
         self.restarting = True
+        self.failed = False
         self.proto = None
         self._remotes = remotes
 
@@ -106,18 +107,28 @@ class _VPNProcess(protocol.ProcessProtocol):
         except Exception as exc:
             print('[!] Error: %s' % exc)
 
-    def _connect_to_management(self):
-        # TODO -- add retries, twisted style, to this.
-        # this sometimes raises 'file not found' error
+    def _connect_to_management(self, retries=30):
+        if retries == 0:
+            self.log.error('Timeout whilte connecting to management')
+            self.failed = True
+            return
+
+        def retry(retries):
+            ctr = retries - 1
+            self.log.warn(
+                'Error connecting to management, retrying. Retries left:  %s' % ctr)
+            reactor.callLater(
+                0.1, self._connect_to_management, ctr)
+
         self._d = connectProtocol(
             self._management_endpoint,
             ManagementProtocol(verbose=True))
-        self._d.addCallback(self._got_management_protocol)
-        self._d.addErrback(self.log.error)
+        self._d.addCallbacks(
+            self._got_management_protocol,
+            lambda f: retry(retries))
 
     def connectionMade(self):
-        # TODO cut this wait time when retries are done
-        reactor.callLater(0.5, self._connect_to_management)
+        reactor.callLater(0.1, self._connect_to_management)
 
     def processExited(self, failure):
         err = failure.trap(
@@ -143,7 +154,9 @@ class _VPNProcess(protocol.ProcessProtocol):
         Called when the child process exits and all file descriptors associated
         with it have been closed.
         """
+        self.proto = None
         exit_code = reason.value.exitCode
+
         if isinstance(exit_code, int):
             self.log.debug('processEnded, status %d' % (exit_code,))
             if self.restarting:
@@ -174,6 +187,7 @@ class _VPNProcess(protocol.ProcessProtocol):
         if not self.proto:
             status = {'status': 'off', 'error': None}
             return status
+
         status = {'status': self.proto.state.simple.lower(),
                   'error': None}
         if self.proto.traffic:
@@ -230,7 +244,6 @@ class _VPNProcess(protocol.ProcessProtocol):
             self.log.debug('Process Exited Already')
 
     def terminate_or_kill(self):
-        # XXX this returns a deferred
         return self._launcher.terminate_or_kill(
             self.terminate, self.kill, self)
 
diff --git a/src/leap/bitmask/vpn/tunnel.py b/src/leap/bitmask/vpn/tunnel.py
index d05a50e..d9a8be3 100644
--- a/src/leap/bitmask/vpn/tunnel.py
+++ b/src/leap/bitmask/vpn/tunnel.py
@@ -76,11 +76,15 @@ class ConfiguredTunnel(object):
         self._port = port
         self._vpnproc = None
 
+    @defer.inlineCallbacks
     def start(self):
-        return self._start_vpn()
+        started = yield self._start_vpn()
+        defer.returnValue(started)
 
+    @defer.inlineCallbacks
     def stop(self):
-        return self._stop_vpn(restart=False)
+        stopped = yield self._stop_vpn(restart=False)
+        defer.returnValue(stopped)
 
     #  status
 
@@ -96,6 +100,7 @@ class ConfiguredTunnel(object):
 
     # VPN Control
 
+    @defer.inlineCallbacks
     def _start_vpn(self):
         self.log.debug('VPN: start')
         args = [self._vpnconfig, self._providerconfig, self._host,
@@ -105,15 +110,15 @@ class ConfiguredTunnel(object):
         vpnproc = VPNProcess(*args, **kwargs)
         self._vpnproc = vpnproc
 
-        self._start_pre_up(vpnproc)
+        self.__start_pre_up(vpnproc)
         cmd = self.__start_get_cmd(vpnproc)
-        # XXX this should be a deferred
-        running = self.__start_spawn_proc(vpnproc, cmd)
+
+        running = yield self.__start_spawn_proc(vpnproc, cmd)
         if running:
             vpnproc.pid = running.pid
-            return True
+            defer.returnValue(True)
         else:
-            return False
+            defer.returnValue(False)
 
     def __start_pre_up(self, proc):
         try:
@@ -147,6 +152,7 @@ class ConfiguredTunnel(object):
         reactor.callLater(
             self.RESTART_WAIT, self.start)
 
+    @defer.inlineCallbacks
     def _stop_vpn(self, restart=False):
         """
         Stops the openvpn subprocess.
@@ -162,11 +168,12 @@ class ConfiguredTunnel(object):
 
         if self._vpnproc is None:
             self.log.debug('Tried to stop VPN but no process found')
-            return
+            defer.returnValue(False)
 
         self._vpnproc.restarting = restart
         self.__stop_pre_down(self._vpnproc)
-        self._vpnproc.terminate_or_kill()
+        stopped = yield self._vpnproc.terminate_or_kill()
+        defer.returnValue(stopped)
 
     def __stop_pre_down(self, proc):
         try: