verygoodsecurity · Iapetus999 · Jun 9, 2023
diff --git a/larky/src/main/java/com/verygood/security/larky/modules/types/structs/StructBinOp.java b/larky/src/main/java/com/verygood/security/larky/modules/types/structs/StructBinOp.java
@@ -27,17 +27,17 @@
 import com.verygood.security.larky.modules.types.LarkyIterator;
 import com.verygood.security.larky.modules.types.PyProtocols;
 import com.verygood.security.larky.parser.StarlarkUtil;
-
+import lombok.extern.slf4j.Slf4j;
 import net.starlark.java.eval.Dict;
 import net.starlark.java.eval.EvalException;
 import net.starlark.java.eval.Starlark;
 import net.starlark.java.eval.StarlarkCallable;
 import net.starlark.java.eval.StarlarkThread;
 import net.starlark.java.eval.Tuple;
 import net.starlark.java.syntax.TokenKind;
-
 import org.jetbrains.annotations.Nullable;
 
+@Slf4j
 public class StructBinOp {
 
   private StructBinOp() {
@@ -145,7 +145,7 @@ public static Object binaryOperation(
     String operator,
     StarlarkThread thread
   ) throws EvalException {
-
+log.info("binaryOperation");
     String lhsType = StarlarkUtil.richType(lhs);
     String rhsType = StarlarkUtil.richType(rhs);
 

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
@@ -5,18 +5,20 @@
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Maps;
 import com.google.common.flogger.FluentLogger;
+import com.verygood.security.larky.ModuleSupplier;
+import com.verygood.security.larky.annot.Library;
+import com.verygood.security.larky.console.Console;
+import com.verygood.security.larky.modules.utils.Reporter;
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
-
-import com.verygood.security.larky.ModuleSupplier;
-import com.verygood.security.larky.annot.Library;
-import com.verygood.security.larky.console.Console;
-import com.verygood.security.larky.modules.utils.Reporter;
-
+import lombok.Builder;
+import lombok.Data;
+import lombok.Getter;
+import lombok.val;
 import net.starlark.java.annot.StarlarkAnnotations;
 import net.starlark.java.annot.StarlarkBuiltin;
 import net.starlark.java.eval.EvalException;
@@ -31,15 +33,10 @@
 import net.starlark.java.syntax.Program;
 import net.starlark.java.syntax.StarlarkFile;
 import net.starlark.java.syntax.SyntaxError;
-
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.annotations.VisibleForTesting;
 
-import lombok.Builder;
-import lombok.Data;
-import lombok.Getter;
-
 /**
  * An utility class for traversing and evaluating the config file dependency graph.
  */
@@ -74,26 +71,29 @@ public LarkyEvaluator(LarkyScript larkyScript, Console console) {
   }
 
   /**
-   * The output of a Larky script is an evaluated {@link Module} that contains
-   * various attributes such as {@link Module#getGlobals()},
-   * {@link Module#getPredeclaredBindings()}, as well as other various items.
-   *
+   * The output of a Larky script is an evaluated {@link Module} that contains various attributes such as
+   * {@link Module#getGlobals()}, {@link Module#getPredeclaredBindings()}, as well as other various items.
+   * <p>
    * Sometimes, when evaluating a Larky script, there is some output that is generated.
-   *
-   * This interface encapsulates an interface that allows the caller to introspect
-   * the result of a Larky script evaluation.
+   * <p>
+   * This interface encapsulates an interface that allows the caller to introspect the result of a Larky script
+   * evaluation.
    */
   public interface EvaluationResult {
-      boolean hasOutput();
-      boolean hasModule();
 
-      Object getOutput();
-      Module getModule();
+    boolean hasOutput();
+
+    boolean hasModule();
+
+    Object getOutput();
+
+    Module getModule();
   }
 
   @Builder
   @Data
   protected static class DefaultEvaluationResult implements EvaluationResult {
+
     private Object output;
     private Module module;
 
@@ -114,9 +114,9 @@ public EvaluationResult eval(StarFile content)
     Module module = loaded.get(content.path());
     if (module != null) {
       return DefaultEvaluationResult.builder()
-        .output(null)
-        .module(module)
-        .build();
+          .output(null)
+          .module(module)
+          .build();
     }
     pending.add(content.path());
 
@@ -137,18 +137,23 @@ 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(Integer.parseInt(environment.get("STEP_LIMIT").toString()));
+      }
+
       try {
         starlarkOutput = Starlark.execFileProgram(prog, module, thread);
-      } catch(EvalException cause) {
+      } catch (EvalException cause) {
         throw new StarlarkEvalWrapper.Exc.RuntimeEvalException(cause, thread);
       }
     }
     pending.remove(content.path());
     loaded.put(content.path(), module);
     return DefaultEvaluationResult.builder()
-      .output(starlarkOutput)
-      .module(module)
-      .build();
+        .output(starlarkOutput)
+        .module(module)
+        .build();
   }
 
   @VisibleForTesting
@@ -184,22 +189,21 @@ public Module load(String moduleToLoad) {
          * Check if the module is in the module set. If it is, return a module with an environment
          * of the module that was passed in via the module set.
          */
-        else if(isNativeJavaModule(targetModule)) {
+        else if (isNativeJavaModule(targetModule)) {
           loadedModule = fromNativeModule(targetModule);
-        }
-        else {
+        } else {
           // try to load from directory...
           ResourceContentStarFile starFile = ResourceContentStarFile.buildStarFile(moduleToLoad);
           loadedModule = evaluator.eval(starFile).getModule();
         }
 
       } catch (IOException | InterruptedException | EvalException e) {
         throw new RuntimeException(
-          String.format(
-            "Encountered error (%s) while attempting to load %s from module: %s.",
-            e.getMessage(),
-            moduleToLoad,
-            this.content.path()), e);
+            String.format(
+                "Encountered error (%s) while attempting to load %s from module: %s.",
+                e.getMessage(),
+                moduleToLoad,
+                this.content.path()), e);
       }
       return loadedModule;
     }
