Modified transformation to use security package

pixee · Dec 17, 2024 · b8b90cd · b8b90cd
1 parent 25f1f06
commit b8b90cd
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 57 deletions.
diff --git a/core-codemods/src/test/resources/codeql-regexdos/RegexUtil.java.after b/core-codemods/src/test/resources/codeql-regexdos/RegexUtil.java.after
@@ -18,11 +18,10 @@
 
 package org.apache.roller.util;
 
+import io.github.pixee.security.ExecuteWithTimeout;
 import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
 import java.util.List;
-import java.util.concurrent.Callable;
-import java.util.concurrent.Executors;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 import org.apache.commons.codec.binary.Hex;
@@ -62,7 +61,7 @@ public final class RegexUtil {
      */
     public static String obfuscateEmail(String str) {
         Matcher emailMatch = EMAIL_PATTERN.matcher(str);
-        while (executeWithTimeout(() -> emailMatch.find(), 5000)) {
+        while (ExecuteWithTimeout.executeWithTimeout(() -> emailMatch.find(), 5000)) {
             String at = emailMatch.group(1);
             //System.out.println("at=" + at);
             str = str.replaceFirst(at, "-AT-");
@@ -115,13 +114,4 @@ public final class RegexUtil {
         return result.toString();
     }
 
-    public <E> E executeWithTimeout(final Callable<E> action, final int timeout) {
-        Future<E> maybeResult = Executors.newSingleThreadExecutor().submit(action);
-        try {
-            return maybeResult.get(timeout, TimeUnit.MILLISECONDS);
-        } catch (Exception e) {
-            throw new RuntimeException("Failed to execute within time limit.");
-        }
-    }
-
 }
diff --git a/...codemodder-base/src/main/java/io/codemodder/remediation/regexdos/RegexDoSFixStrategy.java b/...codemodder-base/src/main/java/io/codemodder/remediation/regexdos/RegexDoSFixStrategy.java
@@ -1,15 +1,10 @@
 package io.codemodder.remediation.regexdos;
 
-import com.github.javaparser.StaticJavaParser;
 import com.github.javaparser.ast.CompilationUnit;
 import com.github.javaparser.ast.Node;
 import com.github.javaparser.ast.NodeList;
-import com.github.javaparser.ast.body.ClassOrInterfaceDeclaration;
-import com.github.javaparser.ast.body.MethodDeclaration;
-import com.github.javaparser.ast.expr.Expression;
-import com.github.javaparser.ast.expr.IntegerLiteralExpr;
-import com.github.javaparser.ast.expr.LambdaExpr;
-import com.github.javaparser.ast.expr.MethodCallExpr;
+import com.github.javaparser.ast.expr.*;
+import io.codemodder.DependencyGAV;
 import io.codemodder.ast.ASTTransforms;
 import io.codemodder.ast.ASTs;
 import io.codemodder.ast.LocalDeclaration;
@@ -18,8 +13,6 @@
 import io.codemodder.remediation.SuccessOrReason;
 import java.util.List;
 import java.util.Optional;
-import java.util.concurrent.Callable;
-import java.util.concurrent.Executors;
 
 /** Adds a timeout function and wraps regex match call with it * */
 final class RegexDoSFixStrategy extends MatchAndFixStrategy {
@@ -51,33 +44,6 @@ public boolean match(final Node node) {
         .isPresent();
   }
 
-  private static void addTimeoutMethodIfMissing(
-      final CompilationUnit cu, final ClassOrInterfaceDeclaration classDecl) {
-    final String method =
-        """
-                public <E> E executeWithTimeout(final Callable<E> action, final int timeout){
-                  Future<E> maybeResult = Executors.newSingleThreadExecutor().submit(action);
-                  try{
-                      return maybeResult.get(timeout, TimeUnit.MILLISECONDS);
-                  }catch(Exception e){
-                      throw new RuntimeException("Failed to execute within time limit.");
-                  }
-                }
-            """;
-    boolean filterMethodPresent =
-        classDecl.findAll(MethodDeclaration.class).stream()
-            .anyMatch(
-                md ->
-                    md.getNameAsString().equals("executeWithTimeout")
-                        && md.getParameters().size() == 2);
-    if (!filterMethodPresent) {
-      classDecl.addMember(StaticJavaParser.parseMethodDeclaration(method));
-    }
-    // Add needed import
-    ASTTransforms.addImportIfMissing(cu, Callable.class.getName());
-    ASTTransforms.addImportIfMissing(cu, Executors.class.getName());
-  }
-
   @Override
   public SuccessOrReason fix(final CompilationUnit cu, final Node node) {
     // indirect case, assigned to a variable
@@ -91,21 +57,20 @@ public SuccessOrReason fix(final CompilationUnit cu, final Node node) {
     if (allValidMethodCalls.isEmpty()) {
       return SuccessOrReason.reason("Couldn't find any matching methods");
     }
-    // Add executeWithTimout method to the encompassing class
-    var classDecl = call.findAncestor(ClassOrInterfaceDeclaration.class);
-    if (classDecl.isEmpty()) {
-      return SuccessOrReason.reason("Couldn't find encompassing class");
-    }
-    classDecl.ifPresent(cd -> addTimeoutMethodIfMissing(cu, cd));
+
     for (var mce : allValidMethodCalls) {
       // Wrap it with executeWithTimeout with a default 5000 of timeout
       var newCall =
           new MethodCallExpr(
+              new NameExpr("ExecuteWithTimeout"),
               "executeWithTimeout",
-              new LambdaExpr(new NodeList<>(), mce.clone()),
-              new IntegerLiteralExpr(DEFAULT_TIMEOUT));
+              new NodeList<>(
+                  new LambdaExpr(new NodeList<>(), mce.clone()),
+                  new IntegerLiteralExpr(DEFAULT_TIMEOUT)));
       mce.replace(newCall);
     }
-    return SuccessOrReason.success();
+
+    ASTTransforms.addImportIfMissing(cu, "io.github.pixee.security.ExecuteWithTimeout");
+    return SuccessOrReason.success(List.of(DependencyGAV.JAVA_SECURITY_TOOLKIT));
   }
 }