Add builder for vulnerability types and fix insecure auth protocol

dd-java-agent/agent-iast/src/main/java/com/datadog/iast/model/Vulnerability.java

@@ -50,9 +50,7 @@ public long getHash() {
   public void updateSpan(final AgentSpan newSpan) {
     if (newSpan != null) {
       location.updateSpan(newSpan);
-      if (type instanceof VulnerabilityType.HeaderVulnerabilityType) {
-        hash = type.calculateHash(this);
-      }
+      hash = type.calculateHash(this);
     }
   }
 

dd-java-agent/agent-iast/src/main/java/com/datadog/iast/model/VulnerabilityType.java

@@ -1,96 +1,107 @@
 package com.datadog.iast.model;
 
+import static com.datadog.iast.util.CRCUtils.update;
+import static datadog.trace.api.iast.VulnerabilityMarks.COMMAND_INJECTION_MARK;
+import static datadog.trace.api.iast.VulnerabilityMarks.HEADER_INJECTION_MARK;
+import static datadog.trace.api.iast.VulnerabilityMarks.LDAP_INJECTION_MARK;
 import static datadog.trace.api.iast.VulnerabilityMarks.NOT_MARKED;
+import static datadog.trace.api.iast.VulnerabilityMarks.PATH_TRAVERSAL_MARK;
+import static datadog.trace.api.iast.VulnerabilityMarks.REFLECTION_INJECTION_MARK;
+import static datadog.trace.api.iast.VulnerabilityMarks.SQL_INJECTION_MARK;
+import static datadog.trace.api.iast.VulnerabilityMarks.SSRF_MARK;
+import static datadog.trace.api.iast.VulnerabilityMarks.TRUST_BOUNDARY_VIOLATION_MARK;
+import static datadog.trace.api.iast.VulnerabilityMarks.UNVALIDATED_REDIRECT_MARK;
+import static datadog.trace.api.iast.VulnerabilityMarks.XPATH_INJECTION_MARK;
+import static datadog.trace.api.iast.VulnerabilityMarks.XSS_MARK;
 
-import datadog.trace.api.iast.VulnerabilityMarks;
 import datadog.trace.api.iast.VulnerabilityTypes;
 import java.io.File;
-import java.nio.charset.StandardCharsets;
+import java.util.function.BiFunction;
 import java.util.zip.CRC32;
 import javax.annotation.Nonnull;
 
 public interface VulnerabilityType {
-  VulnerabilityType WEAK_CIPHER = new VulnerabilityTypeImpl(VulnerabilityTypes.WEAK_CIPHER);
-  VulnerabilityType WEAK_HASH = new VulnerabilityTypeImpl(VulnerabilityTypes.WEAK_HASH);
+
+  VulnerabilityType WEAK_CIPHER = type(VulnerabilityTypes.WEAK_CIPHER).build();
+  VulnerabilityType WEAK_HASH = type(VulnerabilityTypes.WEAK_HASH).build();
   VulnerabilityType INSECURE_COOKIE =
-      new CookieVulnerabilityType(VulnerabilityTypes.INSECURE_COOKIE);
+      type(VulnerabilityTypes.INSECURE_COOKIE).hash(VulnerabilityType::evidenceHash).build();
   VulnerabilityType NO_HTTPONLY_COOKIE =
-      new CookieVulnerabilityType(VulnerabilityTypes.NO_HTTPONLY_COOKIE);
+      type(VulnerabilityTypes.NO_HTTPONLY_COOKIE).hash(VulnerabilityType::evidenceHash).build();
   VulnerabilityType HSTS_HEADER_MISSING =
-      new HeaderVulnerabilityType(VulnerabilityTypes.HSTS_HEADER_MISSING);
+      type(VulnerabilityTypes.HSTS_HEADER_MISSING).hash(VulnerabilityType::serviceHash).build();
   VulnerabilityType XCONTENTTYPE_HEADER_MISSING =
-      new HeaderVulnerabilityType(VulnerabilityTypes.XCONTENTTYPE_HEADER_MISSING);
+      type(VulnerabilityTypes.XCONTENTTYPE_HEADER_MISSING)
+          .hash(VulnerabilityType::serviceHash)
+          .build();
   VulnerabilityType NO_SAMESITE_COOKIE =
-      new CookieVulnerabilityType(VulnerabilityTypes.NO_SAMESITE_COOKIE);
+      type(VulnerabilityTypes.NO_SAMESITE_COOKIE).hash(VulnerabilityType::evidenceHash).build();
 
   VulnerabilityType SQL_INJECTION =
-      new VulnerabilityTypeImpl(
-          VulnerabilityTypes.SQL_INJECTION, VulnerabilityMarks.SQL_INJECTION_MARK);
+      type(VulnerabilityTypes.SQL_INJECTION).mark(SQL_INJECTION_MARK).build();
   VulnerabilityType COMMAND_INJECTION =
-      new VulnerabilityTypeImpl(
-          VulnerabilityTypes.COMMAND_INJECTION, VulnerabilityMarks.COMMAND_INJECTION_MARK);
+      type(VulnerabilityTypes.COMMAND_INJECTION).mark(COMMAND_INJECTION_MARK).build();
   VulnerabilityType PATH_TRAVERSAL =
-      new VulnerabilityTypeImpl(
-          VulnerabilityTypes.PATH_TRAVERSAL,
-          File.separatorChar,
-          VulnerabilityMarks.PATH_TRAVERSAL_MARK);
+      type(VulnerabilityTypes.PATH_TRAVERSAL)
+          .separator(File.separatorChar)
+          .mark(PATH_TRAVERSAL_MARK)
+          .build();
   VulnerabilityType LDAP_INJECTION =
-      new VulnerabilityTypeImpl(
-          VulnerabilityTypes.LDAP_INJECTION, VulnerabilityMarks.LDAP_INJECTION_MARK);
-  VulnerabilityType SSRF =
-      new VulnerabilityTypeImpl(VulnerabilityTypes.SSRF, VulnerabilityMarks.SSRF_MARK);
+      type(VulnerabilityTypes.LDAP_INJECTION).mark(LDAP_INJECTION_MARK).build();
+  VulnerabilityType SSRF = type(VulnerabilityTypes.SSRF).mark(SSRF_MARK).build();
   VulnerabilityType UNVALIDATED_REDIRECT =
-      new VulnerabilityTypeImpl(
-          VulnerabilityTypes.UNVALIDATED_REDIRECT, VulnerabilityMarks.UNVALIDATED_REDIRECT_MARK);
-  VulnerabilityType WEAK_RANDOMNESS = new VulnerabilityTypeImpl(VulnerabilityTypes.WEAK_RANDOMNESS);
+      type(VulnerabilityTypes.UNVALIDATED_REDIRECT).mark(UNVALIDATED_REDIRECT_MARK).build();
+  VulnerabilityType WEAK_RANDOMNESS = type(VulnerabilityTypes.WEAK_RANDOMNESS).build();
 
   VulnerabilityType XPATH_INJECTION =
-      new VulnerabilityTypeImpl(
-          VulnerabilityTypes.XPATH_INJECTION, VulnerabilityMarks.XPATH_INJECTION_MARK);
+      type(VulnerabilityTypes.XPATH_INJECTION).mark(XPATH_INJECTION_MARK).build();
 
   VulnerabilityType TRUST_BOUNDARY_VIOLATION =
-      new VulnerabilityTypeImpl(
-          VulnerabilityTypes.TRUST_BOUNDARY_VIOLATION, VulnerabilityMarks.TRUST_BOUNDARY_VIOLATION);
+      type(VulnerabilityTypes.TRUST_BOUNDARY_VIOLATION).mark(TRUST_BOUNDARY_VIOLATION_MARK).build();
 
-  VulnerabilityType XSS =
-      new VulnerabilityTypeImpl(VulnerabilityTypes.XSS, VulnerabilityMarks.XSS_MARK);
+  VulnerabilityType XSS = type(VulnerabilityTypes.XSS).mark(XSS_MARK).build();
 
   VulnerabilityType HEADER_INJECTION =
-      new VulnerabilityTypeImpl(
-          VulnerabilityTypes.HEADER_INJECTION, VulnerabilityMarks.HEADER_INJECTION_MARK);
+      type(VulnerabilityTypes.HEADER_INJECTION).mark(HEADER_INJECTION_MARK).build();
 
-  VulnerabilityType STACKTRACE_LEAK = new VulnerabilityTypeImpl(VulnerabilityTypes.STACKTRACE_LEAK);
+  VulnerabilityType STACKTRACE_LEAK = type(VulnerabilityTypes.STACKTRACE_LEAK).build();
 
-  VulnerabilityType VERB_TAMPERING = new VulnerabilityTypeImpl(VulnerabilityTypes.VERB_TAMPERING);
+  VulnerabilityType VERB_TAMPERING = type(VulnerabilityTypes.VERB_TAMPERING).build();
 
   VulnerabilityType ADMIN_CONSOLE_ACTIVE =
-      new ServiceVulnerabilityType(VulnerabilityTypes.ADMIN_CONSOLE_ACTIVE, false);
+      type(VulnerabilityTypes.ADMIN_CONSOLE_ACTIVE)
+          .deduplicable(false)
+          .hash(VulnerabilityType::serviceHash)
+          .build();
 
   VulnerabilityType DEFAULT_HTML_ESCAPE_INVALID =
-      new VulnerabilityTypeImpl(VulnerabilityTypes.DEFAULT_HTML_ESCAPE_INVALID);
+      type(VulnerabilityTypes.DEFAULT_HTML_ESCAPE_INVALID).build();
 
-  VulnerabilityType SESSION_TIMEOUT = new VulnerabilityTypeImpl(VulnerabilityTypes.SESSION_TIMEOUT);
+  VulnerabilityType SESSION_TIMEOUT = type(VulnerabilityTypes.SESSION_TIMEOUT).build();
 
   VulnerabilityType DIRECTORY_LISTING_LEAK =
-      new VulnerabilityTypeImpl(VulnerabilityTypes.DIRECTORY_LISTING_LEAK);
-  VulnerabilityType INSECURE_JSP_LAYOUT =
-      new VulnerabilityTypeImpl(VulnerabilityTypes.INSECURE_JSP_LAYOUT);
+      type(VulnerabilityTypes.DIRECTORY_LISTING_LEAK).build();
+  VulnerabilityType INSECURE_JSP_LAYOUT = type(VulnerabilityTypes.INSECURE_JSP_LAYOUT).build();
 
-  VulnerabilityType HARDCODED_SECRET =
-      new VulnerabilityTypeImpl(VulnerabilityTypes.HARDCODED_SECRET);
+  VulnerabilityType HARDCODED_SECRET = type(VulnerabilityTypes.HARDCODED_SECRET).build();
 
   VulnerabilityType INSECURE_AUTH_PROTOCOL =
-      new VulnerabilityTypeImpl(VulnerabilityTypes.INSECURE_AUTH_PROTOCOL);
+      type(VulnerabilityTypes.INSECURE_AUTH_PROTOCOL).hash(VulnerabilityType::evidenceHash).build();
 
   VulnerabilityType REFLECTION_INJECTION =
-      new VulnerabilityTypeImpl(
-          VulnerabilityTypes.REFLECTION_INJECTION, VulnerabilityMarks.REFLECTION_INJECTION_MARK);
+      type(VulnerabilityTypes.REFLECTION_INJECTION).mark(REFLECTION_INJECTION_MARK).build();
 
   VulnerabilityType SESSION_REWRITING =
-      new ServiceVulnerabilityType(VulnerabilityTypes.SESSION_REWRITING, false);
+      type(VulnerabilityTypes.SESSION_REWRITING)
+          .deduplicable(false)
+          .hash(VulnerabilityType::serviceHash)
+          .build();
 
   VulnerabilityType DEFAULT_APP_DEPLOYED =
-      new ServiceVulnerabilityType(VulnerabilityTypes.DEFAULT_APP_DEPLOYED, false);
+      type(VulnerabilityTypes.DEFAULT_APP_DEPLOYED)
+          .deduplicable(false)
+          .hash(VulnerabilityType::serviceHash)
+          .build();
 
   String name();
 
@@ -104,6 +115,10 @@ public interface VulnerabilityType {
   /** A flag to indicate if the vulnerability is deduplicable. */
   boolean isDeduplicable();
 
+  static Builder type(final byte type) {
+    return new Builder(type);
+  }
+
   class VulnerabilityTypeImpl implements VulnerabilityType {
 
     private final byte type;
@@ -114,24 +129,19 @@ class VulnerabilityTypeImpl implements VulnerabilityType {
 
     private final boolean deduplicable;
 
-    public VulnerabilityTypeImpl(final byte type, final int... marks) {
-      this(type, ' ', marks);
-    }
-
-    public VulnerabilityTypeImpl(final byte type, boolean deduplicable, final int... marks) {
-      this(type, ' ', deduplicable, marks);
-    }
-
-    public VulnerabilityTypeImpl(final byte type, final char separator, final int... marks) {
-      this(type, separator, true, marks);
-    }
+    private final BiFunction<VulnerabilityType, Vulnerability, Long> hash;
 
     public VulnerabilityTypeImpl(
-        final byte type, final char separator, final boolean deduplicable, final int... marks) {
+        final byte type,
+        final char separator,
+        final int mark,
+        final boolean deduplicable,
+        final BiFunction<VulnerabilityType, Vulnerability, Long> hash) {
       this.type = type;
       this.separator = separator;
-      mark = computeMarks(marks);
+      this.mark = mark;
       this.deduplicable = deduplicable;
+      this.hash = hash;
     }
 
     @Override
@@ -150,89 +160,86 @@ public char separator() {
     }
 
     @Override
-    public long calculateHash(@Nonnull final Vulnerability vulnerability) {
-      CRC32 crc = new CRC32();
-      update(crc, name());
-      final Location location = vulnerability.getLocation();
-      if (location != null) {
-        crc.update(location.getLine());
-        if (location.getPath() != null) {
-          update(crc, location.getPath());
-        }
-        if (location.getLine() <= -1 && location.getMethod() != null) {
-          update(crc, location.getMethod());
-        }
-      }
-      return crc.getValue();
+    public long calculateHash(@Nonnull Vulnerability vulnerability) {
+      return hash.apply(this, vulnerability);
     }
 
     @Override
     public boolean isDeduplicable() {
       return deduplicable;
     }
+  }
 
-    protected void update(final CRC32 crc, final String value) {
-      final byte[] bytes = value.getBytes(StandardCharsets.UTF_8);
-      crc.update(bytes, 0, bytes.length);
+  class Builder {
+    private final byte type;
+    private char separator = ' ';
+    private int mark = NOT_MARKED;
+    private boolean deduplicable = true;
+    private BiFunction<VulnerabilityType, Vulnerability, Long> hash =
+        VulnerabilityType::fileAndLineHash;
+
+    public Builder(byte type) {
+      this.type = type;
     }
 
-    private static int computeMarks(final int... marks) {
-      int result = NOT_MARKED;
-      for (final int mark : marks) {
-        result |= mark;
-      }
-      return result;
+    public Builder separator(final char separator) {
+      this.separator = separator;
+      return this;
     }
-  }
 
-  class HeaderVulnerabilityType extends VulnerabilityTypeImpl {
-    public HeaderVulnerabilityType(byte type, int... marks) {
-      super(type, marks);
+    public Builder mark(final int mark) {
+      this.mark = mark;
+      return this;
     }
 
-    @Override
-    public long calculateHash(@Nonnull final Vulnerability vulnerability) {
-      CRC32 crc = new CRC32();
-      update(crc, name());
-      String serviceName = vulnerability.getLocation().getServiceName();
-      if (serviceName != null) {
-        update(crc, serviceName);
-      }
-      return crc.getValue();
+    public Builder deduplicable(final boolean deduplicable) {
+      this.deduplicable = deduplicable;
+      return this;
     }
-  }
 
-  class CookieVulnerabilityType extends VulnerabilityTypeImpl {
-    public CookieVulnerabilityType(byte type, int... marks) {
-      super(type, marks);
+    public Builder hash(final BiFunction<VulnerabilityType, Vulnerability, Long> hash) {
+      this.hash = hash;
+      return this;
     }
 
-    @Override
-    public long calculateHash(@Nonnull final Vulnerability vulnerability) {
-      CRC32 crc = new CRC32();
-      update(crc, name());
-      final Evidence evidence = vulnerability.getEvidence();
-      if (evidence != null) {
-        update(crc, evidence.getValue());
+    public VulnerabilityType build() {
+      return new VulnerabilityTypeImpl(type, separator, mark, deduplicable, hash);
+    }
+  }
+
+  static long fileAndLineHash(final VulnerabilityType type, final Vulnerability vulnerability) {
+    CRC32 crc = new CRC32();
+    update(crc, type.name());
+    final Location location = vulnerability.getLocation();
+    if (location != null) {
+      crc.update(location.getLine());
+      if (location.getPath() != null) {
+        update(crc, location.getPath());
+      }
+      if (location.getLine() <= -1 && location.getMethod() != null) {
+        update(crc, location.getMethod());
       }
-      return crc.getValue();
     }
+    return crc.getValue();
   }
 
-  class ServiceVulnerabilityType extends VulnerabilityTypeImpl {
-    public ServiceVulnerabilityType(byte type, boolean deduplicable, int... marks) {
-      super(type, deduplicable, marks);
+  static long evidenceHash(final VulnerabilityType type, final Vulnerability vulnerability) {
+    CRC32 crc = new CRC32();
+    update(crc, type.name());
+    final Evidence evidence = vulnerability.getEvidence();
+    if (evidence != null) {
+      update(crc, evidence.getValue());
     }
+    return crc.getValue();
+  }
 
-    @Override
-    public long calculateHash(@Nonnull final Vulnerability vulnerability) {
-      CRC32 crc = new CRC32();
-      update(crc, name());
-      String serviceName = vulnerability.getLocation().getServiceName();
-      if (serviceName != null) {
-        update(crc, serviceName);
-      }
-      return crc.getValue();
+  static long serviceHash(final VulnerabilityType type, final Vulnerability vulnerability) {
+    CRC32 crc = new CRC32();
+    update(crc, type.name());
+    final String serviceName = vulnerability.getLocation().getServiceName();
+    if (serviceName != null) {
+      update(crc, serviceName);
     }
+    return crc.getValue();
   }
 }

dd-java-agent/agent-iast/src/main/java/com/datadog/iast/util/CRCUtils.java

@@ -0,0 +1,25 @@
+package com.datadog.iast.util;
+
+import java.nio.charset.Charset;
+import java.nio.charset.StandardCharsets;
+import java.util.zip.CRC32;
+
+public abstract class CRCUtils {
+
+  private static final Charset DEFAULT_CHARSET = StandardCharsets.UTF_8;
+
+  private CRCUtils() {}
+
+  public static void update(final CRC32 crc, final String value) {
+    update(crc, value, DEFAULT_CHARSET);
+  }
+
+  public static void update(final CRC32 crc, final String value, final Charset charset) {
+    final byte[] bytes = value.getBytes(charset);
+    update(crc, bytes);
+  }
+
+  public static void update(final CRC32 crc, final byte[] bytes) {
+    crc.update(bytes, 0, bytes.length);
+  }
+}