fix(jetbrains): avoid connection loops

gitpod-io · Aug 31, 2022 · 0e7c17c · 0e7c17c
1 parent db6f5ed
commit 0e7c17c
Showing 1 changed file with 26 additions and 2 deletions.
diff --git a/...ns/gateway-plugin/src/main/kotlin/io/gitpod/jetbrains/gateway/GitpodConnectionProvider.kt b/...ns/gateway-plugin/src/main/kotlin/io/gitpod/jetbrains/gateway/GitpodConnectionProvider.kt
@@ -12,6 +12,7 @@ import com.intellij.ide.BrowserUtil
 import com.intellij.openapi.components.service
 import com.intellij.openapi.diagnostic.thisLogger
 import com.intellij.openapi.progress.ProgressManager
+import com.intellij.openapi.ui.Messages
 import com.intellij.remote.RemoteCredentialsHolder
 import com.intellij.ssh.AskAboutHostKey
 import com.intellij.ssh.OpenSshLikeHostKeyVerifier
@@ -30,6 +31,7 @@ import com.jetbrains.gateway.api.*
 import com.jetbrains.gateway.ssh.ClientOverSshTunnelConnector
 import com.jetbrains.gateway.ssh.SshHostTunnelConnector
 import com.jetbrains.gateway.thinClientLink.ThinClientHandle
+import com.jetbrains.rd.util.ConcurrentHashMap
 import com.jetbrains.rd.util.URI
 import com.jetbrains.rd.util.lifetime.Lifetime
 import io.gitpod.gitpodprotocol.api.entities.WorkspaceInstance
@@ -49,7 +51,7 @@ import kotlin.coroutines.coroutineContext
 
 @Suppress("UnstableApiUsage", "OPT_IN_USAGE")
 class GitpodConnectionProvider : GatewayConnectionProvider {
-
+    private val activeConnections = ConcurrentHashMap<String, Boolean>()
     private val gitpod = service<GitpodConnectionService>()
 
     private val httpClient = HttpClient.newBuilder()
@@ -58,7 +60,7 @@ class GitpodConnectionProvider : GatewayConnectionProvider {
 
     private val jacksonMapper = jacksonObjectMapper()
         .configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false)
-
+    
     override suspend fun connect(
         parameters: Map<String, String>,
         requestor: ConnectionRequestor
@@ -74,6 +76,7 @@ class GitpodConnectionProvider : GatewayConnectionProvider {
             parameters["workspaceId"]!!,
             parameters["backendPort"]
         )
+
         val client = gitpod.obtainClient(connectParams.gitpodHost)
         val connectionLifetime = Lifetime.Eternal.createNested()
         val updates = client.listenToWorkspace(connectionLifetime, connectParams.workspaceId)
@@ -144,6 +147,24 @@ class GitpodConnectionProvider : GatewayConnectionProvider {
             connectionPanel.repaint()
         }
 
+        val connectionKeyId = "${connectParams.gitpodHost}-${connectParams.workspaceId}-${connectParams.backendPort}"
+
+        if (activeConnections.containsKey(connectionKeyId)) {
+            val message = "You are trying to connect to a workspace that has a session already open"
+            val title = "Oops"
+            val okButton = Messages.getOkButton()
+            val options = arrayOf(okButton)
+            val defaultIndex = 0
+            val icon = Messages.getWarningIcon()
+            Messages.showDialog(message, title, options, defaultIndex, icon)
+
+            val errMessage = "A connection to the same workspace already exists: $connectionKeyId"
+            thisLogger().warn(errMessage)
+            throw IllegalStateException(errMessage)
+        } else {
+            activeConnections.putIfAbsent(connectionKeyId, true)
+        }
+
         GlobalScope.launch {
             var thinClient: ThinClientHandle? = null
             var thinClientJob: Job? = null
@@ -219,6 +240,7 @@ class GitpodConnectionProvider : GatewayConnectionProvider {
                                     val sshHostUrl = URL(update.ideUrl.replace(update.workspaceId, "${update.workspaceId}.ssh"));
                                     val hostKeys = resolveHostKeys(updatedIdeUrl, connectParams)
                                     if (hostKeys.isNullOrEmpty()) {
+                                        activeConnections.remove(connectionKeyId)
                                         setErrorMessage("${connectParams.gitpodHost} installation does not allow SSH access, public keys cannot be found")
                                         return@launch
                                     }
@@ -227,6 +249,7 @@ class GitpodConnectionProvider : GatewayConnectionProvider {
                                         resolveCredentials(sshHostUrl, update.workspaceId, ownerToken, hostKeys)
                                     val joinLink = resolveJoinLink(updatedIdeUrl, ownerToken, connectParams)
                                     if (joinLink.isNullOrEmpty()) {
+                                        activeConnections.remove(connectionKeyId)
                                         setErrorMessage("failed to fetch JetBrains Gateway Join Link.")
                                         return@launch
                                     }
@@ -239,6 +262,7 @@ class GitpodConnectionProvider : GatewayConnectionProvider {
                                     clientHandle.clientClosed.advise(connectionLifetime) {
                                         application.invokeLater {
                                             connectionLifetime.terminate()
+                                            activeConnections.remove(connectionKeyId)
                                         }
                                     }
                                     clientHandle.onClientPresenceChanged.advise(connectionLifetime) {