Add step limit and expiration redux (#489)

* Add step limit and expiration redux

* Rollback WS changes

* Rollback WS changes

larky/src/main/java/com/verygood/security/larky/parser/LarkyEvaluator.java

@@ -42,6 +42,8 @@ public final class LarkyEvaluator {
 
   private static final FluentLogger logger = FluentLogger.forEnclosingClass();
   public static final String EXECUTION_STEPS = "_EXECUTION_STEPS_";
+  public static final String STEP_LIMIT = "LARKY_EVAL_STEP_LIMIT";
+  public static final String EXPIRATION_MS = "LARKY_EVAL_EXPIRATION_MS";
 
   private final LinkedHashSet<String> pending = new LinkedHashSet<>();
   private final Map<String, Module> loaded = new HashMap<>();
@@ -136,6 +138,15 @@ public EvaluationResult eval(StarFile content)
       thread.setLoader(loadedModules::get);
       thread.setThreadLocal(Reporter.class, reporter);
       thread.setPrintHandler(reporter::report);
+
+      if (environment.containsKey(STEP_LIMIT) && environment.get(STEP_LIMIT) != null) {
+        thread.setMaxExecutionSteps(Long.parseLong(environment.get(STEP_LIMIT).toString()));
+      }
+
+      if (environment.containsKey(EXPIRATION_MS) && environment.get(EXPIRATION_MS) != null) {
+        thread.setExpirationMs(Long.parseLong(environment.get(EXPIRATION_MS).toString()));
+      }
+
       try {
         starlarkOutput = Starlark.execFileProgram(prog, module, thread);
       } catch (EvalException cause) {

libstarlark/src/main/java/net/starlark/java/eval/Eval.java

@@ -264,6 +264,10 @@ private static TokenKind exec(StarlarkThread.Frame fr, Statement st)
       throw new EvalException("Starlark computation cancelled: too many steps");
     }
 
+    if (fr.thread.isExpired()) {
+      throw new EvalException("Starlark computation cancelled: past expiration date");
+    }
+
     switch (st.kind()) {
       case ASSIGNMENT:
         execAssignment(fr, (AssignmentStatement) st);
@@ -478,6 +482,9 @@ private static Object eval(StarlarkThread.Frame fr, Expression expr)
     if (++fr.thread.steps >= fr.thread.stepLimit) {
       throw new EvalException("Starlark computation cancelled: too many steps");
     }
+    if (fr.thread.isExpired()) {
+      throw new EvalException("Starlark computation cancelled: past expiration date");
+    }
 
     // The switch cases have been split into separate functions
     // to reduce the stack usage during recursion, which is

libstarlark/src/main/java/net/starlark/java/eval/StarlarkThread.java

@@ -14,13 +14,16 @@
 
 package net.starlark.java.eval;
 
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
+import java.time.Clock;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.concurrent.atomic.AtomicInteger;
+import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 import javax.annotation.concurrent.Immutable;
 import net.starlark.java.syntax.Location;
@@ -66,6 +69,8 @@ public final class StarlarkThread {
 
   long steps; // count of logical computation steps executed so far
   long stepLimit = Long.MAX_VALUE; // limit on logical computation steps
+  long expirationMs = Long.MAX_VALUE; // time after which execution should halt
+  Clock clock = Clock.systemUTC(); // Used to check expiration. Can be set for debug purposes
 
   /**
    * Returns the number of Starlark computation steps executed by this thread according to a
@@ -86,6 +91,54 @@ public void setMaxExecutionSteps(long steps) {
     this.stepLimit = steps;
   }
 
+  /**
+   *
+   * @return step limit
+   */
+  public long getStepLimit() {
+    return stepLimit;
+  }
+
+  /**
+   *
+   * @return expiration date in ms since 1970
+   */
+  public long getExpirationMs() {
+    return expirationMs;
+  }
+
+  /**
+   * Sets the expiration date.
+   * Will halt execution on the next evaluation after the given date.
+   * If not called, evals will not expire.
+   *
+   * @param expirationMs expiration date in ms since 1970
+   */
+  public void setExpirationMs(long expirationMs) {
+    this.expirationMs = expirationMs;
+  }
+
+  /**
+   *
+   * @return current Clock object
+   */
+  public Clock getClock() {
+    return clock;
+  }
+
+  /**
+   * Sets a specific clock object for testing purposes
+   * @param clock Clock for testing
+   */
+  @VisibleForTesting
+  public void setClock(@Nonnull Clock clock) {
+    this.clock = clock;
+  }
+
+  public boolean isExpired(){
+    return clock.millis() > expirationMs;
+  }
+
   /**
    * Disables polling of the {@link java.lang.Thread#interrupted} flag during Starlark evaluation.
    */

libstarlark/src/test/java/net/starlark/java/eval/EvaluationTest.java

@@ -15,6 +15,7 @@
 
 import static com.google.common.truth.Truth.assertThat;
 import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.assertTrue;
 
 import com.google.common.collect.ImmutableMap;
 import java.util.ArrayList;
@@ -164,6 +165,32 @@ long run(int n) throws SyntaxError.Exception, EvalException, InterruptedExceptio
     assertThat(ex).hasMessageThat().contains("Starlark computation cancelled: too many steps");
   }
 
+  @Test
+  public void testExpiration() throws Exception {
+    Mutability mu = Mutability.create("test");
+    StarlarkThread thread = new StarlarkThread(mu, StarlarkSemantics.DEFAULT);
+    ParserInput input = ParserInput.fromLines("squares = [x+x for x in range(n)]");
+
+    class C {
+      long run(int n) throws SyntaxError.Exception, EvalException, InterruptedException {
+        Module module =
+            Module.withPredeclared(
+                StarlarkSemantics.DEFAULT, ImmutableMap.of("n", StarlarkInt.of(n)));
+        long steps0 = thread.getExecutedSteps();
+        Starlark.execFile(input, FileOptions.DEFAULT, module, thread);
+        return thread.getExecutedSteps() - steps0;
+      }
+    }
+
+    // Exceeding the limit causes cancellation.
+    thread.setExpirationMs(1);
+    EvalException ex = assertThrows(EvalException.class, () -> new C().run(1000));
+    assertThat(ex).hasMessageThat().contains("Starlark computation cancelled: past expiration date");
+    thread.setExpirationMs(System.currentTimeMillis() + 10000000);
+    long steps = new C().run(1000); // should not throw error
+    assertTrue(steps > 0);
+  }
+
   @Test
   public void testExprs() throws Exception {
     ev.new Scenario()