11package com .datadog .iast .model ;
22
3+ import static com .datadog .iast .util .CRCUtils .update ;
4+ import static datadog .trace .api .iast .VulnerabilityMarks .COMMAND_INJECTION_MARK ;
5+ import static datadog .trace .api .iast .VulnerabilityMarks .HEADER_INJECTION_MARK ;
6+ import static datadog .trace .api .iast .VulnerabilityMarks .LDAP_INJECTION_MARK ;
37import static datadog .trace .api .iast .VulnerabilityMarks .NOT_MARKED ;
8+ import static datadog .trace .api .iast .VulnerabilityMarks .PATH_TRAVERSAL_MARK ;
9+ import static datadog .trace .api .iast .VulnerabilityMarks .REFLECTION_INJECTION_MARK ;
10+ import static datadog .trace .api .iast .VulnerabilityMarks .SQL_INJECTION_MARK ;
11+ import static datadog .trace .api .iast .VulnerabilityMarks .SSRF_MARK ;
12+ import static datadog .trace .api .iast .VulnerabilityMarks .TRUST_BOUNDARY_VIOLATION_MARK ;
13+ import static datadog .trace .api .iast .VulnerabilityMarks .UNVALIDATED_REDIRECT_MARK ;
14+ import static datadog .trace .api .iast .VulnerabilityMarks .XPATH_INJECTION_MARK ;
15+ import static datadog .trace .api .iast .VulnerabilityMarks .XSS_MARK ;
416
5- import datadog .trace .api .iast .VulnerabilityMarks ;
617import datadog .trace .api .iast .VulnerabilityTypes ;
718import java .io .File ;
8- import java .nio . charset . StandardCharsets ;
19+ import java .util . function . BiFunction ;
920import java .util .zip .CRC32 ;
1021import javax .annotation .Nonnull ;
1122
1223public interface VulnerabilityType {
13- VulnerabilityType WEAK_CIPHER = new VulnerabilityTypeImpl (VulnerabilityTypes .WEAK_CIPHER );
14- VulnerabilityType WEAK_HASH = new VulnerabilityTypeImpl (VulnerabilityTypes .WEAK_HASH );
24+
25+ VulnerabilityType WEAK_CIPHER = type (VulnerabilityTypes .WEAK_CIPHER ).build ();
26+ VulnerabilityType WEAK_HASH = type (VulnerabilityTypes .WEAK_HASH ).build ();
1527 VulnerabilityType INSECURE_COOKIE =
16- new CookieVulnerabilityType (VulnerabilityTypes .INSECURE_COOKIE );
28+ type (VulnerabilityTypes .INSECURE_COOKIE ). hash ( VulnerabilityType :: evidenceHash ). build ( );
1729 VulnerabilityType NO_HTTPONLY_COOKIE =
18- new CookieVulnerabilityType (VulnerabilityTypes .NO_HTTPONLY_COOKIE );
30+ type (VulnerabilityTypes .NO_HTTPONLY_COOKIE ). hash ( VulnerabilityType :: evidenceHash ). build ( );
1931 VulnerabilityType HSTS_HEADER_MISSING =
20- new HeaderVulnerabilityType (VulnerabilityTypes .HSTS_HEADER_MISSING );
32+ type (VulnerabilityTypes .HSTS_HEADER_MISSING ). hash ( VulnerabilityType :: serviceHash ). build ( );
2133 VulnerabilityType XCONTENTTYPE_HEADER_MISSING =
22- new HeaderVulnerabilityType (VulnerabilityTypes .XCONTENTTYPE_HEADER_MISSING );
34+ type (VulnerabilityTypes .XCONTENTTYPE_HEADER_MISSING )
35+ .hash (VulnerabilityType ::serviceHash )
36+ .build ();
2337 VulnerabilityType NO_SAMESITE_COOKIE =
24- new CookieVulnerabilityType (VulnerabilityTypes .NO_SAMESITE_COOKIE );
38+ type (VulnerabilityTypes .NO_SAMESITE_COOKIE ). hash ( VulnerabilityType :: evidenceHash ). build ( );
2539
2640 VulnerabilityType SQL_INJECTION =
27- new VulnerabilityTypeImpl (
28- VulnerabilityTypes .SQL_INJECTION , VulnerabilityMarks .SQL_INJECTION_MARK );
41+ type (VulnerabilityTypes .SQL_INJECTION ).mark (SQL_INJECTION_MARK ).build ();
2942 VulnerabilityType COMMAND_INJECTION =
30- new VulnerabilityTypeImpl (
31- VulnerabilityTypes .COMMAND_INJECTION , VulnerabilityMarks .COMMAND_INJECTION_MARK );
43+ type (VulnerabilityTypes .COMMAND_INJECTION ).mark (COMMAND_INJECTION_MARK ).build ();
3244 VulnerabilityType PATH_TRAVERSAL =
33- new VulnerabilityTypeImpl (
34- VulnerabilityTypes . PATH_TRAVERSAL ,
35- File . separatorChar ,
36- VulnerabilityMarks . PATH_TRAVERSAL_MARK );
45+ type ( VulnerabilityTypes . PATH_TRAVERSAL )
46+ . separator ( File . separatorChar )
47+ . mark ( PATH_TRAVERSAL_MARK )
48+ . build ( );
3749 VulnerabilityType LDAP_INJECTION =
38- new VulnerabilityTypeImpl (
39- VulnerabilityTypes .LDAP_INJECTION , VulnerabilityMarks .LDAP_INJECTION_MARK );
40- VulnerabilityType SSRF =
41- new VulnerabilityTypeImpl (VulnerabilityTypes .SSRF , VulnerabilityMarks .SSRF_MARK );
50+ type (VulnerabilityTypes .LDAP_INJECTION ).mark (LDAP_INJECTION_MARK ).build ();
51+ VulnerabilityType SSRF = type (VulnerabilityTypes .SSRF ).mark (SSRF_MARK ).build ();
4252 VulnerabilityType UNVALIDATED_REDIRECT =
43- new VulnerabilityTypeImpl (
44- VulnerabilityTypes .UNVALIDATED_REDIRECT , VulnerabilityMarks .UNVALIDATED_REDIRECT_MARK );
45- VulnerabilityType WEAK_RANDOMNESS = new VulnerabilityTypeImpl (VulnerabilityTypes .WEAK_RANDOMNESS );
53+ type (VulnerabilityTypes .UNVALIDATED_REDIRECT ).mark (UNVALIDATED_REDIRECT_MARK ).build ();
54+ VulnerabilityType WEAK_RANDOMNESS = type (VulnerabilityTypes .WEAK_RANDOMNESS ).build ();
4655
4756 VulnerabilityType XPATH_INJECTION =
48- new VulnerabilityTypeImpl (
49- VulnerabilityTypes .XPATH_INJECTION , VulnerabilityMarks .XPATH_INJECTION_MARK );
57+ type (VulnerabilityTypes .XPATH_INJECTION ).mark (XPATH_INJECTION_MARK ).build ();
5058
5159 VulnerabilityType TRUST_BOUNDARY_VIOLATION =
52- new VulnerabilityTypeImpl (
53- VulnerabilityTypes .TRUST_BOUNDARY_VIOLATION , VulnerabilityMarks .TRUST_BOUNDARY_VIOLATION );
60+ type (VulnerabilityTypes .TRUST_BOUNDARY_VIOLATION ).mark (TRUST_BOUNDARY_VIOLATION_MARK ).build ();
5461
55- VulnerabilityType XSS =
56- new VulnerabilityTypeImpl (VulnerabilityTypes .XSS , VulnerabilityMarks .XSS_MARK );
62+ VulnerabilityType XSS = type (VulnerabilityTypes .XSS ).mark (XSS_MARK ).build ();
5763
5864 VulnerabilityType HEADER_INJECTION =
59- new VulnerabilityTypeImpl (
60- VulnerabilityTypes .HEADER_INJECTION , VulnerabilityMarks .HEADER_INJECTION_MARK );
65+ type (VulnerabilityTypes .HEADER_INJECTION ).mark (HEADER_INJECTION_MARK ).build ();
6166
62- VulnerabilityType STACKTRACE_LEAK = new VulnerabilityTypeImpl (VulnerabilityTypes .STACKTRACE_LEAK );
67+ VulnerabilityType STACKTRACE_LEAK = type (VulnerabilityTypes .STACKTRACE_LEAK ). build ( );
6368
64- VulnerabilityType VERB_TAMPERING = new VulnerabilityTypeImpl (VulnerabilityTypes .VERB_TAMPERING );
69+ VulnerabilityType VERB_TAMPERING = type (VulnerabilityTypes .VERB_TAMPERING ). build ( );
6570
6671 VulnerabilityType ADMIN_CONSOLE_ACTIVE =
67- new ServiceVulnerabilityType (VulnerabilityTypes .ADMIN_CONSOLE_ACTIVE , false );
72+ type (VulnerabilityTypes .ADMIN_CONSOLE_ACTIVE )
73+ .deduplicable (false )
74+ .hash (VulnerabilityType ::serviceHash )
75+ .build ();
6876
6977 VulnerabilityType DEFAULT_HTML_ESCAPE_INVALID =
70- new VulnerabilityTypeImpl (VulnerabilityTypes .DEFAULT_HTML_ESCAPE_INVALID );
78+ type (VulnerabilityTypes .DEFAULT_HTML_ESCAPE_INVALID ). build ( );
7179
72- VulnerabilityType SESSION_TIMEOUT = new VulnerabilityTypeImpl (VulnerabilityTypes .SESSION_TIMEOUT );
80+ VulnerabilityType SESSION_TIMEOUT = type (VulnerabilityTypes .SESSION_TIMEOUT ). build ( );
7381
7482 VulnerabilityType DIRECTORY_LISTING_LEAK =
75- new VulnerabilityTypeImpl (VulnerabilityTypes .DIRECTORY_LISTING_LEAK );
76- VulnerabilityType INSECURE_JSP_LAYOUT =
77- new VulnerabilityTypeImpl (VulnerabilityTypes .INSECURE_JSP_LAYOUT );
83+ type (VulnerabilityTypes .DIRECTORY_LISTING_LEAK ).build ();
84+ VulnerabilityType INSECURE_JSP_LAYOUT = type (VulnerabilityTypes .INSECURE_JSP_LAYOUT ).build ();
7885
79- VulnerabilityType HARDCODED_SECRET =
80- new VulnerabilityTypeImpl (VulnerabilityTypes .HARDCODED_SECRET );
86+ VulnerabilityType HARDCODED_SECRET = type (VulnerabilityTypes .HARDCODED_SECRET ).build ();
8187
8288 VulnerabilityType INSECURE_AUTH_PROTOCOL =
83- new VulnerabilityTypeImpl (VulnerabilityTypes .INSECURE_AUTH_PROTOCOL );
89+ type (VulnerabilityTypes .INSECURE_AUTH_PROTOCOL ). hash ( VulnerabilityType :: evidenceHash ). build ( );
8490
8591 VulnerabilityType REFLECTION_INJECTION =
86- new VulnerabilityTypeImpl (
87- VulnerabilityTypes .REFLECTION_INJECTION , VulnerabilityMarks .REFLECTION_INJECTION_MARK );
92+ type (VulnerabilityTypes .REFLECTION_INJECTION ).mark (REFLECTION_INJECTION_MARK ).build ();
8893
8994 VulnerabilityType SESSION_REWRITING =
90- new ServiceVulnerabilityType (VulnerabilityTypes .SESSION_REWRITING , false );
95+ type (VulnerabilityTypes .SESSION_REWRITING )
96+ .deduplicable (false )
97+ .hash (VulnerabilityType ::serviceHash )
98+ .build ();
9199
92100 VulnerabilityType DEFAULT_APP_DEPLOYED =
93- new ServiceVulnerabilityType (VulnerabilityTypes .DEFAULT_APP_DEPLOYED , false );
101+ type (VulnerabilityTypes .DEFAULT_APP_DEPLOYED )
102+ .deduplicable (false )
103+ .hash (VulnerabilityType ::serviceHash )
104+ .build ();
94105
95106 String name ();
96107
@@ -104,6 +115,10 @@ public interface VulnerabilityType {
104115 /** A flag to indicate if the vulnerability is deduplicable. */
105116 boolean isDeduplicable ();
106117
118+ static Builder type (final byte type ) {
119+ return new Builder (type );
120+ }
121+
107122 class VulnerabilityTypeImpl implements VulnerabilityType {
108123
109124 private final byte type ;
@@ -114,24 +129,19 @@ class VulnerabilityTypeImpl implements VulnerabilityType {
114129
115130 private final boolean deduplicable ;
116131
117- public VulnerabilityTypeImpl (final byte type , final int ... marks ) {
118- this (type , ' ' , marks );
119- }
120-
121- public VulnerabilityTypeImpl (final byte type , boolean deduplicable , final int ... marks ) {
122- this (type , ' ' , deduplicable , marks );
123- }
124-
125- public VulnerabilityTypeImpl (final byte type , final char separator , final int ... marks ) {
126- this (type , separator , true , marks );
127- }
132+ private final BiFunction <VulnerabilityType , Vulnerability , Long > hash ;
128133
129134 public VulnerabilityTypeImpl (
130- final byte type , final char separator , final boolean deduplicable , final int ... marks ) {
135+ final byte type ,
136+ final char separator ,
137+ final int mark ,
138+ final boolean deduplicable ,
139+ final BiFunction <VulnerabilityType , Vulnerability , Long > hash ) {
131140 this .type = type ;
132141 this .separator = separator ;
133- mark = computeMarks ( marks ) ;
142+ this . mark = mark ;
134143 this .deduplicable = deduplicable ;
144+ this .hash = hash ;
135145 }
136146
137147 @ Override
@@ -150,89 +160,86 @@ public char separator() {
150160 }
151161
152162 @ Override
153- public long calculateHash (@ Nonnull final Vulnerability vulnerability ) {
154- CRC32 crc = new CRC32 ();
155- update (crc , name ());
156- final Location location = vulnerability .getLocation ();
157- if (location != null ) {
158- crc .update (location .getLine ());
159- if (location .getPath () != null ) {
160- update (crc , location .getPath ());
161- }
162- if (location .getLine () <= -1 && location .getMethod () != null ) {
163- update (crc , location .getMethod ());
164- }
165- }
166- return crc .getValue ();
163+ public long calculateHash (@ Nonnull Vulnerability vulnerability ) {
164+ return hash .apply (this , vulnerability );
167165 }
168166
169167 @ Override
170168 public boolean isDeduplicable () {
171169 return deduplicable ;
172170 }
171+ }
173172
174- protected void update (final CRC32 crc , final String value ) {
175- final byte [] bytes = value .getBytes (StandardCharsets .UTF_8 );
176- crc .update (bytes , 0 , bytes .length );
173+ class Builder {
174+ private final byte type ;
175+ private char separator = ' ' ;
176+ private int mark = NOT_MARKED ;
177+ private boolean deduplicable = true ;
178+ private BiFunction <VulnerabilityType , Vulnerability , Long > hash =
179+ VulnerabilityType ::fileAndLineHash ;
180+
181+ public Builder (byte type ) {
182+ this .type = type ;
177183 }
178184
179- private static int computeMarks (final int ... marks ) {
180- int result = NOT_MARKED ;
181- for (final int mark : marks ) {
182- result |= mark ;
183- }
184- return result ;
185+ public Builder separator (final char separator ) {
186+ this .separator = separator ;
187+ return this ;
185188 }
186- }
187189
188- class HeaderVulnerabilityType extends VulnerabilityTypeImpl {
189- public HeaderVulnerabilityType ( byte type , int ... marks ) {
190- super ( type , marks ) ;
190+ public Builder mark ( final int mark ) {
191+ this . mark = mark ;
192+ return this ;
191193 }
192194
193- @ Override
194- public long calculateHash (@ Nonnull final Vulnerability vulnerability ) {
195- CRC32 crc = new CRC32 ();
196- update (crc , name ());
197- String serviceName = vulnerability .getLocation ().getServiceName ();
198- if (serviceName != null ) {
199- update (crc , serviceName );
200- }
201- return crc .getValue ();
195+ public Builder deduplicable (final boolean deduplicable ) {
196+ this .deduplicable = deduplicable ;
197+ return this ;
202198 }
203- }
204199
205- class CookieVulnerabilityType extends VulnerabilityTypeImpl {
206- public CookieVulnerabilityType ( byte type , int ... marks ) {
207- super ( type , marks ) ;
200+ public Builder hash ( final BiFunction < VulnerabilityType , Vulnerability , Long > hash ) {
201+ this . hash = hash ;
202+ return this ;
208203 }
209204
210- @ Override
211- public long calculateHash (@ Nonnull final Vulnerability vulnerability ) {
212- CRC32 crc = new CRC32 ();
213- update (crc , name ());
214- final Evidence evidence = vulnerability .getEvidence ();
215- if (evidence != null ) {
216- update (crc , evidence .getValue ());
205+ public VulnerabilityType build () {
206+ return new VulnerabilityTypeImpl (type , separator , mark , deduplicable , hash );
207+ }
208+ }
209+
210+ static long fileAndLineHash (final VulnerabilityType type , final Vulnerability vulnerability ) {
211+ CRC32 crc = new CRC32 ();
212+ update (crc , type .name ());
213+ final Location location = vulnerability .getLocation ();
214+ if (location != null ) {
215+ crc .update (location .getLine ());
216+ if (location .getPath () != null ) {
217+ update (crc , location .getPath ());
218+ }
219+ if (location .getLine () <= -1 && location .getMethod () != null ) {
220+ update (crc , location .getMethod ());
217221 }
218- return crc .getValue ();
219222 }
223+ return crc .getValue ();
220224 }
221225
222- class ServiceVulnerabilityType extends VulnerabilityTypeImpl {
223- public ServiceVulnerabilityType (byte type , boolean deduplicable , int ... marks ) {
224- super (type , deduplicable , marks );
226+ static long evidenceHash (final VulnerabilityType type , final Vulnerability vulnerability ) {
227+ CRC32 crc = new CRC32 ();
228+ update (crc , type .name ());
229+ final Evidence evidence = vulnerability .getEvidence ();
230+ if (evidence != null ) {
231+ update (crc , evidence .getValue ());
225232 }
233+ return crc .getValue ();
234+ }
226235
227- @ Override
228- public long calculateHash (@ Nonnull final Vulnerability vulnerability ) {
229- CRC32 crc = new CRC32 ();
230- update (crc , name ());
231- String serviceName = vulnerability .getLocation ().getServiceName ();
232- if (serviceName != null ) {
233- update (crc , serviceName );
234- }
235- return crc .getValue ();
236+ static long serviceHash (final VulnerabilityType type , final Vulnerability vulnerability ) {
237+ CRC32 crc = new CRC32 ();
238+ update (crc , type .name ());
239+ final String serviceName = vulnerability .getLocation ().getServiceName ();
240+ if (serviceName != null ) {
241+ update (crc , serviceName );
236242 }
243+ return crc .getValue ();
237244 }
238245}
0 commit comments