network: Improve disconnect handling

Make sure disconnecting closes the SSH tunnel connection. Also add keep
alive ping to check if connection to server is still there and time out
if the keep alive ping does not respond.

2 files changed, 28 insertions, 11 deletions

diff --git a/src/main.py b/src/main.py
index cc5a966..a120031 100644
--- a/src/main.py
+++ b/src/main.py
@@ -350,6 +350,9 @@ class WeegtkApplication(Adw.Application):
             self.network.desync_weechat()
         elif message.msgid == '_upgrade_ended':
             self.network.sync_weechat()
+        elif message.msgid == "_pong":
+            # For now don't do anything with pong messages
+            pass
         else:
             print(f"Unknown message with id {message.msgid}")
 
@@ -372,9 +375,8 @@ class WeegtkApplication(Adw.Application):
             self.pages.insert(index, page)
 
     def buffer_input(self, source_object, full_name, text):
-        if self.network.is_connected():
-            message = f"input {full_name} {text}\n"
-            self.network.send_to_weechat(message)
+        message = f"input {full_name} {text}\n"
+        self.network.send_to_weechat(message)
 
 def main(version):
     """The application's entry point."""
diff --git a/src/relay/network.py b/src/relay/network.py
index b27edaa..18d2920 100644
--- a/src/relay/network.py
+++ b/src/relay/network.py
@@ -121,11 +121,6 @@ class Network(GObject.GObject):
         self._socketclient = Gio.SocketClient.new()
         self._socket = None
 
-        # TODO figure out how to deal with these signals
-        #self._socket.connected.self._socket_connected)
-        #self._socket.readyRead.connect(self._socket_read)
-        #self._socket.disconnected.connect(self._socket_disconnected)
-
     def _init_connection(self):
         self.status = STATUS_DISCONNECTED
         self._hostname = None
@@ -138,6 +133,7 @@ class Network(GObject.GObject):
         self._pwd_hash_algo = None
         self._pwd_hash_iter = 0
         self._server_nonce = None
+        self.ssh_tunnel = None
 
     def set_status(self, status):
         """Set current status."""
@@ -208,6 +204,12 @@ class Network(GObject.GObject):
         self.send_to_weechat(_PROTO_HANDSHAKE)
         self._handshake_timer = GLib.timeout_add(2000, self.handshake_timer_expired)
 
+    def keep_alive(self):
+        if self.is_connected():
+            self.send_to_weechat("ping keep_alive\n")
+            GLib.timeout_add(30*1000, self.keep_alive)
+        return False
+
     def _socket_read(self, source_object, res, *user_data):
         """Slot: data available on socket."""
         try:
@@ -230,6 +232,7 @@ class Network(GObject.GObject):
                 self._buffer = self._buffer[0:length]
             self.emit("message_from_weechat", GLib.Bytes(self._buffer))
             if not self.is_connected():
+                self.disconnect()
                 return
             self._buffer.clear()
             if remainder:
@@ -238,10 +241,12 @@ class Network(GObject.GObject):
         self.input.read_bytes_async(
             4096, 0, self.cancel_network_reads, self._socket_read)
 
-    def _socket_disconnected(self):
+    def disconnect(self):
         """Slot: socket disconnected."""
-        if self._handshake_timer:
-            self._handshake_timer.stop()
+        if self._handshake_timer is not None:
+            GLib.source_remove(self._handshake_timer)
+        if self.ssh_tunnel is not None:
+            self.ssh_tunnel.stop(force=True)
         self._init_connection()
         self.set_status(STATUS_DISCONNECTED)
 
@@ -291,6 +296,9 @@ class Network(GObject.GObject):
             Gio.NetworkAddress.new(self._hostname, self._port),
             None,
             self._connected_func, None)
+
+        # TODO probably make the timeout configurable
+        self._socketclient.set_timeout(60)
         self.set_status(STATUS_CONNECTING)
 
     def _connected_func(self, source_object, res, *user_data):
@@ -324,19 +332,25 @@ class Network(GObject.GObject):
             self.set_status(STATUS_DISCONNECTED)
 
     def handle_network_error(self, err):
+        print("net err")
         if err.matches(Gio.io_error_quark(), Gio.IOErrorEnum.CANCELLED):
             print("Connection has been canceled by user.")
+            self.disconnect()
             return
         elif err.matches(Gio.tls_error_quark(), Gio.TlsError.EOF):
             print("Server has closed the connection.")
+            self.disconnect()
             return
         elif err.matches(Gio.io_error_quark(), Gio.IOErrorEnum.BROKEN_PIPE):
             print("Broken pipe, connection lost.")
+            self.disconnect()
             return
         elif err.matches(Gio.io_error_quark(), Gio.IOErrorEnum.TIMED_OUT):
             print("Connection timed out.")
+            self.disconnect()
             return
         else:
+            print("Unhandled net error {}".format(err))
             raise
 
     def send_to_weechat(self, message):
@@ -358,6 +372,7 @@ class Network(GObject.GObject):
                 self.send_to_weechat(cmd)
                 self.sync_weechat()
                 self.set_status(STATUS_CONNECTED)
+                GLib.timeout_add(30*1000, self.keep_alive)
                 return
         # failed to initialize: disconnect
         self.disconnect_weechat()