Ensure messages are not modified before span serialization

Author: manuel-alvarez-alvarez

dd-java-agent/agent-aiguard/src/main/java/com/datadog/aiguard/AIGuardInternal.java

@@ -135,18 +135,17 @@ private static List<Message> messagesForMetaStruct(List<Message> messages) {
     final int maxContent = config.getAiGuardMaxContentSize();
     boolean contentTruncated = false;
     for (int i = 0; i < size; i++) {
-      Message source = messages.get(i);
-      final String content = source.getContent();
+      final Message source = messages.get(i);
+      String content = source.getContent();
       if (content != null && content.length() > maxContent) {
         contentTruncated = true;
-        source =
-            new Message(
-                source.getRole(),
-                content.substring(0, maxContent),
-                source.getToolCalls(),
-                source.getToolCallId());
+        content = content.substring(0, maxContent);
       }
-      result.add(source);
+      List<ToolCall> toolCalls = source.getToolCalls();
+      if (toolCalls != null) {
+        toolCalls = new ArrayList<>(toolCalls);
+      }
+      result.add(new Message(source.getRole(), content, toolCalls, source.getToolCallId()));
     }
     if (contentTruncated) {
       WafMetricCollector.get().aiGuardTruncated(CONTENT);
@@ -240,7 +239,7 @@ public Evaluation evaluate(final List<Message> messages, final Options options)
           span.setTag(BLOCKED_TAG, true);
           throw new AIGuardAbortError(action, reason, tags);
         }
-        return new Evaluation(action, reason);
+        return new Evaluation(action, reason, tags);
       }
     } catch (AIGuardAbortError e) {
       span.addThrowable(e);

dd-java-agent/agent-aiguard/src/test/groovy/com/datadog/aiguard/AIGuardInternalTests.groovy

@@ -10,6 +10,7 @@ import datadog.trace.api.aiguard.AIGuard
 import datadog.trace.api.telemetry.WafMetricCollector
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer.TracerAPI
 import datadog.trace.test.util.DDSpecification
 import okhttp3.Call
 import okhttp3.HttpUrl
@@ -71,6 +72,7 @@ class AIGuardInternalTests extends DDSpecification {
   @Shared
   protected static final PROMPT = TOOL_OUTPUT + [AIGuard.Message.message('assistant', '2 + 2 is 5'), AIGuard.Message.message('user', '')]
 
+  protected TracerAPI tracer
   protected AgentSpan span
 
   void setup() {
@@ -81,7 +83,7 @@ class AIGuardInternalTests extends DDSpecification {
     final builder = Mock(AgentTracer.SpanBuilder) {
       start() >> span
     }
-    final tracer = Stub(AgentTracer.TracerAPI) {
+    tracer = Stub(TracerAPI) {
       buildSpan(_ as String, _ as String) >> builder
     }
     AgentTracer.forceRegister(tracer)
@@ -198,19 +200,18 @@ class AIGuardInternalTests extends DDSpecification {
       1 * span.addThrowable(_ as AIGuard.AIGuardAbortError)
     }
 
-    receivedMeta.messages == suite.messages
-    if (suite.tags) {
-      receivedMeta.attack_categories == suite.tags
-    }
+    assertMeta(receivedMeta, suite)
     assertRequest(request, suite.messages)
     if (throwAbortError) {
       error instanceof AIGuard.AIGuardAbortError
       error.action == suite.action
       error.reason == suite.reason
+      error.tags == suite.tags
     } else {
       error == null
       eval.action == suite.action
       eval.reason == suite.reason
+      eval.tags == suite.tags
     }
     assertTelemetry('ai_guard.requests', "action:$suite.action", "block:$throwAbortError", 'error:false')
 
@@ -221,19 +222,7 @@ class AIGuardInternalTests extends DDSpecification {
   void 'test evaluate with API errors'() {
     given:
     final errors = [[status: 400, title: 'Bad request']]
-    Request request = null
-    final call = Mock(Call) {
-      execute() >> {
-        return mockResponse(request, 404, [errors: errors])
-      }
-    }
-    final client = Mock(OkHttpClient) {
-      newCall(_ as Request) >> {
-        request = (Request) it[0]
-        return call
-      }
-    }
-    final aiguard = new AIGuardInternal(URL, HEADERS, client)
+    final aiguard = mockClient(404, [errors: errors])
 
     when:
     aiguard.evaluate(TOOL_CALL, AIGuard.Options.DEFAULT)
@@ -247,19 +236,7 @@ class AIGuardInternalTests extends DDSpecification {
 
   void 'test evaluate with invalid JSON'() {
     given:
-    Request request = null
-    final call = Mock(Call) {
-      execute() >> {
-        return mockResponse(request, 200, [bad: 'This is an invalid response'])
-      }
-    }
-    final client = Mock(OkHttpClient) {
-      newCall(_ as Request) >> {
-        request = (Request) it[0]
-        return call
-      }
-    }
-    final aiguard = new AIGuardInternal(URL, HEADERS, client)
+    final aiguard = mockClient(200, [bad: 'This is an invalid response'])
 
     when:
     aiguard.evaluate(TOOL_CALL, AIGuard.Options.DEFAULT)
@@ -272,19 +249,7 @@ class AIGuardInternalTests extends DDSpecification {
 
   void 'test evaluate with missing action'() {
     given:
-    Request request = null
-    final call = Mock(Call) {
-      execute() >> {
-        return mockResponse(request, 200, [data: [attributes: [reason: 'I miss something']]])
-      }
-    }
-    final client = Mock(OkHttpClient) {
-      newCall(_ as Request) >> {
-        request = (Request) it[0]
-        return call
-      }
-    }
-    final aiguard = new AIGuardInternal(URL, HEADERS, client)
+    final aiguard = mockClient(200, [data: [attributes: [reason: 'I miss something']]])
 
     when:
     aiguard.evaluate(TOOL_CALL, AIGuard.Options.DEFAULT)
@@ -297,19 +262,7 @@ class AIGuardInternalTests extends DDSpecification {
 
   void 'test evaluate with non JSON response'() {
     given:
-    Request request = null
-    final call = Mock(Call) {
-      execute() >> {
-        return mockResponse(request, 200, [data: [attributes: [reason: 'I miss something']]])
-      }
-    }
-    final client = Mock(OkHttpClient) {
-      newCall(_ as Request) >> {
-        request = (Request) it[0]
-        return call
-      }
-    }
-    final aiguard = new AIGuardInternal(URL, HEADERS, client)
+    final aiguard = mockClient(200, 'I am no JSON')
 
     when:
     aiguard.evaluate(TOOL_CALL, AIGuard.Options.DEFAULT)
@@ -322,19 +275,7 @@ class AIGuardInternalTests extends DDSpecification {
 
   void 'test evaluate with empty response'() {
     given:
-    Request request = null
-    final call = Mock(Call) {
-      execute() >> {
-        return mockResponse(request, 200, null)
-      }
-    }
-    final client = Mock(OkHttpClient) {
-      newCall(_ as Request) >> {
-        request = (Request) it[0]
-        return call
-      }
-    }
-    final aiguard = new AIGuardInternal(URL, HEADERS, client)
+    final aiguard = mockClient(200, null)
 
     when:
     aiguard.evaluate(TOOL_CALL, AIGuard.Options.DEFAULT)
@@ -348,19 +289,7 @@ class AIGuardInternalTests extends DDSpecification {
   void 'test message length truncation'() {
     given:
     final maxMessages = Config.get().getAiGuardMaxMessagesLength()
-    Request request = null
-    final call = Mock(Call) {
-      execute() >> {
-        return mockResponse(request, 200, [data: [attributes: [action: ALLOW, reason: 'It is fine']]])
-      }
-    }
-    final client = Mock(OkHttpClient) {
-      newCall(_ as Request) >> {
-        request = (Request) it[0]
-        return call
-      }
-    }
-    final aiguard = new AIGuardInternal(URL, HEADERS, client)
+    final aiguard = mockClient(200, [data: [attributes: [action: 'ALLOW', reason: 'It is fine']]])
     final messages = (0..maxMessages)
       .collect { AIGuard.Message.message('user', "This is a prompt: ${it}") }
       .toList()
@@ -380,19 +309,7 @@ class AIGuardInternalTests extends DDSpecification {
   void 'test message content truncation'() {
     given:
     final maxContent = Config.get().getAiGuardMaxContentSize()
-    Request request = null
-    final call = Mock(Call) {
-      execute() >> {
-        return mockResponse(request, 200, [data: [attributes: [action: ALLOW, reason: 'It is fine']]])
-      }
-    }
-    final client = Mock(OkHttpClient) {
-      newCall(_ as Request) >> {
-        request = (Request) it[0]
-        return call
-      }
-    }
-    final aiguard = new AIGuardInternal(URL, HEADERS, client)
+    final aiguard = mockClient(200, [data: [attributes: [action: 'ALLOW', reason: 'It is fine']]])
     final message = AIGuard.Message.message("user", (0..maxContent).collect { 'A' }.join())
 
     when:
@@ -426,23 +343,7 @@ class AIGuardInternalTests extends DDSpecification {
 
   void 'test missing tool name'() {
     given:
-    def request
-    final call = Mock(Call) {
-      execute() >> {
-        return mockResponse(
-          request,
-          200,
-          [data: [attributes: [action: 'ALLOW', reason: 'Just do it']]]
-          )
-      }
-    }
-    final client = Mock(OkHttpClient) {
-      newCall(_ as Request) >> {
-        request = (Request) it[0]
-        return call
-      }
-    }
-    final aiguard = new AIGuardInternal(URL, HEADERS, client)
+    final aiguard = mockClient(200, [data: [attributes: [action: 'ALLOW', reason: 'Just do it']]])
 
     when:
     aiguard.evaluate([AIGuard.Message.tool('call_1', 'Content')], AIGuard.Options.DEFAULT)
@@ -460,6 +361,52 @@ class AIGuardInternalTests extends DDSpecification {
     thrown(IllegalArgumentException)
   }
 
+  void 'test message immutability'() {
+    given:
+    final aiguard = mockClient(200, [data: [attributes: [action: 'ALLOW', reason: 'Just do it']]])
+    final message = new AIGuard.Message(
+      "assistant",
+      null,
+      [AIGuard.ToolCall.toolCall('call_1', 'execute_shell', '{"cmd": "ls -lah"}')],
+      null
+      )
+    Map<String, Object> receivedMeta
+
+    when:
+    aiguard.evaluate([message], AIGuard.Options.DEFAULT)
+
+    then:
+    1 * span.finish() >> {
+      // modify the tool calls before flushing
+      message.toolCalls.add(
+        AIGuard.ToolCall.toolCall('call_2', 'execute_shell', '{"cmd": "rm -rf"}')
+        )
+    }
+    1 * span.setMetaStruct(AIGuardInternal.META_STRUCT_TAG, _ as Map) >> {
+      receivedMeta = it[1] as Map<String, Object>
+      return span
+    }
+    final messages = receivedMeta.messages as List<AIGuard.Message>
+    final assistant = messages.first()
+    assistant.toolCalls.size() == 1
+  }
+
+  private AIGuardInternal mockClient(final int status, final Object response) {
+    def request
+    final call = Stub(Call) {
+      execute() >> {
+        return mockResponse(request, status, response)
+      }
+    }
+    final client = Stub(OkHttpClient) {
+      newCall(_ as Request) >> {
+        request = (Request) it[0]
+        return call
+      }
+    }
+    return new AIGuardInternal(URL, HEADERS, client)
+  }
+
   private static assertTelemetry(final String metric, final String...tags) {
     final metrics = WafMetricCollector.get().with {
       prepareMetrics()
@@ -475,6 +422,16 @@ class AIGuardInternalTests extends DDSpecification {
     return true
   }
 
+  private static assertMeta(final Map<String, Object> meta, final TestSuite suite) {
+    if (suite.tags) {
+      assert meta.attack_categories == suite.tags
+    }
+    final receivedMessages = snakeCaseJson(meta.messages)
+    final expectedMessages = snakeCaseJson(suite.messages)
+    JSONAssert.assertEquals(expectedMessages, receivedMessages, JSONCompareMode.NON_EXTENSIBLE)
+    return true
+  }
+
   private static assertRequest(final Request request, final List<AIGuard.Message> messages) {
     assert request.url() == URL
     assert request.method() == 'POST'
@@ -556,7 +513,8 @@ class AIGuardInternalTests extends DDSpecification {
       ", reason='" + reason + '\'' +
       ", blocking=" + blocking +
       ", target='" + target + '\'' +
-      ", messages=" + messages +
+      ", messages=" + messages + '\'' +
+      ", tags=" + tags +
       '}'
     }
   }

dd-trace-api/src/main/java/datadog/trace/api/aiguard/AIGuard.java

@@ -143,16 +143,19 @@ public static class Evaluation {
 
     final Action action;
     final String reason;
+    final List<String> tags;
 
     /**
      * Creates a new evaluation result.
      *
      * @param action the recommended action for the evaluated content
      * @param reason human-readable explanation for the decision
+     * @param tags list of tags associated with the evaluation (e.g. indirect-prompt-injection)
      */
-    public Evaluation(final Action action, final String reason) {
+    public Evaluation(final Action action, final String reason, final List<String> tags) {
       this.action = action;
       this.reason = reason;
+      this.tags = tags;
     }
 
     /**
@@ -172,6 +175,15 @@ public Action getAction() {
     public String getReason() {
       return reason;
     }
+
+    /**
+     * Returns the list of tags associated with the evaluation (e.g. indirect-prompt-injection)
+     *
+     * @return list of tags.
+     */
+    public List<String> getTags() {
+      return tags;
+    }
   }
 
   /**

dd-trace-api/src/main/java/datadog/trace/api/aiguard/noop/NoOpEvaluator.java

@@ -1,6 +1,7 @@
 package datadog.trace.api.aiguard.noop;
 
 import static datadog.trace.api.aiguard.AIGuard.Action.ALLOW;
+import static java.util.Collections.emptyList;
 
 import datadog.trace.api.aiguard.AIGuard.Evaluation;
 import datadog.trace.api.aiguard.AIGuard.Message;
@@ -12,6 +13,6 @@ public final class NoOpEvaluator implements Evaluator {
 
   @Override
   public Evaluation evaluate(final List<Message> messages, final Options options) {
-    return new Evaluation(ALLOW, "AI Guard is not enabled");
+    return new Evaluation(ALLOW, "AI Guard is not enabled", emptyList());
   }
 }