From 8621b2aa8a421232d3af6649c2908323e87655ac Mon Sep 17 00:00:00 2001
From: Simon Bennetts <psiinon@gmail.com>
Date: Thu, 2 Jan 2025 14:36:18 +0000
Subject: [PATCH] Client: report nodes with no components as visited

Signed-off-by: Simon Bennetts <psiinon@gmail.com>
---
 .../addon/client/ClientIntegrationAPI.java     | 18 ++++++++++++------
 .../client/ExtensionClientIntegration.java     |  9 +++++++++
 .../addon/client/internal/ClientMap.java       | 10 ++++++++++
 3 files changed, 31 insertions(+), 6 deletions(-)

diff --git a/addOns/client/src/main/java/org/zaproxy/addon/client/ClientIntegrationAPI.java b/addOns/client/src/main/java/org/zaproxy/addon/client/ClientIntegrationAPI.java
index 8c033142e91..dde7973cb7a 100644
--- a/addOns/client/src/main/java/org/zaproxy/addon/client/ClientIntegrationAPI.java
+++ b/addOns/client/src/main/java/org/zaproxy/addon/client/ClientIntegrationAPI.java
@@ -116,6 +116,16 @@ private void handleReportObject(JSONObject json) {
         }
     }
 
+    private void handleReportEvent(JSONObject json) {
+        ReportedEvent event = new ReportedEvent(json);
+        if (event.getUrl() == null || !ExtensionClientIntegration.isApiUrl(event.getUrl())) {
+            this.extension.addReportedObject(event);
+            if (event.getUrl() != null) {
+                ThreadUtils.invokeAndWaitHandled(() -> this.extension.setVisited(event.getUrl()));
+            }
+        }
+    }
+
     @Override
     public ApiResponse handleApiAction(String name, JSONObject params) throws ApiException {
         JSONObject json;
@@ -131,11 +141,7 @@ public ApiResponse handleApiAction(String name, JSONObject params) throws ApiExc
                 String eventJson = this.getParam(params, PARAM_EVENT_JSON, "");
                 LOGGER.debug("Got event: {}", eventJson);
                 json = JSONObject.fromObject(eventJson);
-                ReportedEvent event = new ReportedEvent(json);
-                if (event.getUrl() == null
-                        || !ExtensionClientIntegration.isApiUrl(event.getUrl())) {
-                    this.extension.addReportedObject(event);
-                }
+                handleReportEvent(json);
                 break;
 
             case ACTION_REPORT_ZEST_STATEMENT:
@@ -193,7 +199,7 @@ public String handleCallBack(HttpMessage msg) throws ApiException {
             } else if (body.startsWith(PARAM_EVENT_JSON)) {
                 JSONObject json = decodeParam(body, PARAM_EVENT_JSON);
                 LOGGER.debug("Got event: {}", json);
-                this.extension.addReportedObject(new ReportedEvent(json));
+                handleReportEvent(json);
             } else if (body.startsWith(PARAM_STATEMENT_JSON)) {
                 try {
                     this.extension.addZestStatement(decodeParamString(body, PARAM_STATEMENT_JSON));
diff --git a/addOns/client/src/main/java/org/zaproxy/addon/client/ExtensionClientIntegration.java b/addOns/client/src/main/java/org/zaproxy/addon/client/ExtensionClientIntegration.java
index 80d79c30e67..c545a0c9cb1 100644
--- a/addOns/client/src/main/java/org/zaproxy/addon/client/ExtensionClientIntegration.java
+++ b/addOns/client/src/main/java/org/zaproxy/addon/client/ExtensionClientIntegration.java
@@ -495,6 +495,15 @@ public boolean setRedirect(String originalUrl, String redirectedUrl) {
         return false;
     }
 
+    public boolean setVisited(String url) {
+        ClientNode node = this.clientTree.setVisited(url);
+        if (node != null) {
+            this.clientNodeChanged(node);
+            return true;
+        }
+        return false;
+    }
+
     public void deleteNodes(List<ClientNode> nodes) {
         this.clientTree.deleteNodes(nodes);
         if (View.isInitialised()) {
diff --git a/addOns/client/src/main/java/org/zaproxy/addon/client/internal/ClientMap.java b/addOns/client/src/main/java/org/zaproxy/addon/client/internal/ClientMap.java
index 8b000eb8dce..3b25f535b05 100644
--- a/addOns/client/src/main/java/org/zaproxy/addon/client/internal/ClientMap.java
+++ b/addOns/client/src/main/java/org/zaproxy/addon/client/internal/ClientMap.java
@@ -198,6 +198,16 @@ public ClientNode setRedirect(String originalUrl, String redirectedUrl) {
         LOGGER.debug("setRedirect, no node for URL {}", originalUrl);
         return null;
     }
+
+    public ClientNode setVisited(String url) {
+        ClientNode node = getNode(url, false, false);
+        if (node != null && !node.getUserObject().isVisited()) {
+            node.getUserObject().setVisited(true);
+            return node;
+        }
+        LOGGER.debug("setVisited, no node for URL or already visited {}", url);
+        return null;
+    }
 }
 
 /**