retry on failure for flaky gadgets/tests

src/main/java/ysoserial/payloads/Wicket1.java

@@ -46,7 +46,7 @@ import ysoserial.payloads.util.Reflections;
  * $ cat /tmp/upload_3805815b_2d50_4e00_9dae_a854d5a0e614_479431761.tmp
  * blue lobster
  */
-@PayloadTest(harness = "ysoserial.payloads.FileUploadTest")
+@PayloadTest(harness = "ysoserial.payloads.FileUploadTest", flaky="possible race condition")
 @Dependencies({"org.apache.wicket:wicket-util:6.23.0", "org.slf4j:slf4j-api:1.6.4"})
 @Authors({ Authors.JACOBAINES })
 public class Wicket1 implements ReleaseableObjectPayload<DiskFileItem> {

src/main/java/ysoserial/payloads/annotation/PayloadTest.java

@@ -9,11 +9,11 @@ import java.lang.annotation.RetentionPolicy;
  */
 @Retention(RetentionPolicy.RUNTIME)
 public @interface PayloadTest {
-
     String skip() default "";
 
     String precondition() default "";
 
     String harness() default "";
 
+    String flaky() default "";
 }

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

@@ -18,13 +18,8 @@ import org.junit.runner.RunWith;
 import org.junit.runners.Parameterized;
 import org.junit.runners.Parameterized.Parameters;
 
-import ysoserial.CustomDeserializer;
-import ysoserial.CustomPayloadArgs;
-import ysoserial.CustomTest;
-import ysoserial.Deserializer;
-import ysoserial.Serializer;
+import ysoserial.*;
 import ysoserial.util.Throwables;
-import ysoserial.WrappedTest;
 import ysoserial.payloads.TestHarnessTest.ExecMockPayload;
 import ysoserial.payloads.TestHarnessTest.NoopMockPayload;
 import ysoserial.payloads.annotation.Dependencies;
@@ -74,6 +69,7 @@ public class PayloadsTest {
 
         PayloadTest t = payloadClass.getAnnotation(PayloadTest.class);
 
+        int tries = 1;
         if ( t != null ) {
             if ( !t.skip().isEmpty() ) {
                 Assume.assumeTrue(t.skip(), false);
@@ -82,6 +78,10 @@ public class PayloadsTest {
             if ( !t.precondition().isEmpty() ) {
                 Assume.assumeTrue("Precondition: " + t.precondition(), checkPrecondition(payloadClass, t.precondition()));
             }
+
+            if (! t.flaky().isEmpty()) {
+                tries = 5;
+            }
         }
 
         String[] deps = buildDeps(payloadClass);
@@ -114,11 +114,22 @@ public class PayloadsTest {
             callable = ( (WrappedTest) testHarness ).createCallable(callable);
         }
 
-        if ( testHarness instanceof CustomTest ) {
-            ( (CustomTest) testHarness ).run(callable);
-            return;
+        if (testHarness instanceof CustomTest) {
+            // if marked as flaky try up to 5 times
+            Exception ex = new Exception();
+            for (int i = 0; i < tries; i++) {
+                try {
+                    ((CustomTest) testHarness).run(callable);
+                    ex = null;
+                    break;
+                } catch (Exception e) {
+                    ex = e;
                 }
             }
+            if (ex != null) throw ex;
+        }
+
+    }