wip - fix

Author: jandro996

dd-java-agent/appsec/src/main/java/com/datadog/appsec/event/data/ObjectIntrospection.java

@@ -2,7 +2,6 @@
 
 import com.datadog.appsec.gateway.AppSecRequestContext;
 import datadog.trace.api.Platform;
-import datadog.trace.api.telemetry.WafMetricCollector;
 import java.lang.reflect.Array;
 import java.lang.reflect.Field;
 import java.lang.reflect.InvocationTargetException;
@@ -38,25 +37,6 @@ public final class ObjectIntrospection {
 
   private ObjectIntrospection() {}
 
-  /** Functional interface to receive truncation flags. */
-  @FunctionalInterface
-  public interface TruncationCallback {
-
-    /**
-     * Called when the object is converted.
-     *
-     * @param requestContext the request context
-     * @param stringTooLong true if a string was truncated
-     * @param listMapTooLarge true if a list or map was truncated
-     * @param objectTooDeep true if an object was too deep
-     */
-    void onTruncation(
-        AppSecRequestContext requestContext,
-        boolean stringTooLong,
-        boolean listMapTooLarge,
-        boolean objectTooDeep);
-  }
-
   /**
    * Converts arbitrary objects compatible with ddwaf_object. Possible types in the result are:
    *
@@ -86,16 +66,18 @@ void onTruncation(
    * @param requestContext the request context
    * @return the converted object
    */
-  public static Object convert(Object obj, AppSecRequestContext requestContext) {
+  public static ConversionResult<Object> convert(Object obj, AppSecRequestContext requestContext) {
     State state = new State(requestContext);
     Object converted = guardedConversion(obj, 0, state);
     if (state.stringTooLong || state.listMapTooLarge || state.objectTooDeep) {
       requestContext.setWafTruncated();
-      //TODO Enable when rebase with master
-//      WafMetricCollector.get()
-//          .wafInputTruncated(state.stringTooLong, state.listMapTooLarge, state.objectTooDeep);
+      // TODO Enable when rebase with master
+      //      WafMetricCollector.get()
+      //          .wafInputTruncated(state.stringTooLong, state.listMapTooLarge,
+      // state.objectTooDeep);
     }
-    return converted;
+    return new ConversionResult<>(
+        converted, state.stringTooLong, state.listMapTooLarge, state.objectTooDeep);
   }
 
   private static class State {
@@ -222,8 +204,8 @@ private static Object doConversion(Object obj, int depth, State state) {
     Map<String, Object> newMap = new HashMap<>();
     List<Field[]> allFields = new ArrayList<>();
     for (Class<?> classToLook = clazz;
-         classToLook != null && classToLook != Object.class;
-         classToLook = classToLook.getSuperclass()) {
+        classToLook != null && classToLook != Object.class;
+        classToLook = classToLook.getSuperclass()) {
       allFields.add(classToLook.getDeclaredFields());
     }
 
@@ -301,5 +283,42 @@ private static String checkStringLength(final String str, final State state) {
     }
     return str;
   }
-}
 
+  public static class ConversionResult<T> {
+    private final T value;
+    private final boolean truncatedByString;
+    private final boolean truncatedByCollection;
+    private final boolean truncatedByDepth;
+
+    public ConversionResult(
+        T value,
+        boolean truncatedByString,
+        boolean truncatedByCollection,
+        boolean truncatedByDepth) {
+      this.value = value;
+      this.truncatedByString = truncatedByString;
+      this.truncatedByCollection = truncatedByCollection;
+      this.truncatedByDepth = truncatedByDepth;
+    }
+
+    public T getValue() {
+      return value;
+    }
+
+    public boolean isStringTruncated() {
+      return truncatedByString;
+    }
+
+    public boolean isCollectionTruncated() {
+      return truncatedByCollection;
+    }
+
+    public boolean isDepthTruncated() {
+      return truncatedByDepth;
+    }
+
+    public boolean isAnyTruncated() {
+      return truncatedByString || truncatedByCollection || truncatedByDepth;
+    }
+  }
+}

dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/GatewayBridge.java

@@ -69,9 +69,9 @@ public class GatewayBridge {
 
   /** User tracking tags that will force the collection of request headers */
   private static final String[] USER_TRACKING_TAGS = {
-      "appsec.events.users.login.success.track",
-      "appsec.events.users.login.failure.track",
-      "appsec.events.users.signup.track"
+    "appsec.events.users.login.success.track",
+    "appsec.events.users.login.failure.track",
+    "appsec.events.users.signup.track"
   };
 
   private static final String USER_COLLECTION_MODE_TAG = "_dd.appsec.user.collection_mode";
@@ -473,7 +473,7 @@ private Flow<Void> onGrpcServerRequestMessage(RequestContext ctx_, Object obj) {
       if (subInfo == null || subInfo.isEmpty()) {
         return NoopFlow.INSTANCE;
       }
-      Object convObj = ObjectIntrospection.convert(obj, ctx);
+      Object convObj = ObjectIntrospection.convert(obj, ctx).getValue();
       DataBundle bundle =
           new SingletonDataBundle<>(KnownAddresses.GRPC_SERVER_REQUEST_MESSAGE, convObj);
       try {
@@ -573,17 +573,16 @@ private Flow<Void> onRequestBodyProcessed(RequestContext ctx_, Object obj) {
       if (subInfo == null || subInfo.isEmpty()) {
         return NoopFlow.INSTANCE;
       }
-      Object converted =  ObjectIntrospection.convert(obj, ctx)
-//      if (Config.get().isAppSecRaspCollectRequestBody()) {
-//        ctx.setProcessedRequestBody(pair.getLeft());
-//        Boolean limitsExceeded = pair.getRight();
-//        if (Boolean.TRUE.equals(limitsExceeded)) {
-//          ctx_.getTraceSegment().setTagTop("_dd.appsec.rasp.request_body_size.exceeded", true);
-//        }
-//      }
+      ObjectIntrospection.ConversionResult<Object> converted =
+          ObjectIntrospection.convert(obj, ctx);
+      if (Config.get().isAppSecRaspCollectRequestBody()) {
+        ctx.setProcessedRequestBody(converted.getValue());
+        if (converted.isAnyTruncated()) {
+          ctx_.getTraceSegment().setTagTop("_dd.appsec.rasp.request_body_size.exceeded", true);
+        }
+      }
       DataBundle bundle =
-          new SingletonDataBundle<>(
-              KnownAddresses.REQUEST_BODY_OBJECT,converted);
+          new SingletonDataBundle<>(KnownAddresses.REQUEST_BODY_OBJECT, converted.getValue());
       try {
         GatewayContext gwCtx = new GatewayContext(false);
         return producerService.publishDataEvent(subInfo, ctx, bundle, gwCtx);
@@ -757,7 +756,7 @@ private NoopFlow onRequestEnded(RequestContext ctx_, IGSpanInfo spanInfo) {
               ctx.isWafRequestBlockFailure(), // blockFailure,
               ctx.isWafRateLimited(), // rateLimited,
               ctx.isWafTruncated() // inputTruncated
-          );
+              );
     }
 
     ctx.close();
@@ -1142,4 +1141,3 @@ static Collection<datadog.trace.api.gateway.EventType<?>> additionalIGEventTypes
     }
   }
 }
-