@@ -232,12 +236,17 @@ private Module fromNativeModule(String moduleToLoad) throws IOException, Interru
       // TODO(mahmoudimus): Move this to ModuleSupplier?
       try (Mutability mu = Mutability.create("InMemoryNativeModule")) {
         StarlarkThread thread = new StarlarkThread(mu, evaluator.getLarkySemantics());
+        val environment1 = evaluator.getEnvironment();
+        if (environment1.containsKey("STEP_LIMIT") && environment1.get("STEP_LIMIT") != null) {
+          thread.setMaxExecutionSteps(Integer.parseInt(environment1.get("STEP_LIMIT").toString()));
+        }
+
         try {
           Starlark.execFile(
-            ParserInput.fromString(String.format("%1$s = _%1$s", moduleToLoad), "<builtin>"),
-            evaluator.getStarlarkValidationOptions(),
-            newModule,
-            thread
+              ParserInput.fromString(String.format("%1$s = _%1$s", moduleToLoad), "<builtin>"),
+              evaluator.getStarlarkValidationOptions(),
+              newModule,
+              thread
           );
         } catch (InterruptedException | EvalException | SyntaxError.Exception e) {
           throw new StarlarkEvalWrapper.Exc.RuntimeEvalException(e, thread);
@@ -276,7 +285,7 @@ Program compileStarlarkProgram(Module module, ParserInput input, FileOptions opt
       throw new EvalException(
           String.format(
               "Error compiling Starlark program: %1$s%n" +
-                  "%2$s",
+              "%2$s",
               input.getFile(),
               String.join("\n", errs)));
     }
@@ -307,10 +316,10 @@ private RuntimeException throwCycleError(String cycleElement) throws EvalExcepti
   }
 
   /**
-   * Create the environment for all evaluations (will be shared between all the dependent files
-   * loaded).
+   * Create the environment for all evaluations (will be shared between all the dependent files loaded).
    */
-  private ImmutableMap<String, Object> createEnvironment(Iterable<Class<?>> globalModules, Map<String, Object> globals) {
+  private ImmutableMap<String, Object> createEnvironment(Iterable<Class<?>> globalModules,
+      Map<String, Object> globals) {
     Map<String, Object> env = Maps.newHashMap();
 
     for (Class<?> module : globalModules) {

diff --git a/libstarlark/src/main/java/net/starlark/java/eval/Eval.java b/libstarlark/src/main/java/net/starlark/java/eval/Eval.java
@@ -261,7 +261,8 @@ private static TokenKind exec(StarlarkThread.Frame fr, Statement st)
     }
 
     if (++fr.thread.steps >= fr.thread.stepLimit) {
-      throw new EvalException("Starlark computation cancelled: too many steps");
+      throw new EvalException("Starlark computation cancelled: too many steps, limit is"
+                              + fr.thread.stepLimit + ", check your loops.");
     }
 
     switch (st.kind()) {
@@ -468,7 +469,7 @@ private static Object inplaceBinaryOp(StarlarkThread.Frame fr, TokenKind op, Obj
       xDict.putEntries(yDict);
       return xDict;
     }
-    return EvalUtils.binaryOp(op, x, y, fr.thread);
+    return EvalUtils.binaryOp(op, x, y, fr);
   }
 
   // ---- expressions ----
@@ -542,7 +543,7 @@ private static Object evalBinaryOperator(StarlarkThread.Frame fr, BinaryOperator
       default:
         Object y = eval(fr, binop.getY());
         try {
-          return EvalUtils.binaryOp(binop.getOperator(), x, y, fr.thread);
+          return EvalUtils.binaryOp(binop.getOperator(), x, y, fr);
         } catch (EvalException ex) {
           fr.setErrorLocation(binop.getOperatorLocation());
           throw ex;

diff --git a/libstarlark/src/main/java/net/starlark/java/eval/EvalUtils.java b/libstarlark/src/main/java/net/starlark/java/eval/EvalUtils.java
@@ -15,12 +15,16 @@
 
 import com.google.common.base.Strings;
 import java.util.IllegalFormatException;
+import net.starlark.java.eval.StarlarkThread.Frame;
 import net.starlark.java.syntax.TokenKind;
 
-/** Internal declarations used by the evaluator. */
+/**
+ * Internal declarations used by the evaluator.
+ */
 final class EvalUtils {
 
-  private EvalUtils() {}
+  private EvalUtils() {
+  }
 
   static void addIterator(Object x) {
     if (x instanceof Mutability.Freezable) {
@@ -37,8 +41,8 @@ static void removeIterator(Object x) {
   // The following functions for indexing and slicing match the behavior of Python.
 
   /**
-   * Resolves a positive or negative index to an index in the range [0, length), or throws
-   * EvalException if it is out of range. If the index is negative, it counts backward from length.
+   * Resolves a positive or negative index to an index in the range [0, length), or throws EvalException if it is out of
+   * range. If the index is negative, it counts backward from length.
    */
   static int getSequenceIndex(int index, int length) throws EvalException {
     int actualIndex = index;
@@ -53,10 +57,9 @@ static int getSequenceIndex(int index, int length) throws EvalException {
   }
 
   /**
-   * Returns the effective index denoted by a user-supplied integer. First, if the integer is
-   * negative, the length of the sequence is added to it, so an index of -1 represents the last
-   * element of the sequence. Then, the integer is "clamped" into the inclusive interval [0,
-   * length].
+   * Returns the effective index denoted by a user-supplied integer. First, if the integer is negative, the length of
+   * the sequence is added to it, so an index of -1 represents the last element of the sequence. Then, the integer is
+   * "clamped" into the inclusive interval [0, length].
    */
   static int toIndex(int index, int length) {
     if (index < 0) {
@@ -72,9 +75,12 @@ static int toIndex(int index, int length) {
     }
   }
 
-  /** Evaluates an eager binary operation, {@code x op y}. (Excludes AND and OR.) */
-  static Object binaryOp(TokenKind op, Object x, Object y, StarlarkThread starlarkThread)
+  /**
+   * Evaluates an eager binary operation, {@code x op y}. (Excludes AND and OR.)
+   */
+  static Object binaryOp(TokenKind op, Object x, Object y, Frame fr)
       throws EvalException {
+    StarlarkThread starlarkThread = fr.thread;
     StarlarkSemantics semantics = starlarkThread.getSemantics();
     Mutability mu = starlarkThread.mutability();
     switch (op) {
@@ -91,6 +97,19 @@ static Object binaryOp(TokenKind op, Object x, Object y, StarlarkThread starlark
 
         } else if (x instanceof String) {
           if (y instanceof String) {
+            // MEMORY ALLOCATION CHECK
+            if (((StarlarkFunction) fr.fn).getModule().getPredeclaredBindings()
+                .containsKey("MAX_STRING_LENGTH")) {
+              long maxLength = Long.parseLong(
+                  ((StarlarkFunction) fr.fn).getModule().getPredeclared("MAX_STRING_LENGTH").toString());
+              final long length = ((String) x).length() + ((String) y).length();
+              if (maxLength > 0 && length > maxLength) {
+                throw new EvalException(
+                    "Exceeded string concatenation limit, wanted "
+                    + length + ", limit is " + maxLength);
+              }
+            }
+
             // string + string
             return (String) x + (String) y;
           }
@@ -357,7 +376,7 @@ static Object binaryOp(TokenKind op, Object x, Object y, StarlarkThread starlark
         break;
 
       case NOT_IN:
-        Object z = binaryOp(TokenKind.IN, x, y, starlarkThread);
+        Object z = binaryOp(TokenKind.IN, x, y, fr);
         if (z != null) {
           return !Starlark.truth(z);
         }
@@ -409,7 +428,9 @@ private static String repeatString(String s, StarlarkInt in) throws EvalExceptio
     }
   }
 
-  /** Evaluates a unary operation. */
+  /**
+   * Evaluates a unary operation.
+   */
   static Object unaryOp(TokenKind op, Object x) throws EvalException {
     switch (op) {
       case NOT:
@@ -497,7 +518,9 @@ static void setIndex(Object object, Object key, Object value) throws EvalExcepti
     }
   }
 
-  /** Updates the named field of x as if by the Starlark statement {@code x.field = value}. */
+  /**
+   * Updates the named field of x as if by the Starlark statement {@code x.field = value}.
+   */
   static void setField(Object x, String field, Object value) throws EvalException {
     if (x instanceof Structure) {
       ((Structure) x).setField(field, value);

diff --git a/libstarlark/src/main/java/net/starlark/java/eval/StarlarkBytes.java b/libstarlark/src/main/java/net/starlark/java/eval/StarlarkBytes.java
@@ -19,15 +19,13 @@
 import java.util.ListIterator;
 import java.util.function.UnaryOperator;
 import java.util.stream.Collectors;
-
 import net.starlark.java.annot.Param;
 import net.starlark.java.annot.ParamType;
 import net.starlark.java.annot.StarlarkBuiltin;
 import net.starlark.java.annot.StarlarkMethod;
 import net.starlark.java.ext.ByteList;
 import net.starlark.java.ext.ByteStringModuleApi;
 import net.starlark.java.syntax.TokenKind;
-
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
@@ -161,8 +159,8 @@ public int hashCode() {
     @Override
     public Object binaryOp(TokenKind op, Object that, boolean thisLeft) throws EvalException {
       try(Mutability mu = Mutability.create("StarlarkBytesBinaryOp")) {
-        StarlarkThread thread = new StarlarkThread(mu, StarlarkSemantics.DEFAULT);
-        return EvalUtils.binaryOp(op, toStarlarkInt(), that, thread);
+//        StarlarkThread thread = new StarlarkThread(mu, StarlarkSemantics.DEFAULT);
+        return null;//EvalUtils.binaryOp(op, toStarlarkInt(), that, thread);
       }
     }