WIP: Reuse the same tunnel for an environment

mustard-mh · mustard-mh · commit 4448660270a9 · 2024-12-02T13:51:37.000+08:00
diff --git a/components/ide/jetbrains/toolbox/src/main/kotlin/io/gitpod/toolbox/auth/GitpodAuthManager.kt b/components/ide/jetbrains/toolbox/src/main/kotlin/io/gitpod/toolbox/auth/GitpodAuthManager.kt
@@ -22,7 +22,7 @@ import java.net.URI
 import java.util.*
 import java.util.concurrent.Future
 
-// TODO: Validate Scopes
+// TODO(hw): Validate Scopes
 val authScopesJetBrainsToolbox = listOf(
     "function:getGitpodTokenScopes",
     "function:getLoggedInUser",
diff --git a/components/ide/jetbrains/toolbox/src/main/kotlin/io/gitpod/toolbox/components/AbstractUiPage.kt b/components/ide/jetbrains/toolbox/src/main/kotlin/io/gitpod/toolbox/components/AbstractUiPage.kt
@@ -6,8 +6,6 @@ package io.gitpod.toolbox.components
 
 import com.jetbrains.toolbox.gateway.ui.UiField
 import com.jetbrains.toolbox.gateway.ui.UiPage
-import java.util.function.BiConsumer
-import java.util.function.Function
 
 abstract class AbstractUiPage : UiPage {
     private var stateAccessor: UiPage.UiFieldStateAccessor? = null
diff --git a/components/ide/jetbrains/toolbox/src/main/kotlin/io/gitpod/toolbox/gateway/GitpodRemoteProvider.kt b/components/ide/jetbrains/toolbox/src/main/kotlin/io/gitpod/toolbox/gateway/GitpodRemoteProvider.kt
@@ -63,9 +63,10 @@ class GitpodRemoteProvider(
             environmentMap[connectParams.uniqueID] = Pair(workspace, env)
             consumer.consumeEnvironments(environmentMap.values.map { it.second })
         }
-        val joinLinkInfo = workspace!!.fetchJoinLink2Info(publicApi.getWorkspaceOwnerToken(workspaceId))
+        val joinLinkInfo = publicApi.fetchJoinLink2Info(workspaceId, workspace!!.getIDEUrl())
+        // TODO(hw): verify if it's working
         Utils.clientHelper.prepareClient(joinLinkInfo.ideVersion)
-        Utils.clientHelper.setAutoConnectOnEnvironmentReady(workspaceId, joinLinkInfo.ideVersion, joinLinkInfo.projectPath)
+        Utils.clientHelper.setAutoConnectOnEnvironmentReady(connectParams.uniqueID, joinLinkInfo.ideVersion, joinLinkInfo.projectPath)
     }
 
     private fun showWorkspacesList() {
diff --git a/components/ide/jetbrains/toolbox/src/main/kotlin/io/gitpod/toolbox/gateway/GitpodRemoteProviderEnvironment.kt b/components/ide/jetbrains/toolbox/src/main/kotlin/io/gitpod/toolbox/gateway/GitpodRemoteProviderEnvironment.kt
@@ -6,12 +6,7 @@ package io.gitpod.toolbox.gateway
 
 import com.jetbrains.toolbox.gateway.AbstractRemoteProviderEnvironment
 import com.jetbrains.toolbox.gateway.EnvironmentVisibilityState
-import com.jetbrains.toolbox.gateway.deploy.CommandExecution
-import com.jetbrains.toolbox.gateway.deploy.HostAccessInterfaces
-import com.jetbrains.toolbox.gateway.deploy.HostCommandExecutor
-import com.jetbrains.toolbox.gateway.deploy.HostFileAccessor
 import com.jetbrains.toolbox.gateway.environments.EnvironmentContentsView
-import com.jetbrains.toolbox.gateway.environments.HostAccessCapableEnvironmentContentsView
 import com.jetbrains.toolbox.gateway.states.EnvironmentStateConsumer
 import com.jetbrains.toolbox.gateway.states.StandardRemoteEnvironmentState
 import com.jetbrains.toolbox.gateway.ui.ActionDescription
@@ -36,13 +31,9 @@ class GitpodRemoteProviderEnvironment(
     private val publicApi: GitpodPublicApiManager, observablePropertiesFactory: ObservablePropertiesFactory?,
 ) : AbstractRemoteProviderEnvironment(observablePropertiesFactory), DisposableHandle {
     private val actionList = Utils.observablePropertiesFactory.emptyObservableList<ActionDescription>();
-    private val contentsViewFuture: CompletableFuture<EnvironmentContentsView> = CompletableFuture.completedFuture(
-        GitpodSSHEnvironmentContentsView(
-            authManager,
-            connectParams,
-            publicApi,
-        )
-    )
+    private val envContentsView = GitpodSSHEnvironmentContentsView(connectParams, publicApi)
+    private val contentsViewFuture: CompletableFuture<EnvironmentContentsView> =
+        CompletableFuture.completedFuture(envContentsView)
     private var watchWorkspaceJob: Job? = null
 
     private val lastWSEnvState = MutableSharedFlow<WorkspaceEnvState>(1, 0, BufferOverflow.DROP_OLDEST)
@@ -65,6 +56,9 @@ class GitpodRemoteProviderEnvironment(
                 lastPhase = status.phase
                 GitpodLogger.debug("${connectParams.workspaceId} status updated: $lastPhase")
                 lastWSEnvState.tryEmit(WorkspaceEnvState(status.phase))
+                Utils.coroutineScope.launch {
+                    envContentsView.updateEnvironmentMeta(status)
+                }
             }
         }
     }
@@ -78,19 +72,17 @@ class GitpodRemoteProviderEnvironment(
     }
 
     override fun getId(): String = connectParams.uniqueID
-    override fun getName(): String = connectParams.resolvedWorkspaceId
+    override fun getName(): String = connectParams.uniqueID
 
-    override fun getContentsView(): CompletableFuture<EnvironmentContentsView> {
-        GitpodLogger.info("=============test.getContentView id: $id")
-        return contentsViewFuture
-    }
+    override fun getContentsView(): CompletableFuture<EnvironmentContentsView> = contentsViewFuture
 
     override fun setVisible(visibilityState: EnvironmentVisibilityState) {
     }
 
     override fun getActionList(): ObservableList<ActionDescription> = actionList
 
     override fun onDelete() {
+        // TODO: delete workspace?
         watchWorkspaceJob?.cancel()
     }
 
@@ -120,4 +112,5 @@ private class WorkspaceEnvState(val phase: WorkspaceInstanceStatus.Phase) {
             WorkspaceInstanceStatus.Phase.PHASE_STOPPED to StandardRemoteEnvironmentState.Hibernated,
         )
     }
