payload test refactors

src/test/java/ysoserial/test/payloads/PayloadsTest.java

@@ -1,29 +1,22 @@
 package ysoserial.test.payloads;
 
 
-import java.io.ByteArrayOutputStream;
 import java.io.File;
-import java.io.OutputStream;
+import java.io.IOException;
-import java.io.PrintStream;
 import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
 import java.net.URL;
 import java.net.URLClassLoader;
 import java.util.Arrays;
-import java.util.HashMap;
-import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.Callable;
 
 import org.jboss.shrinkwrap.resolver.api.maven.Maven;
 import org.junit.Assume;
 import org.junit.Test;
-import org.junit.runner.Description;
 import org.junit.runner.JUnitCore;
 import org.junit.runner.Result;
 import org.junit.runner.RunWith;
-import org.junit.runner.notification.Failure;
-import org.junit.runner.notification.RunListener;
 import org.junit.runners.Parameterized;
 import org.junit.runners.Parameterized.Parameters;
 
@@ -39,6 +32,9 @@ import ysoserial.test.payloads.TestHarnessTest.NoopMockPayload;
 import ysoserial.payloads.annotation.Dependencies;
 import ysoserial.payloads.annotation.PayloadTest;
 import ysoserial.payloads.util.ClassFiles;
+import ysoserial.test.util.Logging;
+import ysoserial.test.util.PayloadListener;
+import ysoserial.test.util.StdIoRedirection;
 
 
 /*
@@ -252,7 +248,9 @@ public class PayloadsTest {
         }
     }
 
-    public static void main(String[] args) {
+    public static void main(String[] args) throws IOException {
+        StdIoRedirection.init();
+        Logging.init();
 
         JUnitCore junit = new JUnitCore();
         PayloadListener listener = new PayloadListener();
@@ -260,83 +258,4 @@ public class PayloadsTest {
         Result result = junit.run(PayloadsTest.class);
         System.exit(result.wasSuccessful() ? 0 : 1);
     }
-
-    public static class StdIo {
-
-        private static final PrintStream realOut = System.out;
-        private static final PrintStream realErr = System.err;
-
-        public static void restoreStreams() {
-            setStreams(realOut, realErr);
-        }
-
-        public static void setStreams(PrintStream out, PrintStream err) {
-            System.setOut(out);
-            System.setErr(err);
-        }
-
-        public static void setStreams(OutputStream out, OutputStream err) {
-            setStreams(new PrintStream(out), new PrintStream(err));
-        }
-    }
-
-    public static class PayloadListener extends RunListener {
-        public enum Status {
-            SUCCESS,
-            FAILURE,
-            IGNORE,
-            ASSUMPTION_FAILURE
-        }
-
-        private Map<Description, ByteArrayOutputStream> outs = new HashMap<Description, ByteArrayOutputStream>();
-        private Map<Description, ByteArrayOutputStream> errs = new HashMap<Description, ByteArrayOutputStream>();
-
-        private Map<Description, Status> statuses = new HashMap<Description, Status>();
-
-        private Map<Description, Failure> failures = new HashMap<Description, Failure>();
-
-        @Override
-        public void testStarted(Description description) throws Exception {
-            System.out.println(getPayload(description.getDisplayName()) + ": STARTED");
-
-            statuses.put(description, Status.SUCCESS);
-
-            ByteArrayOutputStream out = new ByteArrayOutputStream();
-//            ByteArrayOutputStream err = new ByteArrayOutputStream();
-
-            outs.put(description, out);
-//            errs.put(description, err);
-
-            StdIo.setStreams(out, out);
-        }
-
-        @Override
-        public void testFinished(Description description) throws Exception {
-            outs.get(description).close();
-            //errs.get(description).close();
-
-            StdIo.restoreStreams();
-
-            Status status = statuses.get(description);
-            System.out.println(getPayload(description.getDisplayName()) + ": " + status);
-            if (status == Status.FAILURE) System.err.println(outs.get(description).toString());
-        }
-
-        @Override
-        public void testFailure(Failure failure) throws Exception {
-            statuses.put(failure.getDescription(), Status.FAILURE);
-            failures.put(failure.getDescription(), failure);
-        }
-
-        @Override
-        public void testAssumptionFailure(Failure failure) {
-            statuses.put(failure.getDescription(), Status.ASSUMPTION_FAILURE);
-            failures.put(failure.getDescription(), failure);
-        }
-
-        // testPayload[payloadClass: class ysoserial.payloads.JavassistWeld1](ysoserial.test.payloads.PayloadsTest)
-        public static String getPayload(String displayName) {
-            return displayName.replaceAll(".*\\[\\S+: class (\\w+\\.)+(\\w+)\\].*", "$2");
-        }
-    }
 }

src/test/java/ysoserial/test/util/Logging.java

@@ -0,0 +1,12 @@
+package ysoserial.test.util;
+
+import ysoserial.test.payloads.PayloadsTest;
+
+import java.io.IOException;
+import java.util.logging.LogManager;
+
+public class Logging {
+    public static void init() throws IOException {
+        LogManager.getLogManager().readConfiguration(PayloadsTest.class.getResourceAsStream("/logging.properties"));
+    }
+}

src/test/java/ysoserial/test/util/PayloadListener.java

@@ -0,0 +1,69 @@
+package ysoserial.test.util;
+
+import org.junit.runner.Description;
+import org.junit.runner.notification.Failure;
+import org.junit.runner.notification.RunListener;
+
+import java.io.ByteArrayOutputStream;
+import java.util.HashMap;
+import java.util.Map;
+
+public class PayloadListener extends RunListener {
+    public enum Status {
+        SUCCESS,
+        FAILURE,
+        IGNORE,
+        ASSUMPTION_FAILURE
+    }
+
+    private Map<Description, ByteArrayOutputStream> outs = new HashMap<Description, ByteArrayOutputStream>();
+    private Map<Description, ByteArrayOutputStream> errs = new HashMap<Description, ByteArrayOutputStream>();
+
+    private Map<Description, Status> statuses = new HashMap<Description, Status>();
+
+    private Map<Description, Failure> failures = new HashMap<Description, Failure>();
+
+    @Override
+    public void testStarted(Description description) throws Exception {
+        System.out.println(getPayload(description.getDisplayName()) + ": STARTED");
+
+        statuses.put(description, Status.SUCCESS);
+
+        ByteArrayOutputStream out = new ByteArrayOutputStream();
+//            ByteArrayOutputStream err = new ByteArrayOutputStream();
+
+        outs.put(description, out);
+//            errs.put(description, err);
+
+        StdIoRedirection.setStreams(out, out);
+    }
+
+    @Override
+    public void testFinished(Description description) throws Exception {
+        outs.get(description).close();
+        //errs.get(description).close();
+
+        StdIoRedirection.restoreStreams();
+
+        Status status = statuses.get(description);
+        System.out.println(getPayload(description.getDisplayName()) + ": " + status);
+        if (status == Status.FAILURE) System.err.println(outs.get(description).toString());
+    }
+
+    @Override
+    public void testFailure(Failure failure) throws Exception {
+        statuses.put(failure.getDescription(), Status.FAILURE);
+        failures.put(failure.getDescription(), failure);
+    }
+
+    @Override
+    public void testAssumptionFailure(Failure failure) {
+        statuses.put(failure.getDescription(), Status.ASSUMPTION_FAILURE);
+        failures.put(failure.getDescription(), failure);
+    }
+
+    // testPayload[payloadClass: class ysoserial.payloads.JavassistWeld1](ysoserial.test.payloads.PayloadsTest)
+    public static String getPayload(String displayName) {
+        return displayName.replaceAll(".*\\[\\S+: class (\\w+\\.)+(\\w+)\\].*", "$2");
+    }
+}

src/test/java/ysoserial/test/util/StdIoRedirection.java

@@ -0,0 +1,66 @@
+package ysoserial.test.util;
+
+import javassist.util.proxy.MethodHandler;
+import javassist.util.proxy.ProxyFactory;
+
+import java.io.OutputStream;
+import java.io.PrintStream;
+import java.lang.reflect.Method;
+
+/*
+  Replace System.out/err early-on with proxies that delegate to streams controlled here to ensure changes reflected
+  even when references are saved by various writing/logging classes
+ */
+public class StdIoRedirection {
+
+    private static final PrintStream realOut = System.out;
+    private static final PrintStream realErr = System.err;
+
+    private static PrintStream delegateOut = System.out;
+    private static PrintStream delegateErr = System.err;
+
+    private static PrintStream proxyOut;
+    private static PrintStream proxyErr;
+
+    static {
+        try {
+            proxyOut = (PrintStream) new ProxyFactory() {{
+                this.setSuperclass(PrintStream.class);
+            }}.create(new Class[]{OutputStream.class}, new Object[]{delegateOut}, new MethodHandler() {
+                @Override
+                public Object invoke(Object self, Method thisMethod, Method proceed, Object[] args) throws Throwable {
+                    return thisMethod.invoke(delegateOut, args);
+                }
+            });
+
+            proxyErr = (PrintStream) new ProxyFactory() {{
+                this.setSuperclass(PrintStream.class);
+            }}.create(new Class[]{OutputStream.class}, new Object[]{delegateErr}, new MethodHandler() {
+                @Override
+                public Object invoke(Object self, Method thisMethod, Method proceed, Object[] args) throws Throwable {
+                    return thisMethod.invoke(delegateErr, args);
+                }
+            });
+        } catch (Exception e) {
+            e.printStackTrace();
+        }
+    }
+
+    public static void init() {
+        System.setOut(proxyOut);
+        System.setErr(proxyErr);
+    }
+
+    public static void restoreStreams() {
+        setStreams(realOut, realErr);
+    }
+
+    public static void setStreams(PrintStream out, PrintStream err) {
+        delegateOut = out;
+        delegateErr = err;
+    }
+
+    public static void setStreams(OutputStream out, OutputStream err) {
+        setStreams(new PrintStream(out), new PrintStream(err));
+    }
+}