pkl/build-logic/src/main/kotlin/InstallGraalVm.kt

/*
 * Copyright © 2024-2026 Apple Inc. and the Pkl project authors. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     https://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
import java.nio.file.AtomicMoveNotSupportedException
import java.nio.file.Files
import java.nio.file.Path
import java.nio.file.Paths
import java.nio.file.StandardCopyOption
import java.util.*
import java.util.zip.ZipInputStream
import javax.inject.Inject
import kotlin.io.path.createDirectories
import org.gradle.api.DefaultTask
import org.gradle.api.internal.file.FileOperations
import org.gradle.api.provider.Property
import org.gradle.api.tasks.Input
import org.gradle.api.tasks.TaskAction
import org.gradle.process.ExecOperations

abstract class InstallGraalVm
@Inject
constructor(
  private val fileOperations: FileOperations,
  private val execOperations: ExecOperations,
) : DefaultTask() {
  @get:Input abstract val graalVm: Property<BuildInfo.GraalVm>

  init {
    @Suppress("LeakingThis") onlyIf("GraalVM not installed") { !graalVm.get().installDir.exists() }
  }

  @TaskAction
  @Suppress("unused")
  fun run() {
    // minimize chance of corruption by extract-to-random-dir-and-publish
    val distroDir = Paths.get(graalVm.get().homeDir, UUID.randomUUID().toString())
    try {
      distroDir.createDirectories()
      println("Extracting ${graalVm.get().downloadFile} into $distroDir")
      val os = org.gradle.internal.os.OperatingSystem.current()
      if (os.isWindows) extractZip(distroDir) else extractTarGz(distroDir)

      val distroBinDir =
        if (os.isMacOsX) distroDir.resolve("Contents/Home/bin") else distroDir.resolve("bin")

      println("Installing native-image into $distroDir")
      val gvmVersionMajor =
        requireNotNull(graalVm.get().version.split(".").first().toIntOrNull()) {
          "Invalid GraalVM JDK version: ${graalVm.get().graalVmJdkVersion}"
        }
      if (gvmVersionMajor < 24) {
        execOperations.exec {
          val executableName = if (os.isWindows) "gu.cmd" else "gu"
          executable = distroBinDir.resolve(executableName).toString()
          args("install", "--no-progress", "native-image")
        }
      }

      publishInstallDir(distroDir, os)
    } catch (e: Exception) {
      try {
        fileOperations.delete(distroDir)
      } catch (ignored: Exception) {}
      throw e
    }
  }

  private fun extractTarGz(distroDir: Path) {
    // faster and more reliable than Gradle's `copy { from tarTree() }`
    execOperations.exec {
      workingDir = distroDir.toFile()
      executable = "tar"
      args("--strip-components=1", "-xzf", graalVm.get().downloadFile)
    }
  }

  private fun extractZip(distroDir: Path) {
    val targetDir = distroDir.toAbsolutePath().normalize()
    ZipInputStream(graalVm.get().downloadFile.inputStream().buffered()).use { zipInput ->
      var entry = zipInput.nextEntry
      while (entry != null) {
        try {
          val strippedPath = stripFirstPathComponent(entry.name)
          if (strippedPath != null) {
            val target = targetDir.resolve(strippedPath).normalize()
            require(target.startsWith(targetDir)) {
              "GraalVM archive entry escapes destination directory: ${entry.name}"
            }
            if (entry.isDirectory) {
              target.createDirectories()
            } else {
              target.parent?.createDirectories()
              Files.copy(zipInput, target)
            }
          }
        } finally {
          zipInput.closeEntry()
        }
        entry = zipInput.nextEntry
      }
    }
  }

  private fun stripFirstPathComponent(path: String): String? {
    val normalizedPath = path.replace('\\', '/').trimStart('/')
    val separatorIndex = normalizedPath.indexOf('/')
    if (separatorIndex == -1) return null
    return normalizedPath.substring(separatorIndex + 1).takeIf { it.isNotEmpty() }
  }

  private fun publishInstallDir(distroDir: Path, os: org.gradle.internal.os.OperatingSystem) {
    if (os.isWindows) {
      println("Installing ${graalVm.get().installDir} from $distroDir")
      moveAtomicallyOrRegularly(distroDir, graalVm.get().installDir.toPath())
      return
    }

    println("Creating symlink ${graalVm.get().installDir} for $distroDir")
    val tempLink = Paths.get(graalVm.get().homeDir, UUID.randomUUID().toString())
    Files.createSymbolicLink(tempLink, distroDir)
    try {
      moveAtomicallyOrRegularly(tempLink, graalVm.get().installDir.toPath())
    } catch (e: Exception) {
      try {
        fileOperations.delete(tempLink.toFile())
      } catch (ignored: Exception) {}
      throw e
    }
  }

  private fun moveAtomicallyOrRegularly(source: Path, target: Path) {
    try {
      Files.move(source, target, StandardCopyOption.ATOMIC_MOVE)
    } catch (e: AtomicMoveNotSupportedException) {
      Files.move(source, target)
    }
  }
}