io.codemodder.remediation.xss.PrintingMethodFixer Maven / Gradle / Ivy

Go to download
package io.codemodder.remediation.xss;

import static io.codemodder.javaparser.JavaParserTransformer.wrap;

import com.github.javaparser.ast.CompilationUnit;
import com.github.javaparser.ast.Node;
import com.github.javaparser.ast.expr.Expression;
import com.github.javaparser.ast.expr.MethodCallExpr;
import com.github.javaparser.ast.expr.NameExpr;
import com.github.javaparser.ast.nodeTypes.NodeWithType;
import io.codemodder.ast.ASTs;
import io.codemodder.codetf.DetectorRule;
import io.codemodder.remediation.RemediationMessages;
import java.util.List;
import java.util.Optional;
import java.util.Set;
import java.util.function.Function;

/**
 * Fixes a method invocation that prints {@link String} expression.
 *
 * {@code
 * out.println(foo); // should become return out.println(Encode.forHtml(foo));
 * }
 */
final class PrintingMethodFixer implements XSSCodeShapeFixer {

  private static final Set writingMethodNames = Set.of("print", "println", "write");

  @Override
  public  XSSCodeShapeFixResult fixCodeShape(
      final CompilationUnit cu,
      final String path,
      final DetectorRule detectorRule,
      final List issues,
      final Function getKey,
      final Function getLine,
      final Function getColumn) {
    int line = getLine.apply(issues.get(0));
    Integer column = getColumn.apply(issues.get(0));

    List writingMethodCalls =
        cu.findAll(MethodCallExpr.class).stream()
            .filter(mce -> writingMethodNames.contains(mce.getNameAsString()))
            .filter(mce -> mce.getArguments().size() == 1)
            .filter(mce -> mce.getRange().isPresent())
            .filter(mce -> mce.getRange().get().begin.line == line)
            .filter(mce -> column == null || mce.getRange().get().begin.column == column)
            .filter(
                mce -> {
                  Expression arg = mce.getArgument(0);
                  if (arg instanceof NameExpr nameExpr) {
                    Optional source =
                        ASTs.findNonCallableSimpleNameSource(nameExpr.getName());
                    if (source.isPresent()
                        && source.get() instanceof NodeWithType typedNode) {
                      String typeAsString = typedNode.getTypeAsString();
                      return "String".equals(typeAsString)
                          || "java.lang.String".equals(typeAsString);
                    }
                  }
                  return false;
                })
            .toList();

    if (writingMethodCalls.isEmpty()) {
      return new XSSCodeShapeFixResult(false, false, null, line);
    } else if (writingMethodCalls.size() > 1) {
      return new XSSCodeShapeFixResult(true, false, RemediationMessages.multipleNodesFound, line);
    }

    MethodCallExpr call = writingMethodCalls.get(0);
    wrap(call.getArgument(0)).withStaticMethod("org.owasp.encoder.Encode", "forHtml", false);
    return new XSSCodeShapeFixResult(true, true, null, line);
  }
}