andrecsilva · Apr 12, 2024 · Apr 5, 2024 · Apr 8, 2024 · Apr 8, 2024 · Apr 9, 2024
diff --git a/core-codemods/src/main/java/io/codemodder/codemods/DefaultCodemods.java b/core-codemods/src/main/java/io/codemodder/codemods/DefaultCodemods.java
@@ -47,6 +47,7 @@ public static List<Class<? extends CodeChanger>> asList() {
         RemoveUselessParenthesesCodemod.class,
         ReplaceDefaultHttpClientCodemod.class,
         ReplaceStreamCollectorsToListCodemod.class,
+        // ResourceLeakCodemod.class,
         SanitizeApacheMultipartFilenameCodemod.class,
         SanitizeHttpHeaderCodemod.class,
         SanitizeSpringMultipartFilenameCodemod.class,

diff --git a/...-codemods/src/main/java/io/codemodder/codemods/PreventFileWriterLeakWithFilesCodemod.java b/...-codemods/src/main/java/io/codemodder/codemods/PreventFileWriterLeakWithFilesCodemod.java
@@ -54,6 +54,10 @@ public ChangesResult onResultFound(
     newFilesBufferedWriterCall.setArguments(NodeList.nodeList(pathArgument));
     parent.replace(newBufferedWriterCall, newFilesBufferedWriterCall);
     addImportIfMissing(cu, "java.nio.file.Files");
+
+    // try to wrap it in a try statement
+    ResourceLeakFixer.checkAndFix(newFilesBufferedWriterCall);
+
     return ChangesResult.changesApplied;
   }
 }
diff --git a/core-codemods/src/main/java/io/codemodder/codemods/ResourceLeakCodemod.java b/core-codemods/src/main/java/io/codemodder/codemods/ResourceLeakCodemod.java
@@ -0,0 +1,37 @@
+package io.codemodder.codemods;
+
+import com.github.javaparser.ast.CompilationUnit;
+import com.github.javaparser.ast.expr.Expression;
+import io.codemodder.*;
+import io.codemodder.javaparser.JavaParserChanger;
+import java.util.List;
+import java.util.Optional;
+import java.util.stream.Collectors;
+
+/** A codemod that wraps AutoCloseable objects whenever possible. */
+@Codemod(
+    id = "pixee:java/resource-leak",
+    reviewGuidance = ReviewGuidance.MERGE_WITHOUT_REVIEW,
+    importance = Importance.MEDIUM,
+    executionPriority = CodemodExecutionPriority.LOW)
+public final class ResourceLeakCodemod extends JavaParserChanger {
+
+  private Optional<CodemodChange> onNodeFound(final Expression expr) {
+    int originalLine = expr.getBegin().get().line;
+    if (ResourceLeakFixer.checkAndFix(expr).isPresent()) {
+      return Optional.of(CodemodChange.from(originalLine));
+    } else {
+      return Optional.empty();
+    }
+  }
+
+  @Override
+  public CodemodFileScanningResult visit(
+      final CodemodInvocationContext context, final CompilationUnit cu) {
+    List<CodemodChange> changes =
+        cu.findAll(Expression.class).stream()
+            .flatMap(expr -> onNodeFound(expr).stream())
+            .collect(Collectors.toList());
+    return CodemodFileScanningResult.withOnlyChanges(changes);
+  }
+}
diff --git a/core-codemods/src/main/java/io/codemodder/codemods/ResourceLeakFixer.java b/core-codemods/src/main/java/io/codemodder/codemods/ResourceLeakFixer.java
@@ -5,6 +5,7 @@
 import com.github.javaparser.ast.body.Parameter;
 import com.github.javaparser.ast.body.VariableDeclarator;
 import com.github.javaparser.ast.expr.*;
+import com.github.javaparser.ast.stmt.TryStmt;
 import com.github.javaparser.resolution.UnsolvedSymbolException;
 import io.codemodder.Either;
 import io.codemodder.ast.*;
@@ -14,6 +15,7 @@
 import java.util.HashSet;
 import java.util.List;
 import java.util.Optional;
+import java.util.Set;
 import java.util.function.Predicate;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
