Parse JSON

Author: piochelepiotr

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/messaging/DatadogAttributeParser.java

@@ -26,6 +26,7 @@ public static void forEachProperty(AgentPropagation.KeyClassifier classifier, St
         acceptJsonProperty(classifier, json, "x-datadog-sampling-priority");
       }
       if (Config.get().isDataStreamsEnabled()) {
+        System.out.println("DSM ENABLED!!!: " + json);
         acceptJsonProperty(classifier, json, "dd-pathway-ctx-base64");
       }
     } catch (Exception e) {

dd-java-agent/instrumentation/aws-java-sqs-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sqs/MessageExtractAdapter.java

@@ -2,8 +2,14 @@
 
 import com.amazonaws.services.sqs.model.Message;
 import com.amazonaws.services.sqs.model.MessageAttributeValue;
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.ObjectMapper;
 import datadog.trace.bootstrap.instrumentation.api.AgentPropagation;
 import datadog.trace.bootstrap.instrumentation.messaging.DatadogAttributeParser;
+
+import java.nio.ByteBuffer;
+import java.io.IOException;
+import java.util.Base64;
 import java.util.Map;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -28,6 +34,32 @@ public void forEachKey(Message carrier, AgentPropagation.KeyClassifier classifie
       } else if ("Binary".equals(datadog.getDataType())) {
         DatadogAttributeParser.forEachProperty(classifier, datadog.getBinaryValue());
       }
+    } else {
+      try {
+        this.forEachKeyInBody(carrier.getBody(), classifier);
+      } catch (IOException e) {
+        log.warn("Error extracting Datadog context from SQS message body", e);
+      }
+    }
+  }
+
+  public void forEachKeyInBody(String body, AgentPropagation.KeyClassifier classifier) throws IOException {
+    ObjectMapper objectMapper = new ObjectMapper();
+
+    // Parse the JSON string into a JsonNode
+    JsonNode rootNode = objectMapper.readTree(body);
+
+    // Navigate to MessageAttributes._datadog
+    JsonNode messageAttributes = rootNode.path("MessageAttributes").path("_datadog");
+
+    // Extract Value and Type
+    String value = messageAttributes.path("Value").asText();
+    String type = messageAttributes.path("Type").asText();
+    if ("String".equals(type)) {
+      DatadogAttributeParser.forEachProperty(classifier, value);
+    } else if ("Binary".equals(type)) {
+      ByteBuffer decodedValue = ByteBuffer.wrap(Base64.getDecoder().decode(value));
+      DatadogAttributeParser.forEachProperty(classifier, decodedValue);
     }
   }