Detect assemblies with too many entries to fail shell script prepending

example/basic/4-builtin-commands/build.sc

@@ -202,7 +202,8 @@ foo.sources
 foo.allSources
 foo.allSourceFiles
 foo.compile
-foo.localClasspath
+foo.finalMainClassOpt
+foo.prependShellScript
 foo.assembly
 
 */

scalalib/src/mill/scalalib/Assembly.scala

@@ -16,8 +16,12 @@ import scala.jdk.CollectionConverters._
 import scala.tools.nsc.io.Streamable
 import scala.util.Using
 
+case class Assembly(pathRef: PathRef, addedEntries: Int)
+
 object Assembly {
 
+ implicit val assemblyJsonRW: upickle.default.ReadWriter[Assembly] = upickle.default.macroRW
+
  val defaultRules: Seq[Rule] = Seq(
  Rule.Append("reference.conf", separator = "\n"),
  Rule.Exclude(JarFile.MANIFEST_NAME),
@@ -195,13 +199,38 @@ object Assembly {
  base: Option[os.Path] = None,
  assemblyRules: Seq[Assembly.Rule] = Assembly.defaultRules
  )(implicit ctx: Ctx.Dest with Ctx.Log): PathRef = {
- val tmp = ctx.dest / "out-tmp.jar"
+ create(
+ destJar = ctx.dest / "out.jar",
+ inputPaths = inputPaths,
+ manifest = manifest,
+ prependShellScript = Option(prependShellScript).filter(_ != ""),
+ base = base,
+ assemblyRules = assemblyRules
+ ).pathRef
+ }
+
+ def create(
+ destJar: os.Path,
+ inputPaths: Agg[os.Path],
+ manifest: mill.api.JarManifest = mill.api.JarManifest.MillDefault,
+ prependShellScript: Option[String] = None,
+ base: Option[os.Path] = None,
+ assemblyRules: Seq[Assembly.Rule] = Assembly.defaultRules
+ ): Assembly = {
+ val rawJar = os.temp("out-tmp", deleteOnExit = false)
+ // we create the file later
+ os.remove(rawJar)
 
- val baseUri = "jar:" + tmp.toIO.getCanonicalFile.toURI.toASCIIString
+ // use the `base` (the upstream assembly) as a start
+ val baseUri = "jar:" + rawJar.toIO.getCanonicalFile.toURI.toASCIIString
  val hm = base.fold(Map("create" -> "true")) { b =>
- os.copy(b, tmp)
+ os.copy(b, rawJar)
  Map.empty
  }
+
+ var addedEntryCount = 0
+
+ // Add more files by copying files to a JAR file system
  Using.resource(FileSystems.newFileSystem(URI.create(baseUri), hm.asJava)) { zipFs =>
  val manifestPath = zipFs.getPath(JarFile.MANIFEST_NAME)
  Files.createDirectories(manifestPath.getParent)
@@ -225,37 +254,47 @@ object Assembly {
  Seq(new ByteArrayInputStream(entry.separator.getBytes), inputStream())
  )
  val cleaned = if (Files.exists(path)) separated else separated.drop(1)
- val concatenated =
-  new SequenceInputStream(Collections.enumeration(cleaned.asJava))
+ val concatenated = new SequenceInputStream(Collections.enumeration(cleaned.asJava))
+ addedEntryCount += 1
  writeEntry(path, concatenated, append = true)
- case entry: WriteOnceEntry => writeEntry(path, entry.inputStream(), append = false)
+ case entry: WriteOnceEntry =>
+ addedEntryCount += 1
+ writeEntry(path, entry.inputStream(), append = false)
  }
  }
  } finally {
  resourceCleaner()
  }
  }
 
- val output = ctx.dest / "out.jar"
  // Prepend shell script and make it executable