@@ -31,11 +33,20 @@ final class ResourceLeakFixer {
 
   private static final String rootPrefix = "resource";
 
+  private static Set<String> initMethodsList =
+      Set.of(
+          "newBufferedReader",
+          "newBufferedWriter",
+          "newByteChannel",
+          "newDirectoryStream",
+          "newInputStream",
+          "newOutputStream");
+
   /**
    * Detects if an {@link Expression} that creates a resource type is fixable and tries to fix it.
    * Combines {@code isFixable} and {@code tryToFix}.
    */
-  public static Optional<Integer> checkAndFix(final Expression expr) {
+  public static Optional<TryStmt> checkAndFix(final Expression expr) {
     if (expr instanceof ObjectCreationExpr) {
       // finds the root expression in a chain of new AutoCloseable objects
       ObjectCreationExpr root = findRootExpression(expr.asObjectCreationExpr());
@@ -53,15 +64,6 @@ public static Optional<Integer> checkAndFix(final Expression expr) {
 
   /**
    * Detects if a {@link Expression} detected by CodeQL that creates a leaking resource is fixable.
-   * The following can be assumed.
-   *
-   * <p>(1) No close() is called within its scope.
-   *
-   * <p>(2) There does not exist a "root" expression that is closed. (e.g. new BufferedReader(r), r
-   * is the root, stmt.executeQuery(...), stmt is the root).
-   *
-   * <p>(3) It does not escape the contained method's scope. That is, it is not assigned to a field
-   * or returned.
    *
    * <p>A resource R is dependent of another resource S if closing S will also close R. The
    * following suffices to check if a resource R can be closed. (*) A resource R can be closed if no
@@ -70,6 +72,15 @@ public static Optional<Integer> checkAndFix(final Expression expr) {
    * (+): S is assigned to a variable that escapes.
    */
   public static boolean isFixable(final Expression expr) {
+    // Can it be wrapped in a try statement?
+    if (!isAutoCloseableType(expr)) {
+      return false;
+    }
+    // is it already closed? does it escape?
+    // TODO depends on another resource that is already closed?
+    if (isClosedOrEscapes(expr)) {
+      return false;
+    }
     // For ResultSet objects, still need to check if the generating *Statement object does
     // not escape due to the getResultSet() method.
     try {
@@ -97,18 +108,46 @@ public static boolean isFixable(final Expression expr) {
       return false;
     }
 
+    // For any dependent resource s of r if s is not closed and escapes, then r cannot be closed
     final var maybeLVD = immediatelyFlowsIntoLocalVariable(expr);
-    final var allDependent = findDependentResources(expr);
     if (maybeLVD.isPresent()) {
       final var scope = maybeLVD.get().getScope();
       final Predicate<Node> isInScope = scope::inScope;
+      final var allDependent = findDependentResources(expr);
       return allDependent.stream().noneMatch(e -> escapesRootScope(e, isInScope));
     } else {
-      final Predicate<Node> isInScope = n -> true;
-      return allDependent.stream().noneMatch(e -> escapesRootScope(e, isInScope));
+      final var allDependent = findDependentResources(expr);
+      return allDependent.stream().noneMatch(e -> escapesRootScope(e, n -> true));
     }
   }
 
+  private static boolean isClosedOrEscapes(final Expression expr) {
+    if (immediatelyEscapesMethodScope(expr)) {
+      return true;
+    }
+    // find all the variables it flows into
+    final var allVariables = flowsInto(expr);
+
+    // flows into anything that is not a local variable or parameter
+    if (allVariables.stream().anyMatch(Either::isRight)) {
+      return true;
+    }
+
+    // is any of the assigned variables closed?
+    if (allVariables.stream()
+        .filter(Either::isLeft)
+        .map(Either::getLeft)
+        .anyMatch(ld -> !notClosed(ld))) {
+      return true;
+    }
+
+    // If any of the assigned variables escapes
+    return allVariables.stream()
+        .filter(Either::isLeft)
+        .map(Either::getLeft)
+        .anyMatch(ld -> escapesRootScope(ld, x -> true));
+  }
+
   public static ObjectCreationExpr findRootExpression(final ObjectCreationExpr creationExpr) {
     ObjectCreationExpr current = creationExpr;
     var maybeInner = Optional.of(current);
@@ -121,8 +160,7 @@ public static ObjectCreationExpr findRootExpression(final ObjectCreationExpr cre
     return current;
   }
 
-  public static Optional<Integer> tryToFix(final ObjectCreationExpr creationExpr) {
-    final int originalLine = creationExpr.getRange().get().begin.line;
+  public static Optional<TryStmt> tryToFix(final ObjectCreationExpr creationExpr) {
     final var deque = findInnerExpressions(creationExpr);
     int count = 0;
     final var maybeLVD =
@@ -143,6 +181,7 @@ public static Optional<Integer> tryToFix(final ObjectCreationExpr creationExpr)
       var cu = creationExpr.findCompilationUnit().get();
       for (var resource : deque) {
         var name = count == 0 ? rootPrefix : rootPrefix + count;
+        // TODO try to resolve type, failing to do that use var declaration?
         var type = resource.calculateResolvedType().describe();
         resource.replace(new NameExpr(name));
 
@@ -155,14 +194,13 @@ public static Optional<Integer> tryToFix(final ObjectCreationExpr creationExpr)
         ASTTransforms.addImportIfMissing(cu, type);
         count++;
       }
-      return Optional.of(originalLine);
+      return Optional.of(tryStmt);
     }
     return Optional.empty();
   }
 
   /** Tries to fix the leak of {@code expr} and returns its line if it does. */
-  public static Optional<Integer> tryToFix(final MethodCallExpr mce) {
-    final int originalLine = mce.getRange().get().begin.line;
+  public static Optional<TryStmt> tryToFix(final MethodCallExpr mce) {
 
     // Is LocalDeclarationStmt and Never Assigned -> Wrap as a try resource
     List<Expression> resources = new ArrayList<>();
@@ -187,6 +225,7 @@ public static Optional<Integer> tryToFix(final MethodCallExpr mce) {
       var cu = mce.findCompilationUnit().get();
       for (var resource : resources) {
         var name = count == 0 ? rootPrefix : rootPrefix + count;
+        // TODO try to resolve type, failing to do that use var declaration?
         var type = resource.calculateResolvedType().describe();
 
         resource.replace(new NameExpr(name));
@@ -200,7 +239,7 @@ public static Optional<Integer> tryToFix(final MethodCallExpr mce) {
         ASTTransforms.addImportIfMissing(cu, type);
         count++;
       }
-      return Optional.of(originalLine);
+      return Optional.of(tryStmt);
       // TODO if vde is multiple declarations, extract the relevant vd and wrap it
     }
     // other cases here...
@@ -253,9 +292,14 @@ private static Optional<LocalVariableDeclaration> immediatelyFlowsIntoLocalVaria
 
   /** Checks if the expression implements the {@link java.lang.AutoCloseable} interface. */
   private static boolean isAutoCloseableType(final Expression expr) {
-    return expr.calculateResolvedType().isReferenceType()
-        && expr.calculateResolvedType().asReferenceType().getAllAncestors().stream()
-            .anyMatch(t -> t.describe().equals("java.lang.AutoCloseable"));
+    try {
+      return expr.calculateResolvedType().isReferenceType()
+          && expr.calculateResolvedType().asReferenceType().getAllAncestors().stream()
+              .anyMatch(t -> t.describe().equals("java.lang.AutoCloseable"));
+    } catch (RuntimeException e) {
+      LOG.error("Problem resolving type of : {}", expr);
+      return false;
+    }
   }
 
   /** Checks if the expression creates a {@link java.lang.AutoCloseable} */
@@ -268,7 +312,9 @@ private static Optional<ObjectCreationExpr> isAutoCloseableCreation(final Expres
 
   /** Checks if the expression creates a {@link java.lang.AutoCloseable} */
   private static boolean isResourceInit(final Expression expr) {
-    return (expr.isMethodCallExpr() && isJDBCResourceInit(expr.asMethodCallExpr()))
+    return (expr.isMethodCallExpr()
+            && (isJDBCResourceInit(expr.asMethodCallExpr())
+                || isFilesResourceInit(expr.asMethodCallExpr())))
         || isAutoCloseableCreation(expr).isPresent();
   }
 
@@ -285,40 +331,40 @@ private static Either<LocalDeclaration, Node> isLocalDeclaration(Node n) {
     return Either.right(n);
   }
 
+  /** Checks if {@code expr} creates an AutoCloseable Resource. */
+  private static boolean isFilesResourceInit(final MethodCallExpr expr) {
+    try {
+      var hasFilesScope =
+          expr.getScope()
+              .map(mce -> mce.calculateResolvedType().describe())
+              .filter("java.nio.file.Files"::equals);
+      return hasFilesScope.isPresent() && initMethodsList.contains(expr.getNameAsString());
+    } catch (final UnsolvedSymbolException e) {
+      LOG.error("Problem resolving type of : {}", expr, e);
+    }
+    return false;
+  }
+
   /** Checks if {@code expr} creates a JDBC Resource. */
   private static boolean isJDBCResourceInit(final MethodCallExpr expr) {
     final Predicate<MethodCallExpr> isResultSetGen =
-        mce -> {
-          switch (mce.getNameAsString()) {
-            case "executeQuery":
-            case "getResultSet":
-            case "getGeneratedKeys":
-              return true;
-            default:
-              return false;
-          }
-        };
+        mce ->
+            switch (mce.getNameAsString()) {
+              case "executeQuery", "getResultSet", "getGeneratedKeys" -> true;
+              default -> false;
+            };
     final Predicate<MethodCallExpr> isStatementGen =
-        mce -> {
-          switch (mce.getNameAsString()) {
-            case "createStatement":
-            case "prepareCall":
-            case "prepareStatement":
-              return true;
-            default:
-              return false;
-          }
-        };
+        mce ->
+            switch (mce.getNameAsString()) {
+              case "createStatement", "prepareCall", "prepareStatement" -> true;
+              default -> false;
+            };
     final Predicate<MethodCallExpr> isReaderGen =
-        mce -> {
-          switch (mce.getNameAsString()) {
-            case "getCharacterStream":
-            case "getNCharacterStream":
-              return true;
-            default:
-              return false;
-          }
-        };
+        mce ->
+            switch (mce.getNameAsString()) {
+              case "getCharacterStream", "getNCharacterStream" -> true;
+              default -> false;
+            };
     final Predicate<MethodCallExpr> isDependent = isResultSetGen.or(isStatementGen.or(isReaderGen));
     return isDependent.test(expr);
   }
@@ -530,6 +576,7 @@ private static boolean immediatelyEscapesMethodScope(final Expression expr) {
     if (!isResourceInit(expr)) {
       return true;
     }
+
     // Returned or argument of a MethodCallExpr
     if (ASTs.isReturnExpr(expr).isPresent() || ASTs.isArgumentOfMethodCall(expr).isPresent()) {
       return true;

diff --git a/...ds/src/main/resources/io/codemodder/codemods/ResourceLeakCodemod/description.md b/...ds/src/main/resources/io/codemodder/codemods/ResourceLeakCodemod/description.md
@@ -0,0 +1,11 @@
+This change adds [try-with-resources](https://docs.oracle.com/javase/tutorial/essential/exceptions/tryResourceClose.html) to code to prevent resources from being leaked, which could lead to denial-of-service conditions like connection pool or file handle exhaustion.
+
+Our changes look something like this:
+
+```diff
+- BufferedReader br = new BufferedReader(new FileReader("C:\\test.txt"));
+- System.out.println(br.readLine());
++ try(FileReader input = new FileReader("C:\\test.txt"); BufferedReader br = new BufferedReader(input)){
++   System.out.println(br.readLine());
++ }
+```
diff --git a/core-codemods/src/main/resources/io/codemodder/codemods/ResourceLeakCodemod/report.json b/core-codemods/src/main/resources/io/codemodder/codemods/ResourceLeakCodemod/report.json
@@ -0,0 +1,9 @@
+{
+  "summary" : "Prevent resource leaks",
+  "change": "Added a try-with-resources statement to automatically close resources",
+  "reviewGuidanceIJustification" : "This codemod causes resources to be cleaned up immediately after use instead of at garbage collection time, and we don't believe this change entails any risk.",
+  "references" : [
+    "https://cwe.mitre.org/data/definitions/404.html",
+    "https://cwe.mitre.org/data/definitions/772.html"
+  ]
+}
diff --git a/core-codemods/src/test/java/io/codemodder/codemods/ResourceLeakCodemodTest.java b/core-codemods/src/test/java/io/codemodder/codemods/ResourceLeakCodemodTest.java
@@ -0,0 +1,10 @@
+package io.codemodder.codemods;
+
+import io.codemodder.testutils.CodemodTestMixin;
+import io.codemodder.testutils.Metadata;
+
+@Metadata(
+    codemodType = ResourceLeakCodemod.class,
+    testResourceDir = "resource-leak",
+    dependencies = {})
+final class ResourceLeakCodemodTest implements CodemodTestMixin {}
diff --git a/core-codemods/src/test/resources/prevent-filewriter-leak-with-nio/Test.java.after b/core-codemods/src/test/resources/prevent-filewriter-leak-with-nio/Test.java.after
@@ -9,18 +9,14 @@ public class Test
 {
     private Response response;
     void foo(File f) {
-        // ruleid: prevent-filewriter-leak-with-nio
-        Writer change1 = Files.newBufferedWriter(f.toPath());
-        // ruleid: prevent-filewriter-leak-with-nio
-        Writer change2 = Files.newBufferedWriter(f.toPath());
+        try (Writer change1 = Files.newBufferedWriter(f.toPath())) {
+            try (Writer change2 = Files.newBufferedWriter(f.toPath())) {
+                Writer dontChange1 = new BufferedWriter(w);
+                Writer dontChange2 = new BufferedWriter(w, 256);
+                Writer dontChange3 = new BufferedWriter(response.getWriter(), 256);
+                Writer dontChange3 = Files.newBufferedWriter(f.toPath());
+            }
+        }
 
-        // ok: prevent-filewriter-leak-with-nio
-        Writer dontChange1 = new BufferedWriter(w);
-        // ok: prevent-filewriter-leak-with-nio
-        Writer dontChange2 = new BufferedWriter(w, 256);
-        // ok: prevent-filewriter-leak-with-nio
-        Writer dontChange3 = new BufferedWriter(response.getWriter(), 256);
-        // ok: prevent-filewriter-leak-with-nio
-        Writer dontChange3 = Files.newBufferedWriter(f.toPath());
     }
 }