From 262dfa7773c3e54402fe0d30b33126798a1a7b33 Mon Sep 17 00:00:00 2001 From: Andrew Rosenberg <2112318+Iapetus999@users.noreply.github.com> Date: Thu, 18 Apr 2024 08:50:29 -0700 Subject: [PATCH] Add step limit and expiration redux (#489) * Add step limit and expiration redux * Rollback WS changes * Rollback WS changes --- .../security/larky/parser/LarkyEvaluator.java | 11 ++++ .../java/net/starlark/java/eval/Eval.java | 7 +++ .../starlark/java/eval/StarlarkThread.java | 53 +++++++++++++++++++ .../starlark/java/eval/EvaluationTest.java | 27 ++++++++++ 4 files changed, 98 insertions(+) diff --git a/larky/src/main/java/com/verygood/security/larky/parser/LarkyEvaluator.java b/larky/src/main/java/com/verygood/security/larky/parser/LarkyEvaluator.java index e34921ea4..31a76dd2b 100644 --- a/larky/src/main/java/com/verygood/security/larky/parser/LarkyEvaluator.java +++ b/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 pending = new LinkedHashSet<>(); private final Map 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) { diff --git a/libstarlark/src/main/java/net/starlark/java/eval/Eval.java b/libstarlark/src/main/java/net/starlark/java/eval/Eval.java index d3eca897e..2b3ab571b 100644 --- a/libstarlark/src/main/java/net/starlark/java/eval/Eval.java +++ b/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 diff --git a/libstarlark/src/main/java/net/starlark/java/eval/StarlarkThread.java b/libstarlark/src/main/java/net/starlark/java/eval/StarlarkThread.java index b6816fdca..ae5f4bdd9 100644 --- a/libstarlark/src/main/java/net/starlark/java/eval/StarlarkThread.java +++ b/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. */ diff --git a/libstarlark/src/test/java/net/starlark/java/eval/EvaluationTest.java b/libstarlark/src/test/java/net/starlark/java/eval/EvaluationTest.java index 7fbc0beb4..88b0b1058 100644 --- a/libstarlark/src/test/java/net/starlark/java/eval/EvaluationTest.java +++ b/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()