Add test for proxy authorization

code-asher · code-asher · commit 250db54cf551 · 2024-02-09T20:13:46.000-09:00
diff --git a/src/main/kotlin/com/coder/gateway/CoderGatewayConnectionProvider.kt b/src/main/kotlin/com/coder/gateway/CoderGatewayConnectionProvider.kt
@@ -6,6 +6,8 @@ import com.coder.gateway.models.TokenSource
 import com.coder.gateway.models.WorkspaceAgentModel
 import com.coder.gateway.sdk.CoderCLIManager
 import com.coder.gateway.sdk.CoderRestClient
+import com.coder.gateway.sdk.defaultProxy
+import com.coder.gateway.sdk.defaultVersion
 import com.coder.gateway.sdk.ex.AuthenticationResponseException
 import com.coder.gateway.sdk.toURL
 import com.coder.gateway.sdk.v2.models.Workspace
@@ -140,7 +142,7 @@ class CoderGatewayConnectionProvider : GatewayConnectionProvider {
         if (token == null) { // User aborted.
             throw IllegalArgumentException("Unable to connect to $deploymentURL, $TOKEN is missing")
         }
-        val client = CoderRestClient(deploymentURL, token.first, null, settings)
+        val client = CoderRestClient(deploymentURL, token.first, defaultVersion(), settings, defaultProxy())
         return try {
             Pair(client, client.me().username)
         } catch (ex: AuthenticationResponseException) {
diff --git a/src/main/kotlin/com/coder/gateway/sdk/CoderRestClientService.kt b/src/main/kotlin/com/coder/gateway/sdk/CoderRestClientService.kt
@@ -34,6 +34,7 @@ import java.io.File
 import java.io.FileInputStream
 import java.net.HttpURLConnection.HTTP_CREATED
 import java.net.InetAddress
+import java.net.ProxySelector
 import java.net.Socket
 import java.net.URL
 import java.nio.file.Path
@@ -75,63 +76,91 @@ class CoderRestClientService {
      * @throws [AuthenticationResponseException] if authentication failed.
      */
     fun initClientSession(url: URL, token: String, settings: CoderSettingsState): User {
-        client = CoderRestClient(url, token, null, settings)
+        client = CoderRestClient(url, token, defaultVersion(), settings, defaultProxy())
         me = client.me()
         buildVersion = client.buildInfo().version
         isReady = true
         return me
     }
 }
 
-class CoderRestClient(
+/**
+ * Holds proxy information.  Exists only to interface with tests since they
+ * cannot create an HttpConfigurable instance.
+ */
+data class ProxyValues (
+    val username: String?,
+    val password: String?,
+    val useAuth: Boolean,
+    val selector: ProxySelector,
+)
+
+fun defaultProxy(): ProxyValues {
+    val inst = HttpConfigurable.getInstance()
+    return ProxyValues(
+        inst.proxyLogin,
+        inst.plainProxyPassword,
+        inst.PROXY_AUTHENTICATION,
+        inst.onlyBySettingsSelector
+    )
+}
+
+fun defaultVersion(): String {
+    // This is the id from the plugin.xml.
+    return PluginManagerCore.getPlugin(PluginId.getId("com.coder.gateway"))!!.version
+}
+
+class CoderRestClient @JvmOverloads constructor(
     var url: URL, var token: String,
-    private var pluginVersion: String?,
-    private var settings: CoderSettingsState,
+    private val pluginVersion: String,
+    private val settings: CoderSettingsState,
+    private val proxyValues: ProxyValues? = null,
 ) {
-    private var httpClient: OkHttpClient
-    private var retroRestClient: CoderV2RestFacade
+    private val httpClient: OkHttpClient
+    private val retroRestClient: CoderV2RestFacade
 
     init {
         val gson: Gson = GsonBuilder().registerTypeAdapter(Instant::class.java, InstantConverter()).setPrettyPrinting().create()
-        if (pluginVersion.isNullOrBlank()) {
-            pluginVersion = PluginManagerCore.getPlugin(PluginId.getId("com.coder.gateway"))!!.version // this is the id from the plugin.xml
-        }
-
-        val proxy = HttpConfigurable.getInstance()
 
         val socketFactory = coderSocketFactory(settings)
         val trustManagers = coderTrustManagers(settings.tlsCAPath)
-        httpClient = OkHttpClient.Builder()
-            .proxySelector(proxy.onlyBySettingsSelector)
-            .proxyAuthenticator { _, response ->
-                val login = proxy.proxyLogin
-                val pass = proxy.plainProxyPassword
-                if (proxy.PROXY_AUTHENTICATION && login != null && pass != null) {
-                    val credentials = Credentials.basic(login, pass)
-                    response.request.newBuilder()
-                        .header("Proxy-Authorization", credentials)
-                        .build()
-                } else null
-            }
+        var builder = OkHttpClient.Builder()
+
+        if (proxyValues != null) {
+            builder = builder
+                .proxySelector(proxyValues.selector)
+                .proxyAuthenticator { _, response ->
+                    if (proxyValues.useAuth && proxyValues.username != null && proxyValues.password != null) {
+                        val credentials = Credentials.basic(proxyValues.username, proxyValues.password)
+                        response.request.newBuilder()
+                            .header("Proxy-Authorization", credentials)
+                            .build()
+                    } else null
+                }
+        }
+
+        httpClient = builder
             .sslSocketFactory(socketFactory, trustManagers[0] as X509TrustManager)
             .hostnameVerifier(CoderHostnameVerifier(settings.tlsAlternateHostname))
             .addInterceptor { it.proceed(it.request().newBuilder().addHeader("Coder-Session-Token", token).build()) }
             .addInterceptor { it.proceed(it.request().newBuilder().addHeader("User-Agent", "Coder Gateway/${pluginVersion} (${SystemInfo.getOsNameAndVersion()}; ${SystemInfo.OS_ARCH})").build()) }
             .addInterceptor {
                 var request = it.request()
                 val headers = getHeaders(url, settings.headerCommand)
-                if (headers.size > 0) {
-                    val builder = request.newBuilder()
-                    headers.forEach { h -> builder.addHeader(h.key, h.value) }
-                    request = builder.build()
+                if (headers.isNotEmpty()) {
+                    val reqBuilder = request.newBuilder()
+                    headers.forEach { h -> reqBuilder.addHeader(h.key, h.value) }
+                    request = reqBuilder.build()
                 }
                 it.proceed(request)
             }
-            // this should always be last if we want to see previous interceptors logged
+            // This should always be last if we want to see previous interceptors logged.
             .addInterceptor(HttpLoggingInterceptor().apply { setLevel(HttpLoggingInterceptor.Level.BASIC) })
             .build()
 
-        retroRestClient = Retrofit.Builder().baseUrl(url.toString()).client(httpClient).addConverterFactory(GsonConverterFactory.create(gson)).build().create(CoderV2RestFacade::class.java)
+        retroRestClient = Retrofit.Builder().baseUrl(url.toString()).client(httpClient)
+            .addConverterFactory(GsonConverterFactory.create(gson))
+            .build().create(CoderV2RestFacade::class.java)
     }
 
     /**
diff --git a/src/main/kotlin/com/coder/gateway/views/CoderGatewayRecentWorkspaceConnectionsView.kt b/src/main/kotlin/com/coder/gateway/views/CoderGatewayRecentWorkspaceConnectionsView.kt
@@ -9,6 +9,8 @@ import com.coder.gateway.icons.CoderIcons
 import com.coder.gateway.models.RecentWorkspaceConnection
 import com.coder.gateway.models.WorkspaceAgentModel
 import com.coder.gateway.sdk.CoderRestClient
+import com.coder.gateway.sdk.defaultProxy
+import com.coder.gateway.sdk.defaultVersion
 import com.coder.gateway.sdk.toURL
 import com.coder.gateway.sdk.v2.models.WorkspaceStatus
 import com.coder.gateway.sdk.v2.models.toAgentModels
@@ -254,7 +256,7 @@ class CoderGatewayRecentWorkspaceConnectionsView(private val setContentCallback:
                 deployments[dir] ?: try {
                     val url = Path.of(dir).resolve("url").toFile().readText()
                     val token = Path.of(dir).resolve("session").toFile().readText()
-                    DeploymentInfo(CoderRestClient(url.toURL(), token, null, settings))
+                    DeploymentInfo(CoderRestClient(url.toURL(), token, defaultVersion(), settings, defaultProxy()))
                 } catch (e: Exception) {
                     logger.error("Unable to create client from $dir", e)
                     DeploymentInfo(error = "Error trying to read $dir: ${e.message}")
diff --git a/src/test/groovy/CoderRestClientTest.groovy b/src/test/groovy/CoderRestClientTest.groovy
@@ -26,6 +26,7 @@ import java.time.Instant
 @Unroll
 class CoderRestClientTest extends Specification {
     private CoderSettingsState settings = new CoderSettingsState()
+
     /**
      * Create, start, and return a server that mocks the Coder API.
      *
@@ -99,6 +100,48 @@ class CoderRestClientTest extends Specification {
         return [srv, "https://localhost:" + srv.address.port]
     }
 
+    def mockProxy() {
+        HttpServer srv = HttpServer.create(new InetSocketAddress(0), 0)
+        srv.createContext("/", new HttpHandler() {
+            void handle(HttpExchange exchange) {
+                int code
+                String response
+
+                if (exchange.requestHeaders.getFirst("Proxy-Authorization") != "Basic Zm9vOmJhcg==") {
+                    code = HttpURLConnection.HTTP_PROXY_AUTH
+                    response = "authentication required"
+                } else {
+                    try {
+                        HttpURLConnection conn = new URL(exchange.getRequestURI().toString()).openConnection()
+                        exchange.requestHeaders.each{
+                            conn.setRequestProperty(it.key, it.value.join(","))
+                        }
+                        BufferedReader br = new BufferedReader(new InputStreamReader(conn.inputStream))
+                        StringBuilder responseBuilder = new StringBuilder();
+                        String line
+                        while ((line = br.readLine()) != null) {
+                            responseBuilder.append(line)
+                        }
+                        br.close()
+                        response = responseBuilder.toString()
+                        code = conn.responseCode
+                    } catch (Exception error) {
+                        code = HttpURLConnection.HTTP_INTERNAL_ERROR
+                        response = error.message
+                        println(error) // Print since it will not show up in the error.
+                    }
+                }
+
+                byte[] body = response.getBytes()
+                exchange.sendResponseHeaders(code, body.length)
+                exchange.responseBody.write(body)
+                exchange.close()
+            }
+        })
+        srv.start()
+        return srv
+    }
+
     def "gets workspaces"() {
         given:
         def (srv, url) = mockServer(workspaces)
@@ -278,4 +321,33 @@ class CoderRestClientTest extends Specification {
         cleanup:
         srv.stop(0)
     }
+
+    def "uses proxy"() {
+        given:
+        def (srv1, url1) = mockServer([DataGen.workspace("ws1")])
+        def srv2 = mockProxy()
+        def client = new CoderRestClient(new URL(url1), "token", "test", settings, new ProxyValues(
+                "foo",
+                "bar",
+                true,
+                new ProxySelector() {
+                    @Override
+                    List<Proxy> select(URI uri) {
+                        return [new Proxy(Proxy.Type.HTTP, new InetSocketAddress("localhost", srv2.address.port))]
+                    }
+
+                    @Override
+                    void connectFailed(URI uri, SocketAddress sa, IOException ioe) {
+                        getDefault().connectFailed(uri, sa, ioe);
+                    }
+                }
+        ))
+
+        expect:
+        client.workspaces()*.name == ["ws1"]
+
+        cleanup:
+        srv1.stop(0)
+        srv2.stop(0)
+    }
 }
