From 288903e0e914ebfb3acd8991953442cd5a0d1405 Mon Sep 17 00:00:00 2001
From: Rokas Puzonas <rokas.puz@gmail.com>
Date: Sun, 14 May 2023 14:13:07 +0300
Subject: [PATCH] add sending keepalive packets

---
 Cargo.toml |  2 +-
 src/app.rs | 49 ++++++++++++++++++++++++++++++++++++-------------
 2 files changed, 37 insertions(+), 14 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index ae90350..f547e3f 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -21,7 +21,7 @@ smol = "1.3.0"
 ssh2 = "0.9.4"
 syntect = "5.0.0"
 thiserror = "1.0.40"
-tokio = { version = "1.26.0", features = ["macros", "rt-multi-thread"] }
+tokio = { version = "1.26.0", features = ["macros", "rt-multi-thread", "time"] }
 toml = "0.7.3"
 version = "3.0.0"
 
diff --git a/src/app.rs b/src/app.rs
index f5a4712..0d7eb0f 100644
--- a/src/app.rs
+++ b/src/app.rs
@@ -1,10 +1,10 @@
 use std::{
     net::{SocketAddr, SocketAddrV4},
     sync::{mpsc::{Receiver, Sender}, Arc},
-    vec, rc::Rc, time::SystemTime, path::{PathBuf, Path}, fs,
+    vec, rc::Rc, time::{SystemTime, Duration}, path::{PathBuf, Path}, fs,
 };
 
-use anyhow::{Result, anyhow};
+use anyhow::Result;
 use async_ssh2_lite::{AsyncIoTcpStream, AsyncSession};
 use directories_next::ProjectDirs;
 use eframe::CreationContext;
@@ -64,6 +64,7 @@ pub struct App {
     session: Option<AsyncSession<AsyncIoTcpStream>>,
     ubus_call_handle: Option<JoinHandle<()>>,
     last_ubus_call_at: SystemTime,
+    keepalive_handle: Option<JoinHandle<()>>,
 
     selected_object: Option<Rc<ubus::Object>>,
     selected_method: Option<String>,
@@ -104,6 +105,7 @@ impl Default for App {
             },
             session: None,
             ubus_call_handle: None,
+            keepalive_handle: None,
             last_ubus_call_at: SystemTime::UNIX_EPOCH,
 
             object_filter: "".into(),
@@ -131,7 +133,7 @@ async fn connect<A>(
     password: String,
 ) -> Result<AsyncSession<AsyncIoTcpStream>>
 where
-    A: Into<SocketAddr>,
+    A: Into<SocketAddr>
 {
     let mut session = AsyncSession::<AsyncIoTcpStream>::connect(socket_addr, None).await?;
     session.handshake().await?;
@@ -176,15 +178,12 @@ impl App {
                 return;
             }
 
-            let address = self.settings.address.parse();
-            if address.is_err() {
-                return;
-            }
-            let address = address.unwrap();
-            let port = self.settings.port;
+            if let Ok(addr) = self.settings.address.parse() {
+                let port = self.settings.port;
 
-            let socket_addr = SocketAddrV4::new(address, port);
-            self.start_connect(socket_addr, username.clone(), password.clone());
+                let socket_addr = SocketAddrV4::new(addr, port);
+                self.start_connect(socket_addr, username.clone(), password.clone());
+            }
         }
 
         self.copy_texture = Some(cc.egui_ctx.load_texture(
@@ -235,17 +234,19 @@ impl App {
         return false;
     }
 
-    fn handle_events(&mut self, _ctx: &egui::Context) {
+    fn handle_events(&mut self, ctx: &egui::Context) {
         use AsyncEvent::*;
 
         if let Ok(event) = self.rx.try_recv() {
+            ctx.request_repaint();
             match event {
                 Connect(result) => {
                     self.is_connecting = false;
                     match result {
                         Ok(session) => {
                             self.session = Some(session);
-                            self.start_list_objects()
+                            self.start_list_objects();
+                            self.start_keepalive();
                         }
                         Err(err) => todo!("{}", err),
                     }
@@ -253,6 +254,7 @@ impl App {
 
                 Disconnect(result) => {
                     self.is_disconnecting = false;
+                    self.stop_keepalive();
 
                     if let Err(err) = result {
                         todo!("{}", err)
@@ -272,6 +274,27 @@ impl App {
         }
     }
 
+    fn start_keepalive(&mut self) {
+        if let Some(session) = &self.session {
+            if !session.authenticated() { return; }
+
+            let session = session.clone();
+            let handle = tokio::spawn(async move {
+                let next_keepalive = session.keepalive_send().await.unwrap();
+                tokio::time::sleep(Duration::from_secs(next_keepalive as u64)).await
+            });
+
+            self.keepalive_handle = Some(handle);
+        }
+    }
+
+    fn stop_keepalive(&mut self) {
+        if let Some(handle) = &self.keepalive_handle {
+            handle.abort();
+            self.keepalive_handle = None;
+        }
+    }
+
     fn start_connect<A>(&mut self, socket_addr: A, username: String, password: String)
     where
         A: Into<SocketAddr>,