- if (prependShellScript.isEmpty) os.move(tmp, output)
- else {
- val lineSep = if (!prependShellScript.endsWith("\n")) "\n\r\n" else ""
- os.write(output, prependShellScript + lineSep)
- os.write.append(output, os.read.inputStream(tmp))
-
- if (!scala.util.Properties.isWin) {
- os.perms.set(
- output,
- os.perms(output)
- + PosixFilePermission.GROUP_EXECUTE
- + PosixFilePermission.OWNER_EXECUTE
- + PosixFilePermission.OTHERS_EXECUTE
- )
- }
+ prependShellScript match {
+ case None =>
+ os.move(rawJar, destJar)
+ case Some(prependShellScript) =>
+ val lineSep = if (!prependShellScript.endsWith("\n")) "\n\r\n" else ""
+ val prepend = prependShellScript + lineSep
+ // Write the prepend-part into the final jar file
+ // https://en.wikipedia.org/wiki/Zip_(file_format)#Combination_with_other_file_formats
+ os.write(destJar, prepend)
+ // Append the actual JAR content
+ Using.resource(os.read.inputStream(rawJar)) { is =>
+ os.write.append(destJar, is)
+ }
+ os.remove(rawJar)
+
+ if (!scala.util.Properties.isWin) {
+ os.perms.set(
+ destJar,
+ os.perms(destJar)
+ + PosixFilePermission.GROUP_EXECUTE
+ + PosixFilePermission.OWNER_EXECUTE
+ + PosixFilePermission.OTHERS_EXECUTE
+ )
+ }
  }
 
- PathRef(output)
+ Assembly(PathRef(destJar), addedEntryCount)
  }
 
  private def writeEntry(p: java.nio.file.Path, inputStream: InputStream, append: Boolean): Unit = {

scalalib/src/mill/scalalib/JavaModule.scala

@@ -18,6 +18,8 @@ import mill.scalalib.publish.Artifact
 import mill.util.Jvm
 import os.{Path, ProcessOutput}
 
+import scala.annotation.nowarn
+
 /**
  * Core configuration required to compile a single Java compilation target
  */
@@ -569,11 +571,31 @@ trait JavaModule
  *
  * This should allow much faster assembly creation in the common case where
  * upstream dependencies do not change
+ *
+ * This implementation is deprecated because of it's return value.
+ * Please use [[upstreamAssembly2]] instead.
  */
+ @deprecated("Use upstreamAssembly2 instead, which has a richer return value", "Mill 0.11.8")
  def upstreamAssembly: T[PathRef] = T {
- Assembly.createAssembly(
- upstreamAssemblyClasspath().map(_.path),
- manifest(),
+ T.log.error(
+ s"upstreamAssembly target is deprecated and should no longer used." +
+ s" Please make sure to use upstreamAssembly2 instead."
+ )
+ upstreamAssembly2().pathRef
+ }
+
+ /**
+ * Build the assembly for upstream dependencies separate from the current
+ * classpath
+ *
+ * This should allow much faster assembly creation in the common case where
+ * upstream dependencies do not change
+ */
+ def upstreamAssembly2: T[Assembly] = T {
+ Assembly.create(
+ destJar = T.dest / "out.jar",
+ inputPaths = upstreamAssemblyClasspath().map(_.path),
+ manifest = manifest(),
  assemblyRules = assemblyRules
  )
  }
@@ -583,13 +605,45 @@ trait JavaModule
  * classfiles from this module and all it's upstream modules and dependencies
  */
  def assembly: T[PathRef] = T {
- Assembly.createAssembly(
+ // detect potential inconsistencies due to `upstreamAssembly` deprecation after 0.11.7
+ if (
+ (upstreamAssembly.ctx.enclosing: @nowarn) != s"${classOf[JavaModule].getName}#upstreamAssembly"
+ ) {
+ T.log.error(
+ s"${upstreamAssembly.ctx.enclosing: @nowarn} is overriding a deprecated target which is no longer used." +
+ s" Please make sure to override upstreamAssembly2 instead."
+ )
+ }
+
+ val prependScript = Option(prependShellScript()).filter(_ != "")
+ val upstream = upstreamAssembly2()
+
+ val created = Assembly.create(
+ destJar = T.dest / "out.jar",
  Agg.from(localClasspath().map(_.path)),
  manifest(),
- prependShellScript(),
- Some(upstreamAssembly().path),
+ prependScript,
+ Some(upstream.pathRef.path),
  assemblyRules
  )
+ // See https:/com-lihaoyi/mill/pull/2655#issuecomment-1672468284
+ val problematicEntryCount = 65535
+ if (
+ prependScript.isDefined &&
+ (upstream.addedEntries + created.addedEntries) > problematicEntryCount
+ ) {
+ Result.Failure(
+ s"""The created assembly jar contains more than ${problematicEntryCount} ZIP entries.
+ |JARs of that size are known to not work correctly with a prepended shell script.
+ |Either reduce the entries count of the assembly or disable the prepended shell script with:
+ |
+ | def prependShellScript = ""
+ |""".stripMargin,
+ Some(created.pathRef)
+ )
+ } else {
+ Result.Success(created.pathRef)
+ }
  }
 
  /**

scalalib/test/src/mill/scalalib/AssemblyTests.scala

@@ -0,0 +1,147 @@
+package mill.scalalib
+
+import mill._
+import mill.api.Result
+import mill.eval.Evaluator
+import mill.util.{Jvm, TestEvaluator, TestUtil}
+import utest._
+import utest.framework.TestPath
+
+import java.io.PrintStream
+
+// Ensure the assembly is runnable, even if we have assembled lots of dependencies into it
+// Reproduction of issues:
+// - https:/com-lihaoyi/mill/issues/528
+// - https:/com-lihaoyi/mill/issues/2650
+
+object AssemblyTests extends TestSuite {
+
+ object TestCase extends TestUtil.BaseModule {
+ trait Setup extends ScalaModule {
+ def scalaVersion = "2.13.11"
+ def sources = T.sources(T.workspace / "src")
+ def ivyDeps = super.ivyDeps() ++ Agg(
+ ivy"com.lihaoyi::scalatags:0.8.2",
+ ivy"com.lihaoyi::mainargs:0.4.0",
+ ivy"org.apache.avro:avro:1.11.1"
+ )
+ }
+ trait ExtraDeps extends ScalaModule {
+ def ivyDeps = super.ivyDeps() ++ Agg(
+ ivy"dev.zio::zio:2.0.15",
+ ivy"org.typelevel::cats-core:2.9.0",
+ ivy"org.apache.spark::spark-core:3.4.0",
+ ivy"dev.zio::zio-metrics-connectors:2.0.8",
+ ivy"dev.zio::zio-http:3.0.0-RC2"
+ )
+ }
+
+ object noExe extends Module {
+ object small extends Setup {
+ override def prependShellScript: T[String] = ""
+ }
+ object large extends Setup with ExtraDeps {
+ override def prependShellScript: T[String] = ""
+ }
+ }
+
+ object exe extends Module {
+ object small extends Setup
+ object large extends Setup with ExtraDeps
+ }
+
+ }
+
+ val sources = Map(
+ os.rel / "src" / "Main.scala" ->
+ """package ultra
+ |
+ |import scalatags.Text.all._
+ |import mainargs.{main, ParserForMethods}
+ |
+ |object Main {
+ | def generateHtml(text: String) = {
+ | h1(text).toString
+ | }
+ |
+ | @main
+ | def main(text: String) = {
+ | println(generateHtml(text))
+ | }
+ |
+ | def main(args: Array[String]): Unit = ParserForMethods(this).runOrExit(args)
+ |}""".stripMargin
+ )
+
+ def workspaceTest[T](
+ m: TestUtil.BaseModule,
+ env: Map[String, String] = Evaluator.defaultEnv,
+ debug: Boolean = false,
+ errStream: PrintStream = System.err
+ )(t: TestEvaluator => T)(implicit tp: TestPath): T = {
+ val eval = new TestEvaluator(m, env = env, debugEnabled = debug, errStream = errStream)
+ os.remove.all(m.millSourcePath)
+ sources.foreach { case (file, content) =>
+ os.write(m.millSourcePath / file, content, createFolders = true)
+ }
+ os.remove.all(eval.outPath)
+ os.makeDir.all(m.millSourcePath / os.up)
+ t(eval)
+ }
+
+ def runAssembly(file: os.Path, wd: os.Path, checkExe: Boolean = false): Unit = {
+ println(s"File size: ${os.stat(file).size}")
+ Jvm.runSubprocess(
+ commandArgs = Seq(Jvm.javaExe, "-jar", file.toString(), "--text", "tutu"),
+ envArgs = Map.empty[String, String],
+ workingDir = wd
+ )
+ if (checkExe) {
+ Jvm.runSubprocess(
+ commandArgs = Seq(file.toString(), "--text", "tutu"),
+ envArgs = Map.empty[String, String],
+ workingDir = wd
+ )
+ }
+ }
+
+ def tests: Tests = Tests {
+ test("Assembly") {
+ test("noExe") {
+ test("small") {
+ workspaceTest(TestCase) { eval =>
+ val Right((res, _)) = eval(TestCase.noExe.small.assembly)
+ runAssembly(res.path, TestCase.millSourcePath)
+ }
+ }
+ test("large") {
+ workspaceTest(TestCase) { eval =>
+ val Right((res, _)) = eval(TestCase.noExe.large.assembly)
+ runAssembly(res.path, TestCase.millSourcePath)
+ }
+ }
+ }
+ test("exe") {
+ test("small") {
+ workspaceTest(TestCase) { eval =>
+ val Right((res, _)) = eval(TestCase.exe.small.assembly)
+ runAssembly(res.path, TestCase.millSourcePath, checkExe = true)
+ }
+ }
+ test("large-should-fail") {
+ workspaceTest(TestCase) { eval =>
+ val Left(Result.Failure(msg, Some(res))) = eval(TestCase.exe.large.assembly)
+ val expectedMsg =
+ """The created assembly jar contains more than 65535 ZIP entries.
+ |JARs of that size are known to not work correctly with a prepended shell script.
+ |Either reduce the entries count of the assembly or disable the prepended shell script with:
+ |
+ | def prependShellScript = ""
+ |""".stripMargin
+ assert(msg == expectedMsg)
+ }
+ }
+ }
+ }
+ }
+}