Use ETag for downloading binary

code-asher · code-asher · commit ae7b8dfee821 · 2023-04-05T11:06:49.000-08:00
This way we will not need to name the binaries with their versions or
otherwise extract the version from each binary since we can just grab
their hash and set that in the request.  This also means we will not
need to remove old versions as there will only be the one file.
diff --git a/src/main/kotlin/com/coder/gateway/sdk/CoderCLIManager.kt b/src/main/kotlin/com/coder/gateway/sdk/CoderCLIManager.kt
@@ -3,22 +3,29 @@ package com.coder.gateway.sdk
 import com.coder.gateway.views.steps.CoderWorkspacesStepView
 import com.intellij.openapi.diagnostic.Logger
 import org.zeroturnaround.exec.ProcessExecutor
-import java.io.InputStream
+import java.io.BufferedInputStream
+import java.io.FileInputStream
+import java.io.FileNotFoundException
+import java.net.HttpURLConnection
 import java.net.URL
 import java.nio.file.Files
 import java.nio.file.Path
 import java.nio.file.Paths
+import java.nio.file.StandardCopyOption
 import java.nio.file.attribute.PosixFilePermissions
+import java.security.DigestInputStream
+import java.security.MessageDigest
+import javax.xml.bind.annotation.adapters.HexBinaryAdapter
+
 
 /**
  * Manage the CLI for a single deployment.
  */
