From 6909d5269dff17c3f6ffcd22f697fa08e3e00ac1 Mon Sep 17 00:00:00 2001
From: Chris Frohoff <cfrohoff@qualcomm.com>
Date: Fri, 16 Nov 2018 11:40:23 -0800
Subject: [PATCH] retry on failure for flaky gadgets/tests

---
 src/main/java/ysoserial/payloads/Wicket1.java |  2 +-
 .../payloads/annotation/PayloadTest.java      |  6 ++--
 .../java/ysoserial/payloads/PayloadsTest.java | 29 +++++++++++++------
 3 files changed, 24 insertions(+), 13 deletions(-)

diff --git a/src/main/java/ysoserial/payloads/Wicket1.java b/src/main/java/ysoserial/payloads/Wicket1.java
index abf617c..29e5d9b 100644
--- a/src/main/java/ysoserial/payloads/Wicket1.java
+++ b/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> {
diff --git a/src/main/java/ysoserial/payloads/annotation/PayloadTest.java b/src/main/java/ysoserial/payloads/annotation/PayloadTest.java
index f729988..b54a310 100644
--- a/src/main/java/ysoserial/payloads/annotation/PayloadTest.java
+++ b/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 "";
 }
diff --git a/src/test/java/ysoserial/payloads/PayloadsTest.java b/src/test/java/ysoserial/payloads/PayloadsTest.java
index 3713486..8613c5f 100644
--- a/src/test/java/ysoserial/payloads/PayloadsTest.java
+++ b/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,10 +114,21 @@ 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;
         }
+
     }