+    // TODO(hw): add customized state
 }
diff --git a/components/ide/jetbrains/toolbox/src/main/kotlin/io/gitpod/toolbox/gateway/GitpodSSHEnvironmentContentsView.kt b/components/ide/jetbrains/toolbox/src/main/kotlin/io/gitpod/toolbox/gateway/GitpodSSHEnvironmentContentsView.kt
@@ -9,62 +9,63 @@ import com.jetbrains.toolbox.gateway.environments.CachedProjectStub
 import com.jetbrains.toolbox.gateway.environments.ManualEnvironmentContentsView
 import com.jetbrains.toolbox.gateway.environments.SshEnvironmentContentsView
 import com.jetbrains.toolbox.gateway.ssh.SshConnectionInfo
-import io.gitpod.toolbox.auth.GitpodAuthManager
-import io.gitpod.toolbox.service.ConnectParams
-import io.gitpod.toolbox.service.GitpodConnectionProvider
-import io.gitpod.toolbox.service.GitpodPublicApiManager
-import io.gitpod.toolbox.service.Utils
-import io.gitpod.toolbox.utils.GitpodLogger
-import kotlinx.coroutines.future.future
+import io.gitpod.publicapi.experimental.v1.Workspaces.WorkspaceInstanceStatus
+import io.gitpod.toolbox.service.*
 import java.util.concurrent.CompletableFuture
 
 class GitpodSSHEnvironmentContentsView(
-    private val authManager: GitpodAuthManager,
     private val connectParams: ConnectParams,
     private val publicApi: GitpodPublicApiManager,
 ) : SshEnvironmentContentsView, ManualEnvironmentContentsView {
     private var cancel = {}
     private val stateListeners = mutableSetOf<ManualEnvironmentContentsView.Listener>()
+    private val provider = GitpodConnectionProvider(object : ConnectionInfoProvider {
+        override fun getUniqueID() = connectParams.uniqueID
 
-    override fun getConnectionInfo(): CompletableFuture<SshConnectionInfo> {
-        return Utils.coroutineScope.future {
-            val provider = GitpodConnectionProvider(authManager, connectParams, publicApi)
-            val (connInfo, cancel) = provider.connect()
-            this@GitpodSSHEnvironmentContentsView.cancel = cancel
-            GitpodLogger.info("=============test.getConnectionInfo port: ${connInfo.port}")
-            return@future connInfo
+        override suspend fun getWebsocketTunnelUrl(): String {
+            val workspace = publicApi.getWorkspace(connectParams.workspaceId)
+            return workspace.getTunnelUrl()
         }
-    }
 
-    override fun addEnvironmentContentsListener(p0: ManualEnvironmentContentsView.Listener) {
-        stateListeners += p0
-        stateListeners.forEach{
-            // TODO: get from fetchJoinLink2Info
-            it.onProjectListUpdated(listOf(object : CachedProjectStub {
-                override fun getPath(): String {
-                    return "/workspace/template-golang-cli"
-                }
+        override suspend fun getOwnerToken(): String {
+            return publicApi.getWorkspaceOwnerToken(connectParams.workspaceId)
+        }
+    })
 
-                override fun getName(): String? {
-                    return "template-golang-cli"
-                }
+    private val connectionInfo = CompletableFuture.supplyAsync {
+        val (connInfo, cancel) = provider.connect()
+        this.cancel = cancel
+        connInfo
+    }
 
-                override fun getIdeHint(): String? {
-                    return "GO-243.21565.208"
-                }
-            }))
-            it.onIdeListUpdated(listOf(object: CachedIdeStub {
-                override fun getProductCode(): String {
-                    return "GO-243.21565.208"
-                }
+    override fun getConnectionInfo(): CompletableFuture<SshConnectionInfo> = connectionInfo
 
-                override fun isRunning(): Boolean? {
-                    return true
-                }
+    var metadata: GitpodPublicApiManager.JoinLink2Response? = null
+    suspend fun updateEnvironmentMeta(status: WorkspaceInstanceStatus) {
+        if (metadata == null && status.phase == WorkspaceInstanceStatus.Phase.PHASE_RUNNING) {
+            metadata = publicApi.fetchJoinLink2Info(connectParams.workspaceId, status.getIDEUrl())
+        }
+        if (metadata == null) {
+            // TODO(hw): restore from cache?
+            return
+        }
+        stateListeners.forEach {
+            it.onProjectListUpdated(listOf(object : CachedProjectStub {
+                override fun getPath() = metadata!!.projectPath
+                override fun getName() = metadata!!.projectPath.split("/").last()
+                override fun getIdeHint() = metadata!!.ideVersion
+            }))
+            it.onIdeListUpdated(listOf(object : CachedIdeStub {
+                override fun getProductCode() = metadata!!.ideVersion
+                override fun isRunning() = status.phase == WorkspaceInstanceStatus.Phase.PHASE_RUNNING
             }))
         }
     }
 
+    override fun addEnvironmentContentsListener(p0: ManualEnvironmentContentsView.Listener) {
+        stateListeners += p0
+    }
+
     override fun removeEnvironmentContentsListener(p0: ManualEnvironmentContentsView.Listener) {
         stateListeners -= p0
     }
diff --git a/components/ide/jetbrains/toolbox/src/main/kotlin/io/gitpod/toolbox/service/DataManager.kt b/components/ide/jetbrains/toolbox/src/main/kotlin/io/gitpod/toolbox/service/DataManager.kt
@@ -5,11 +5,8 @@
 package io.gitpod.toolbox.service
 
 import io.gitpod.publicapi.experimental.v1.Workspaces.Workspace
-import io.gitpod.toolbox.utils.await
-import kotlinx.serialization.Serializable
-import kotlinx.serialization.decodeFromString
-import kotlinx.serialization.json.Json
-import okhttp3.Request
+import io.gitpod.publicapi.experimental.v1.Workspaces.WorkspaceInstanceStatus
+import java.net.URI
 import java.net.URL
 
 fun Workspace.getConnectParams(): ConnectParams {
@@ -26,20 +23,12 @@ fun Workspace.getGitpodHost(): String {
     return hostSegments.takeLast(2).joinToString(".")
 }
 
-@Serializable
-class JoinLink2Response(val appPid: Int, val joinLink: String, val ideVersion: String, val projectPath: String)
-
-suspend fun Workspace.fetchJoinLink2Info(ownerToken: String): JoinLink2Response {
-    val backendUrl = "https://24000-${URL(getIDEUrl()).host}/joinLink2"
-    val client = Utils.httpClient
-    val req = Request.Builder().url(backendUrl).header("x-gitpod-owner-token", ownerToken)
-    val response = client.newCall(req.build()).await()
-    if (!response.isSuccessful) {
-        throw IllegalStateException("Failed to get join link $backendUrl info: ${response.code} ${response.message}")
-    }
-    if (response.body == null) {
-        throw IllegalStateException("Failed to get join link $backendUrl info: no body")
-    }
-    return Json.decodeFromString<JoinLink2Response>(response.body!!.string())
+fun Workspace.getTunnelUrl(): String {
+    val workspaceHost = URI.create(status.instance.status.url).host
+    return "wss://${workspaceHost}/_supervisor/tunnel/ssh"
+}
+
+fun WorkspaceInstanceStatus.getIDEUrl(): String {
+    return url
 }
 
diff --git a/components/ide/jetbrains/toolbox/src/main/kotlin/io/gitpod/toolbox/service/GitpodConnectionProvider.kt b/components/ide/jetbrains/toolbox/src/main/kotlin/io/gitpod/toolbox/service/GitpodConnectionProvider.kt
@@ -5,26 +5,19 @@
 package io.gitpod.toolbox.service
 
 import com.jetbrains.rd.util.ConcurrentHashMap
-import com.jetbrains.rd.util.URI
 import com.jetbrains.toolbox.gateway.ssh.SshConnectionInfo
-import io.gitpod.publicapi.experimental.v1.Workspaces
-import io.gitpod.toolbox.auth.GitpodAuthManager
-import io.gitpod.toolbox.utils.GitpodLogger
-import kotlinx.serialization.Serializable
 
-class GitpodConnectionProvider(
-    private val authManager: GitpodAuthManager,
-    private val connectParams: ConnectParams,
-    private val publicApi: GitpodPublicApiManager,
-) {
-    private val activeConnections = ConcurrentHashMap<String, Boolean>()
+interface ConnectionInfoProvider {
+    fun getUniqueID(): String
+    suspend fun getWebsocketTunnelUrl(): String
+    suspend fun getOwnerToken(): String
+}
 
-    suspend fun connect(): Pair<SshConnectionInfo, () -> Unit> {
-        val workspaceId = connectParams.workspaceId
-        val workspace = publicApi.getWorkspace(workspaceId)
-        val ownerToken = publicApi.getWorkspaceOwnerToken(workspaceId)
+class GitpodConnectionProvider(private val provider: ConnectionInfoProvider) {
+    private val activeConnections = ConcurrentHashMap<String, Boolean>()
 
-        val (serverPort, cancel) = tunnelWithWebSocket(workspace, connectParams, ownerToken)
+    fun connect(): Pair<SshConnectionInfo, () -> Unit> {
+        val (serverPort, cancel) = tunnelWithWebSocket(provider)
 
         val connInfo = GitpodWebSocketSshConnectionInfo(
             "gitpod",
@@ -34,12 +27,8 @@ class GitpodConnectionProvider(
         return (connInfo to cancel)
     }
 
-    private fun tunnelWithWebSocket(
-        workspace: Workspaces.Workspace,
-        connectParams: ConnectParams,
-        ownerToken: String,
-    ): Pair<Int, () -> Unit> {
-        val connectionKeyId = connectParams.uniqueID
+    private fun tunnelWithWebSocket(provider: ConnectionInfoProvider): Pair<Int, () -> Unit> {
+        val connectionKeyId = provider.getUniqueID()
 
         var found = true
         activeConnections.computeIfAbsent(connectionKeyId) {
@@ -52,9 +41,7 @@ class GitpodConnectionProvider(
             throw IllegalStateException(errMessage)
         }
 
-        val workspaceHost = URI.create(workspace.status.instance.status.url).host
-        val server =
-            GitpodWebSocketTunnelServer("wss://${workspaceHost}/_supervisor/tunnel/ssh", ownerToken)
+        val server = GitpodWebSocketTunnelServer(provider)
 
         val cancelServer = server.start()
 
@@ -82,13 +69,5 @@ data class ConnectParams(
     val host: String,
     val debugWorkspace: Boolean = false,
 ) {
-    val resolvedWorkspaceId = "${if (debugWorkspace) "debug-" else ""}$workspaceId"
-    val title = "$resolvedWorkspaceId"
-    val uniqueID = "$workspaceId-$debugWorkspace"
+    val uniqueID = if (debugWorkspace) "debug-$workspaceId" else workspaceId
 }
-
-@Serializable
-private data class SSHPublicKey(
-    val type: String,
-    val value: String
-)
diff --git a/components/ide/jetbrains/toolbox/src/main/kotlin/io/gitpod/toolbox/service/GitpodPublicApiManager.kt b/components/ide/jetbrains/toolbox/src/main/kotlin/io/gitpod/toolbox/service/GitpodPublicApiManager.kt
@@ -14,8 +14,14 @@ import io.gitpod.publicapi.experimental.v1.*
 import io.gitpod.toolbox.auth.GitpodAccount
 import io.gitpod.toolbox.auth.GitpodAuthManager
 import io.gitpod.toolbox.utils.GitpodLogger
+import io.gitpod.toolbox.utils.await
 import kotlinx.coroutines.Job
 import kotlinx.coroutines.launch
+import kotlinx.serialization.Serializable
+import kotlinx.serialization.decodeFromString
+import kotlinx.serialization.json.Json
+import okhttp3.Request
+import java.net.URL
 import java.time.Duration
 
 class GitpodPublicApiManager(private val authManger: GitpodAuthManager) {
@@ -81,6 +87,24 @@ class GitpodPublicApiManager(private val authManger: GitpodAuthManager) {
         return this.handleResp("getOwnerToken", resp).token
     }
 
+    @Serializable
+    class JoinLink2Response(val appPid: Int, val joinLink: String, val ideVersion: String, val projectPath: String)
+
+    suspend fun fetchJoinLink2Info(workspaceId: String, ideURL: String): JoinLink2Response {
+        val token = getWorkspaceOwnerToken(workspaceId)
+        val backendUrl = "https://24000-${URL(ideURL).host}/joinLink2"
+        val client = Utils.httpClient
+        val req = Request.Builder().url(backendUrl).header("x-gitpod-owner-token", token)
+        val response = client.newCall(req.build()).await()
+        if (!response.isSuccessful) {
+            throw IllegalStateException("Failed to get join link $backendUrl info: ${response.code} ${response.message}")
+        }
+        if (response.body == null) {
+            throw IllegalStateException("Failed to get join link $backendUrl info: no body")
+        }
+        return Json.decodeFromString<JoinLink2Response>(response.body!!.string())
+    }
+
     suspend fun getAuthenticatedUser(): UserOuterClass.User {
         return tryGetAuthenticatedUser(userApi)
     }
diff --git a/components/ide/jetbrains/toolbox/src/main/kotlin/io/gitpod/toolbox/service/GitpodWebSocketTunnelServer.kt b/components/ide/jetbrains/toolbox/src/main/kotlin/io/gitpod/toolbox/service/GitpodWebSocketTunnelServer.kt

Original file line number	Diff line number	Diff line change
`@@ -63,9 +63,10 @@ class GitpodRemoteProvider(`
`63`	`63`	`environmentMap[connectParams.uniqueID] = Pair(workspace, env)`
`64`	`64`	`consumer.consumeEnvironments(environmentMap.values.map { it.second })`
`65`	`65`	`}`
`66`		`- val joinLinkInfo = workspace!!.fetchJoinLink2Info(publicApi.getWorkspaceOwnerToken(workspaceId))`
	`66`	`+ val joinLinkInfo = publicApi.fetchJoinLink2Info(workspaceId, workspace!!.getIDEUrl())`
	`67`	`+ // TODO(hw): verify if it's working`
`67`	`68`	`Utils.clientHelper.prepareClient(joinLinkInfo.ideVersion)`
`68`		`- Utils.clientHelper.setAutoConnectOnEnvironmentReady(workspaceId, joinLinkInfo.ideVersion, joinLinkInfo.projectPath)`
	`69`	`+ Utils.clientHelper.setAutoConnectOnEnvironmentReady(connectParams.uniqueID, joinLinkInfo.ideVersion, joinLinkInfo.projectPath)`
`69`	`70`	`}`
`70`	`71`
`71`	`72`	`private fun showWorkspacesList() {`