From 6a422ddac0b5ca908ac6e3586bfa230b9a0cd19e Mon Sep 17 00:00:00 2001
From: Joan Goyeau <jgoyeau@netflix.com>
Date: Wed, 19 Aug 2020 13:57:55 -0700
Subject: [PATCH] Ability to create InputStream shader function

---
 .../com/eed3si9n/jarjarabrams/Shader.scala    | 38 ++++++++++++-------
 1 file changed, 24 insertions(+), 14 deletions(-)

diff --git a/core/src/main/scala/com/eed3si9n/jarjarabrams/Shader.scala b/core/src/main/scala/com/eed3si9n/jarjarabrams/Shader.scala
index 91334f6..2740391 100644
--- a/core/src/main/scala/com/eed3si9n/jarjarabrams/Shader.scala
+++ b/core/src/main/scala/com/eed3si9n/jarjarabrams/Shader.scala
@@ -28,6 +28,14 @@ object Shader {
       mappings: Iterable[(InputStream, String)],
       verbose: Boolean
   ): Iterable[(InputStream, String)] = {
+    val shader = inputStreamShader(rules, verbose)
+    mappings.flatMap(f => shader(f._1, f._2))
+  }
+
+  def inputStreamShader(
+    rules: Seq[ShadeRule],
+    verbose: Boolean
+  ): (InputStream, String) => Option[(InputStream, String)] = {
     val jjrules = rules.flatMap { r =>
       r.shadePattern match {
         case ShadePattern.Rename(patterns) =>
@@ -55,24 +63,26 @@ object Shader {
 
     val proc = new JJProcessor(jjrules, verbose, true, null)
 
-    /*
-    jarjar MisplacedClassProcessor class transforms byte[] to a class using org.objectweb.asm.ClassReader.getClassName
-    which always translates class names containing '.' into '/', regardless of OS platform.
-    We need to transform any windows file paths in order for jarjar to match them properly and not omit them.
-     */
-    val sanitizedMappings =
-      mappings.map(f => if (f._2.contains('\\')) (f._1, f._2.replace('\\', '/')) else f)
-    val shadedInputStreams = sanitizedMappings.flatMap { f =>
+    { (inputStream, mapping) =>
+      /*
+      jarjar MisplacedClassProcessor class transforms byte[] to a class using org.objectweb.asm.ClassReader.getClassName
+      which always translates class names containing '.' into '/', regardless of OS platform.
+      We need to transform any windows file paths in order for jarjar to match them properly and not omit them.
+       */
+      val sanitizedMapping = if (mapping.contains('\\')) mapping.replace('\\', '/') else mapping
       val entry = new EntryStruct
-      entry.data = readAllBytes(f._1)
-      entry.name = f._2
+      entry.data = readAllBytes(inputStream)
+      entry.name = sanitizedMapping
       entry.time = -1
       entry.skipTransform = false
-      if (proc.process(entry)) Some(new ByteArrayInputStream(entry.data) -> entry.name)
-      else None
+      val shadedInputStream =
+        if (proc.process(entry)) Some(new ByteArrayInputStream(entry.data) {
+          override def close(): Unit = inputStream.close()
+        } -> entry.name)
+        else None
+      val excludes = proc.getExcludes
+      shadedInputStream.filterNot(a => excludes.contains(a._2))
     }
-    val excludes = proc.getExcludes
-    shadedInputStreams.filterNot(mapping => excludes.contains(mapping._2))
   }
 }