-class CoderCLIManager(deployment: URL, buildVersion: String) {
+class CoderCLIManager(deployment: URL) {
     private var remoteBinaryUrl: URL
     var localBinaryPath: Path
     private var binaryNamePrefix: String
     private var destinationDir: Path
-    private var localBinaryName: String
 
     init {
         // TODO: Should use a persistent path to avoid needing to download on
@@ -28,23 +35,21 @@ class CoderCLIManager(deployment: URL, buildVersion: String) {
         val os = getOS()
         binaryNamePrefix = getCoderCLIForOS(os, getArch())
         val binaryName = if (os == OS.WINDOWS) "$binaryNamePrefix.exe" else binaryNamePrefix
-        localBinaryName =
-            if (os == OS.WINDOWS) "${binaryNamePrefix}-${buildVersion}.exe" else "${binaryNamePrefix}-${buildVersion}"
         remoteBinaryUrl = URL(
             deployment.protocol,
             deployment.host,
             deployment.port,
             "/bin/$binaryName"
         )
-        localBinaryPath = destinationDir.resolve(localBinaryName)
+        localBinaryPath = destinationDir.resolve(binaryName)
     }
 
     /**
      * Return the name of the binary (sans extension) for the provided OS and
      * architecture.
      */
     private fun getCoderCLIForOS(os: OS?, arch: Arch?): String {
-        logger.info("Resolving coder cli for $os $arch")
+        logger.info("Resolving binary for $os $arch")
         if (os == null) {
             logger.error("Could not resolve client OS and architecture, defaulting to WINDOWS AMD64")
             return "coder-windows-amd64"
@@ -75,42 +80,59 @@ class CoderCLIManager(deployment: URL, buildVersion: String) {
      * Download the CLI from the deployment if necessary.
      */
     fun downloadCLI(): Boolean {
-        Files.createDirectories(destinationDir)
-        try {
-            logger.info("Downloading Coder CLI to ${localBinaryPath.toAbsolutePath()}")
-            remoteBinaryUrl.openStream().use {
-                Files.copy(it as InputStream, localBinaryPath)
-            }
-        } catch (e: java.nio.file.FileAlreadyExistsException) {
-            // This relies on the provided build version being the latest.  It
-            // must be freshly fetched immediately before downloading.
-            // TODO: Use etags instead?
-            logger.info("${localBinaryPath.toAbsolutePath()} already exists, skipping download")
-            return false
+        val etag = getBinaryETag()
+        val conn = remoteBinaryUrl.openConnection() as HttpURLConnection
+        if (etag != null) {
+            logger.info("Found existing binary at ${localBinaryPath.toAbsolutePath()}; calculated hash as $etag")
+            conn.setRequestProperty("If-None-Match", "\"$etag\"")
         }
-        if (getOS() != OS.WINDOWS) {
-            Files.setPosixFilePermissions(
-                localBinaryPath,
-                PosixFilePermissions.fromString("rwxr-x---")
-            )
+        conn.connect()
+        logger.info("GET ${conn.responseCode} $remoteBinaryUrl")
+        when (conn.responseCode) {
+            200 -> {
+                logger.info("Downloading binary to ${localBinaryPath.toAbsolutePath()}")
+                Files.createDirectories(destinationDir)
+                conn.inputStream.use {
+                    Files.copy(it, localBinaryPath, StandardCopyOption.REPLACE_EXISTING)
+                }
+                if (getOS() != OS.WINDOWS) {
+                    Files.setPosixFilePermissions(
+                        localBinaryPath,
+                        PosixFilePermissions.fromString("rwxr-x---")
+                    )
+                }
+                conn.disconnect()
+                return true
+            }
+
+            304 -> {
+                logger.info("Using cached binary at ${localBinaryPath.toAbsolutePath()}")
+                conn.disconnect()
+                return false
+            }
         }
-        return true
+        conn.disconnect()
+        throw Exception("Unable to download $remoteBinaryUrl (got response code `${conn.responseCode}`)")
     }
 
     /**
-     * Remove all versions of the CLI for this deployment that do not match the
-     * current build version.
+     * Return the entity tag for the binary on disk, if any.
      */
-    fun removeOldCli() {
-        if (Files.isReadable(destinationDir)) {
-            Files.walk(destinationDir, 1).use {
-                it.sorted().map { pt -> pt.toFile() }
-                    .filter { fl -> fl.name.contains(binaryNamePrefix) && fl.name != localBinaryName }
-                    .forEach { fl ->
-                        logger.info("Removing $fl because it is an old version")
-                        fl.delete()
-                    }
+    private fun getBinaryETag(): String? {
+        try {
+            val md = MessageDigest.getInstance("SHA-1")
+            val fis = FileInputStream(localBinaryPath.toFile())
+            val dis = DigestInputStream(BufferedInputStream(fis), md)
+            fis.use {
+                while (dis.read() != -1) {
+                }
             }
+            return HexBinaryAdapter().marshal(md.digest()).lowercase()
+        } catch (e: FileNotFoundException) {
+            return null
+        } catch (e: Exception) {
+            logger.warn("Unable to calculate hash for ${localBinaryPath.toAbsolutePath()}", e)
+            return null
         }
     }
 
diff --git a/src/main/kotlin/com/coder/gateway/views/steps/CoderWorkspacesStepView.kt b/src/main/kotlin/com/coder/gateway/views/steps/CoderWorkspacesStepView.kt
@@ -425,7 +425,7 @@ class CoderWorkspacesStepView(val setNextButtonEnabled: (Boolean) -> Unit) : Cod
                 return@launchUnderBackgroundProgress
             }
 
-            val cliManager = CoderCLIManager(localWizardModel.coderURL.toURL(), coderClient.buildVersion)
+            val cliManager = CoderCLIManager(localWizardModel.coderURL.toURL())
             localWizardModel.token = token
 
             this.indicator.apply {
@@ -459,12 +459,6 @@ class CoderWorkspacesStepView(val setNextButtonEnabled: (Boolean) -> Unit) : Cod
             }
             cliManager.configSsh()
 
-            this.indicator.apply {
-                text = "Removing old Coder CLI versions..."
-                fraction = 0.9
-            }
-            cliManager.removeOldCli()
-
             this.indicator.fraction = 1.0
             updateWorkspaceActions()
             triggerWorkspacePolling(false)
diff --git a/src/test/groovy/CoderCLIManagerTest.groovy b/src/test/groovy/CoderCLIManagerTest.groovy
@@ -8,68 +8,59 @@ import java.nio.file.Path
 
 @Unroll
 class CoderCLIManagerTest extends spock.lang.Specification {
-    def "deletes old versions"() {
+    // TODO: Probably not good to depend on dev.coder.com being up...should use
+    //       a mock?  Or spin up a Coder deployment in CI?
+    def "downloads a working cli"() {
         given:
-        // Simulate downloading an old version.
-        def oldVersion = new CoderCLIManager(new URL("https://test.coder.invalid"), "0.0.1").localBinaryPath.toFile()
-        def dir = oldVersion.toPath().getParent()
-        dir.toFile().deleteDir()
-        Files.createDirectories(dir)
-        oldVersion.write("old-version")
-
-        // Simulate downloading a new version.
-        def ccm = new CoderCLIManager(new URL("https://test.coder.invalid"), "1.0.2")
-        def newVersion = ccm.localBinaryPath.toFile()
-        newVersion.write("new-version")
-
-        // Anything that does not start with the expected prefix is ignored.
-        def otherOsVersion = dir.resolve("coder-alpine-amd64-1.0.2").toFile()
-        otherOsVersion.write("alpine")
-
-        // Anything else matching the prefix is deleted.
-        def invalidVersion = dir.resolve(newVersion.getName() + "-extra-random-text").toFile()
-        invalidVersion.write("invalid")
+        def ccm = new CoderCLIManager(new URL("https://dev.coder.com"))
+        ccm.localBinaryPath.getParent().toFile().deleteDir()
 
         when:
-        ccm.removeOldCli()
+        def downloaded = ccm.downloadCLI()
 
         then:
-        !oldVersion.exists()
-        newVersion.exists()
-        otherOsVersion.exists()
-        !invalidVersion.exists()
+        downloaded
+        ccm.version().contains("Coder")
     }
 
-    // TODO: Probably not good to depend on dev.coder.com being up...should use
-    //       a mock?  Or spin up a Coder deployment in CI?
-    def "downloads a working cli"() {
+    def "overwrites cli if incorrect version"() {
         given:
-        def ccm = new CoderCLIManager(new URL("https://dev.coder.com"), "1.0.1")
-        def dir = Path.of(System.getProperty("java.io.tmpdir")).resolve("coder-gateway/dev.coder.com")
-        ccm.localBinaryPath.getParent().toFile().deleteDir()
+        def ccm = new CoderCLIManager(new URL("https://dev.coder.com"))
+        Files.createDirectories(ccm.localBinaryPath.getParent())
+        ccm.localBinaryPath.toFile().write("cli")
 
         when:
         def downloaded = ccm.downloadCLI()
 
         then:
         downloaded
-        ccm.localBinaryPath.getParent() == dir
-        ccm.localBinaryPath.toFile().exists()
         ccm.version().contains("Coder")
     }
 
     def "skips cli download if it already exists"() {
         given:
-        def ccm = new CoderCLIManager(new URL("https://dev.coder.com"), "1.0.1")
-        Files.createDirectories(ccm.localBinaryPath.getParent())
-        ccm.localBinaryPath.toFile().write("cli")
+        def ccm = new CoderCLIManager(new URL("https://dev.coder.com"))
 
         when:
+        ccm.downloadCLI()
         def downloaded = ccm.downloadCLI()
 
         then:
         !downloaded
-        ccm.localBinaryPath.toFile().readLines() == ["cli"]
+        ccm.version().contains("Coder")
+    }
+
+    def "does not clobber other deployments"() {
+        given:
+        def ccm1 = new CoderCLIManager(new URL("https://oss.demo.coder.com"))
+        def ccm2 = new CoderCLIManager(new URL("https://dev.coder.com"))
+
+        when:
+        ccm1.downloadCLI()
+        ccm2.downloadCLI()
+
+        then:
+        ccm1.localBinaryPath != ccm2.localBinaryPath
     }
 
     /**
