[incubator-servicecomb-java-chassis] branch master updated: [SCB-725] support get main class package when run with "java -jar xxx.jar"

[wu...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

wujimin pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/incubator-servicecomb-java-chassis.git


The following commit(s) were added to refs/heads/master by this push:
     new b00f4c7  [SCB-725] support get main class package when run with "java -jar xxx.jar"
b00f4c7 is described below

commit b00f4c748122460ba1484815ed9c98eadf9eba31
Author: wujimin <wu...@huawei.com>
AuthorDate: Tue Jul 10 12:51:43 2018 +0800

    [SCB-725] support get main class package when run with "java -jar xxx.jar"
---
 .../foundation/common/utils/JvmUtils.java          | 41 ++++++++++++-
 .../foundation/common/utils/TestJvmUtils.java      | 68 +++++++++++++++++++++-
 2 files changed, 105 insertions(+), 4 deletions(-)

diff --git a/foundations/foundation-common/src/main/java/org/apache/servicecomb/foundation/common/utils/JvmUtils.java b/foundations/foundation-common/src/main/java/org/apache/servicecomb/foundation/common/utils/JvmUtils.java
index 48a9baa..bb02d72 100644
--- a/foundations/foundation-common/src/main/java/org/apache/servicecomb/foundation/common/utils/JvmUtils.java
+++ b/foundations/foundation-common/src/main/java/org/apache/servicecomb/foundation/common/utils/JvmUtils.java
@@ -16,6 +16,12 @@
  */
 package org.apache.servicecomb.foundation.common.utils;
 
+import java.io.File;
+import java.io.InputStream;
+import java.net.URL;
+import java.util.jar.JarFile;
+import java.util.jar.Manifest;
+
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -36,21 +42,50 @@ public final class JvmUtils {
    * @return main class or null, never throw exception
    */
   public static Class<?> findMainClass() {
+    // 1.run with java -cp ......
+    //   command is main class and args
+    // 2.run with java -jar ......
+    //   command is jar file name and args
     String command = System.getProperty(SUN_JAVA_COMMAND);
     if (command == null || command.isEmpty()) {
       return null;
     }
 
-    // command is main class and args
-    String mainClass = command.trim().split(" ")[0];
+    String mainClassOrJar = command.trim().split(" ")[0];
+    String mainClass = readFromJar(mainClassOrJar);
+    if (mainClass == null || mainClass.isEmpty()) {
+      return null;
+    }
+
     try {
-      return Class.forName(mainClass);
+      Class<?> cls = Class.forName(mainClass);
+      LOGGER.info("Found main class \"{}\".", mainClass);
+      return cls;
     } catch (Throwable e) {
       LOGGER.warn("\"{}\" is not a valid class.", mainClass, e);
       return null;
     }
   }
 
+  private static String readFromJar(String mainClassOrJar) {
+    if (!mainClassOrJar.endsWith(".jar")) {
+      return mainClassOrJar;
+    }
+
+    String manifestUri = "jar:file:/" + new File(mainClassOrJar).getAbsolutePath() + "!/" + JarFile.MANIFEST_NAME;
+
+    try {
+      URL url = new URL(manifestUri);
+      try (InputStream inputStream = url.openStream()) {
+        Manifest manifest = new Manifest(inputStream);
+        return manifest.getMainAttributes().getValue("Main-Class");
+      }
+    } catch (Throwable e) {
+      LOGGER.warn("Failed to read Main-Class from \"{}\".", manifestUri, e);
+      return null;
+    }
+  }
+
   /**
    * find a property class loader to avoid null
    */
diff --git a/foundations/foundation-common/src/test/java/org/apache/servicecomb/foundation/common/utils/TestJvmUtils.java b/foundations/foundation-common/src/test/java/org/apache/servicecomb/foundation/common/utils/TestJvmUtils.java
index d16792f..257c8f4 100644
--- a/foundations/foundation-common/src/test/java/org/apache/servicecomb/foundation/common/utils/TestJvmUtils.java
+++ b/foundations/foundation-common/src/test/java/org/apache/servicecomb/foundation/common/utils/TestJvmUtils.java
@@ -16,12 +16,24 @@
  */
 package org.apache.servicecomb.foundation.common.utils;
 
+import java.io.ByteArrayInputStream;
+import java.io.File;
+import java.io.InputStream;
+import java.net.URL;
+import java.util.jar.JarFile;
+
 import org.apache.servicecomb.foundation.test.scaffolding.log.LogCollector;
 import org.junit.AfterClass;
 import org.junit.Assert;
 import org.junit.Before;
 import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.powermock.api.mockito.PowerMockito;
+import org.powermock.core.classloader.annotations.PrepareForTest;
+import org.powermock.modules.junit4.PowerMockRunner;
 
+@RunWith(PowerMockRunner.class)
+@PrepareForTest({JvmUtils.class})
 public class TestJvmUtils {
   static String orgCmd = System.getProperty(JvmUtils.SUN_JAVA_COMMAND);
 
@@ -63,9 +75,63 @@ public class TestJvmUtils {
   }
 
   @Test
-  public void findMainClass_normal() {
+  public void findMainClass_class_normal() {
     System.setProperty(JvmUtils.SUN_JAVA_COMMAND, TestJvmUtils.class.getName() + " arg");
 
     Assert.assertEquals(TestJvmUtils.class, JvmUtils.findMainClass());
   }
+
+  @Test
+  public void findMainClass_jar_normal() throws Exception {
+    String content = String.format("Manifest-Version: 1.0\nMain-Class: %s\n", TestJvmUtils.class.getName());
+    InputStream inputStream = new ByteArrayInputStream(content.getBytes());
+
+    URL url = PowerMockito.mock(URL.class);
+
+    String command = "a.jar";
+    String manifestUri = "jar:file:/" + new File(command).getAbsolutePath() + "!/" + JarFile.MANIFEST_NAME;
+
+    PowerMockito.whenNew(URL.class).withParameterTypes(String.class)
+        .withArguments(manifestUri).thenReturn(url);
+    PowerMockito.when(url.openStream()).thenReturn(inputStream);
+
+    System.setProperty(JvmUtils.SUN_JAVA_COMMAND, command + " arg");
+
+    Assert.assertEquals(TestJvmUtils.class, JvmUtils.findMainClass());
+  }
+
+  @Test
+  public void findMainClass_jar_null() throws Exception {
+    String content = "Manifest-Version: 1.0\n";
+    InputStream inputStream = new ByteArrayInputStream(content.getBytes());
+
+    URL url = PowerMockito.mock(URL.class);
+
+    String command = "a.jar";
+    String manifestUri = "jar:file:/" + new File(command).getAbsolutePath() + "!/" + JarFile.MANIFEST_NAME;
+
+    PowerMockito.whenNew(URL.class).withParameterTypes(String.class)
+        .withArguments(manifestUri).thenReturn(url);
+    PowerMockito.when(url.openStream()).thenReturn(inputStream);
+
+    System.setProperty(JvmUtils.SUN_JAVA_COMMAND, command + " arg");
+
+    Assert.assertNull(JvmUtils.findMainClass());
+  }
+
+  @Test
+  public void findMainClass_jar_readFailed() throws Exception {
+    URL url = PowerMockito.mock(URL.class);
+
+    String command = "a.jar";
+    String manifestUri = "jar:file:/" + new File(command).getAbsolutePath() + "!/" + JarFile.MANIFEST_NAME;
+
+    PowerMockito.whenNew(URL.class).withParameterTypes(String.class)
+        .withArguments(manifestUri).thenReturn(url);
+    PowerMockito.when(url.openStream()).thenThrow(new Error());
+
+    System.setProperty(JvmUtils.SUN_JAVA_COMMAND, command + " arg");
+
+    Assert.assertNull(JvmUtils.findMainClass());
+  }
 }