From dc74a31be121fc98e156a166b7a020d599ef41ed Mon Sep 17 00:00:00 2001
From: Daniel Zou <dzou@users.noreply.github.com>
Date: Wed, 16 Jun 2021 10:56:48 -0400
Subject: [PATCH] Modify netty-shaded resources to reference shaded class names
 (#8258)

---
 netty/shaded/build.gradle                     | 55 +++++++++++++++++++
 .../io/grpc/netty/shaded/ShadingTest.java     | 18 ++++++
 2 files changed, 73 insertions(+)

diff --git a/netty/shaded/build.gradle b/netty/shaded/build.gradle
index 6162706869e..af0c9648c6f 100644
--- a/netty/shaded/build.gradle
+++ b/netty/shaded/build.gradle
@@ -1,3 +1,19 @@
+import com.github.jengelman.gradle.plugins.shadow.transformers.Transformer
+import com.github.jengelman.gradle.plugins.shadow.transformers.TransformerContext
+import org.gradle.api.file.FileTreeElement
+import shadow.org.apache.tools.zip.ZipOutputStream
+import shadow.org.apache.tools.zip.ZipEntry
+
+
+buildscript {
+    repositories {
+        jcenter()
+    }
+    dependencies {
+        classpath "com.github.jengelman.gradle.plugins:shadow:6.1.0"
+    }
+}
+
 plugins {
     id "java"
     id "maven-publish"
@@ -41,6 +57,7 @@ shadowJar {
     // this includes concatenation of string literals and constants.
     relocate 'META-INF/native/libnetty', 'META-INF/native/libio_grpc_netty_shaded_netty'
     relocate 'META-INF/native/netty', 'META-INF/native/io_grpc_netty_shaded_netty'
+    transform(NettyResourceTransformer.class)
     mergeServiceFiles()
 }
 
@@ -73,3 +90,41 @@ compileTestShadowJava.options.compilerArgs = compileTestJava.options.compilerArg
 compileTestShadowJava.options.encoding = compileTestJava.options.encoding
 
 test.dependsOn testShadow
+
+/**
+ * A Transformer which updates the Netty JAR META-INF/ resources to accurately
+ * reference shaded class names.
+ */
+class NettyResourceTransformer implements Transformer {
+
+    // A map of resource file paths to be modified
+    private Map<String, String> resources = [:]
+
+    @Override
+    boolean canTransformResource(FileTreeElement fileTreeElement) {
+        fileTreeElement.name.startsWith("META-INF/native-image/io.netty")
+    }
+
+    @Override
+    void transform(TransformerContext context) {
+        String updatedContent = context.is.getText().replace("io.netty", "io.grpc.netty.shaded.io.netty")
+        resources.put(context.path, updatedContent)
+    }
+
+    @Override
+    boolean hasTransformedResource() {
+        resources.size() > 0
+    }
+
+    @Override
+    void modifyOutputStream(ZipOutputStream outputStream, boolean preserveFileTimestamps) {
+        for (resourceEntry in resources) {
+            ZipEntry entry = new ZipEntry(resourceEntry.key)
+            entry.time = TransformerContext.getEntryTimestamp(preserveFileTimestamps, entry.time)
+
+            outputStream.putNextEntry(entry)
+            outputStream.write(resourceEntry.value.getBytes())
+            outputStream.closeEntry()
+        }
+    }
+}
diff --git a/netty/shaded/src/testShadow/java/io/grpc/netty/shaded/ShadingTest.java b/netty/shaded/src/testShadow/java/io/grpc/netty/shaded/ShadingTest.java
index e4a84986976..d3bdc4394ca 100644
--- a/netty/shaded/src/testShadow/java/io/grpc/netty/shaded/ShadingTest.java
+++ b/netty/shaded/src/testShadow/java/io/grpc/netty/shaded/ShadingTest.java
@@ -39,6 +39,11 @@
 import io.grpc.testing.protobuf.SimpleServiceGrpc;
 import io.grpc.testing.protobuf.SimpleServiceGrpc.SimpleServiceBlockingStub;
 import io.grpc.testing.protobuf.SimpleServiceGrpc.SimpleServiceImplBase;
+
+import java.io.IOException;
+import java.io.InputStream;
+import java.nio.charset.StandardCharsets;
+import java.util.Scanner;
 import java.util.concurrent.TimeUnit;
 import org.junit.After;
 import org.junit.Test;
@@ -69,6 +74,19 @@ public void noNormalNetty() throws Exception {
     Class.forName("io.grpc.netty.NettyServerBuilder");
   }
 
+  /** Verify that resources under META-INF/native-image reference shaded class names. */
+  @Test
+  public void nettyResourcesUpdated() throws IOException {
+    InputStream inputStream = NettyChannelBuilder.class.getClassLoader()
+        .getResourceAsStream("META-INF/native-image/io.netty/transport/reflection-config.json");
+    assertThat(inputStream).isNotNull();
+
+    Scanner s = new Scanner(inputStream, StandardCharsets.UTF_8.name()).useDelimiter("\\A");
+    String reflectionConfig = s.hasNext() ? s.next() : "";
+
+    assertThat(reflectionConfig).contains("io.grpc.netty.shaded.io.netty");
+  }
+
   @Test
   public void serviceLoaderFindsNetty() throws Exception {
     assertThat(Grpc.newServerBuilderForPort(0, InsecureServerCredentials.create()))