From 55bbf89b514da50256ebbb01598b90dfa7b0b407 Mon Sep 17 00:00:00 2001
From: Richardas Kuchinskas <r.kuchinskas@vk.team>
Date: Thu, 3 Oct 2024 03:13:18 +0300
Subject: [PATCH] fixed exception whe yubikey have no key and trying to connect

---
 .../vk/admstorm/ssh/SshConnectionService.kt   | 56 ++++++++++++-------
 1 file changed, 37 insertions(+), 19 deletions(-)

diff --git a/src/main/kotlin/com/vk/admstorm/ssh/SshConnectionService.kt b/src/main/kotlin/com/vk/admstorm/ssh/SshConnectionService.kt
index 4e8946e8..356729d5 100644
--- a/src/main/kotlin/com/vk/admstorm/ssh/SshConnectionService.kt
+++ b/src/main/kotlin/com/vk/admstorm/ssh/SshConnectionService.kt
@@ -169,25 +169,7 @@ class SshConnectionService(private var myProject: Project) : Disposable {
                                 "Plugin can try to automatically reset the Yubikey or you can do it yourself with " +
                                 code("ssh-agent")
 
-                        AdmErrorNotification(message, true)
-                            .withTitle("Failed to connect to server")
-                            .withActions(AdmNotification.Action("Reset Yubikey and Connect...") { _, notification ->
-                                notification.expire()
-
-                                val success = YubikeyHandler().autoReset(project) {
-                                    connectWithConnector(connector, onSuccessful)
-                                }
-
-                                if (!success) {
-                                    return@Action
-                                }
-                                connectWithConnector(connector, onSuccessful)
-                            })
-                            .withActions(AdmNotification.Action("Connect...") { _, notification ->
-                                notification.expire()
-                                connectWithConnector(connector, onSuccessful)
-                            })
-                            .show()
+                        resetYubikeyNotification(message)
 
                         LOG.warn("Failed to connect", ex)
                     } catch (ex: TimeoutException) {
@@ -235,6 +217,11 @@ class SshConnectionService(private var myProject: Project) : Disposable {
                     } catch (ex: Exception) {
                         scheduler.shutdownNow()
 
+                        if (ex.message?.contains("Failed to open SSH channel in a just created TCP session") == true) {
+                            handleOpenTcpSessionError()
+                            return
+                        }
+
                         val exceptionName = ex.javaClass.name
                         LOG.error("Unhandled exception", ex)
                         AdmErrorNotification("Unhandled exception $exceptionName").show()
@@ -242,6 +229,37 @@ class SshConnectionService(private var myProject: Project) : Disposable {
                     }
                 }
 
+                fun resetYubikeyNotification(message: String) {
+                    AdmErrorNotification(message, true)
+                        .withTitle("Failed to connect to server")
+                        .withActions(AdmNotification.Action("Reset Yubikey and Connect...") { _, notification ->
+                            notification.expire()
+
+                            val success = YubikeyHandler().autoReset(project) {
+                                connectWithConnector(connector, onSuccessful)
+                            }
+
+                            if (!success) {
+                                return@Action
+                            }
+                            connectWithConnector(connector, onSuccessful)
+                        })
+                        .withActions(AdmNotification.Action("Connect...") { _, notification ->
+                            notification.expire()
+                            connectWithConnector(connector, onSuccessful)
+                        })
+                        .show()
+                }
+
+                fun handleOpenTcpSessionError() {
+                    val message = "Failed to open SSH channel in a just created TCP session"
+                    resetYubikeyNotification(message)
+                    LOG.warn(
+                        "Failed to connect",
+                        OpenFailException("TCP", OpenFailException.Reason.CONNECT_FAILED, message)
+                    )
+                }
+
                 override fun onCancel() {
                     super.onCancel()
                     cancelled = true