Add extended data collection action support

src/main/c/java_call.c

@@ -53,7 +53,7 @@ bool java_meth_init_checked(JNIEnv *env, struct j_method *jmeth,
         goto end;
     }
 
-    *jmeth = (struct j_method){
+    *jmeth = (struct j_method) {
             .type = type != JMETHOD_VIRTUAL_RETRIEVE_CLASS ? type
                                                            : JMETHOD_VIRTUAL,
             .class_glob = class_global,

src/main/c/waf_jni.c

@@ -1877,7 +1877,7 @@ static struct _limits _fetch_limits_checked(JNIEnv *env, jobject limits_obj)
 
     return l;
 error:
-    return (struct _limits){0};
+    return (struct _limits) {0};
 }
 
 static bool _get_time_checked(JNIEnv *env, struct timespec *time)
@@ -2081,7 +2081,7 @@ static bool _convert_ddwaf_config_checked(JNIEnv *env, jobject jconfig,
     }
     if (!JNI(IsSameObject, key_regex_jstr, NULL)) {
         key_regex = java_to_utf8_checked(env, (jstring) key_regex_jstr,
-                                         &(size_t){0});
+                                         &(size_t) {0});
         if (!key_regex) {
             return false;
         }
@@ -2097,7 +2097,7 @@ static bool _convert_ddwaf_config_checked(JNIEnv *env, jobject jconfig,
     }
     if (!JNI(IsSameObject, value_regex_jstr, NULL)) {
         value_regex = java_to_utf8_checked(env, (jstring) value_regex_jstr,
-                                           &(size_t){0});
+                                           &(size_t) {0});
         if (!value_regex) {
             free(key_regex);
             return false;
@@ -2108,7 +2108,7 @@ static bool _convert_ddwaf_config_checked(JNIEnv *env, jobject jconfig,
         }
     }
 
-    *out_config = (ddwaf_config){
+    *out_config = (ddwaf_config) {
             // disable these checks. We also have our own
             // limits given at rule run time
             .limits =

src/test/groovy/com/datadog/ddwaf/BasicTests.groovy

@@ -284,7 +284,7 @@ class BasicTests implements WafTrait {
     handle = builder.buildWafHandleInstance()
     context = new WafContext(handle)
     assertThat handle.knownActions as List,
-      containsInAnyOrder('block_request', 'generate_stack', 'redirect_request')
+      containsInAnyOrder('block_request', 'generate_stack', 'redirect_request', 'extended_data_collection')
   }
 
   @Test
@@ -513,7 +513,8 @@ class BasicTests implements WafTrait {
             operator  : 'ip_match',
             parameters: [
               data  : 'suspicious_ips_data_id',
-              inputs: [[address: 'http.client_ip']]]
+              inputs: [[address: 'http.client_ip']]
+            ]
           ]
         ],
       ]
@@ -639,5 +640,81 @@ class BasicTests implements WafTrait {
       ['server.request.headers.no_cookies': ['user-agent': 'foobar']], limits, metrics)
     assertThat res.result, is(Waf.Result.MATCH)
   }
-}
 
+  @Test
+  void 'test extended_data_collection action'() {
+    def ruleSet = slurper.parseText(JsonOutput.toJson(ARACHNI_ATOM_BLOCK))
+    ruleSet['rules'][0]['on_match'] = ['extended_data']
+
+    wafDiagnostics = builder.addOrUpdateConfig('test', ruleSet)
+    handle = builder.buildWafHandleInstance()
+    context = new WafContext(handle)
+    final params = ['server.request.headers.no_cookies': ['user-agent': 'Arachni/v1']]
+    ResultWithData res = context.run(params, limits, metrics)
+    assertThat res.result, is(Waf.Result.MATCH)
+    assertThat res.actions.keySet(), hasItem('extended_data_collection')
+  }
+
+  @Test
+  void 'test extended_data_collection action with parameters'() {
+    def ruleSet = slurper.parseText(JsonOutput.toJson(ARACHNI_ATOM_BLOCK))
+    ruleSet.putAt('actions', [
+      [
+        id: 'extended_data',
+        parameters: [
+          headers_redaction: true,
+          max_collected_headers: 10
+        ],
+        type: 'extended_data_collection'
+      ]
+    ])
+    ruleSet['rules'][0]['on_match'] = ['extended_data']
+
+    wafDiagnostics = builder.addOrUpdateConfig('test', ruleSet)
+    handle = builder.buildWafHandleInstance()
+    context = new WafContext(handle)
+    final params = ['server.request.headers.no_cookies': ['user-agent': 'Arachni/v1']]
+    ResultWithData res = context.run(params, limits, metrics)
+    assertThat res.result, is(Waf.Result.MATCH)
+    assertThat res.actions.keySet(), hasItem('extended_data_collection')
+    assertThat res.actions.get('extended_data_collection').headers_redaction, is('true')
+    assertThat res.actions.get('extended_data_collection').max_collected_headers, is('10')
+  }
+
+  @Test
+  void 'test multiple actions with extended_data_collection'() {
+    def ruleSet = ARACHNI_ATOM_V2_1
+    ruleSet.putAt('actions', [
+      [
+        id: 'extended_data',
+        parameters: [
+          headers_redaction: true,
+          max_collected_headers: 10
+        ],
+        type: 'extended_data_collection'
+      ]
+    ])
+    ruleSet['rules'][0]['on_match'] = ['block', 'stack_trace', 'extended_data']
+
+    wafDiagnostics = builder.addOrUpdateConfig('test', ruleSet)
+    handle = builder.buildWafHandleInstance()
+    context = new WafContext(handle)
+    final params = ['server.request.headers.no_cookies': ['user-agent': 'Arachni/v1']]
+    ResultWithData res = context.run(params, limits, metrics)
+    assertThat res.result, is(Waf.Result.MATCH)
+    assertThat res.actions.size(), is(3)
+
+    // block action
+    assertThat res.actions.keySet(), hasItem('block_request')
+    assertThat res.actions.get('block_request').type, is('auto')
+    assertThat res.actions.get('block_request').status_code, is('403')
+    assertThat res.actions.get('block_request').grpc_status_code, is('10')
+
+    // stack_trace action
+    assertThat res.actions.keySet(), hasItem('generate_stack')
+    assertThat res.actions.get('generate_stack').stack_id, is(notNullValue())
+
+    // extended_data_collection action
+    assertThat res.actions.keySet(), hasItem('extended_data_collection')
+  }
+}

src/test/groovy/com/datadog/ddwaf/WafTrait.groovy

@@ -108,6 +108,14 @@ trait WafTrait extends JNITrait {
                 "status_code": 400,
                 "location": "https://example3.com/"
               }
+            },
+            {
+              "id": "extended_data",
+              "type": "extended_data_collection",
+              "parameters": {
+                "headers_redaction": true,
+                "max_collected_headers": 5
+              }
             }
           ],
           "rules": [
@@ -158,7 +166,7 @@ trait WafTrait extends JNITrait {
                         "operator": "match_regex"
                     }
                 ],
-                "on_match": ["stack_trace", "redirect2"]
+                "on_match": ["stack_trace", "redirect2", "extended_data"]
             }
           ]
         }''')
@@ -249,4 +257,3 @@ trait WafTrait extends JNITrait {
     ] as Map<String, Object>, limits, metrics)
   }
 }
-