Skip to content

client: persist client map #7371

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
thc202 wants to merge 1 commit into
base: main
Choose a base branch
from
+995 −59
Draft

client: persist client map #7371

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion addOns/client/CHANGELOG.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,8 @@ All notable changes to this add-on will be documented in this file.
The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).

## Unreleased

### Added
- Persist Client Map in the session.

## [0.26.0] - 2026-05-27
### Added
Expand Down
72 changes: 71 additions & 1 deletion addOns/client/src/main/java/org/zaproxy/addon/client/ExtensionClientIntegration.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,9 @@
import java.text.MessageFormat;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import javax.swing.ImageIcon;
import javax.swing.SwingUtilities;
import org.apache.commons.httpclient.URIException;
Expand Down Expand Up @@ -71,6 +73,7 @@
import org.zaproxy.addon.client.internal.ClientSideDetails;
import org.zaproxy.addon.client.internal.ReportedObject;
import org.zaproxy.addon.client.internal.db.ClientHistoryDao;
import org.zaproxy.addon.client.internal.db.ClientMapDao;
import org.zaproxy.addon.client.internal.db.TableJdo;
import org.zaproxy.addon.client.pscan.ClientPassiveScanController;
import org.zaproxy.addon.client.pscan.ClientPassiveScanHelper;
Expand Down Expand Up @@ -183,7 +186,28 @@ public void initModel(Model model) {
new ClientSideDetails(
Constant.messages.getString("client.tree.title"), null),
this.getModel().getSession()));
clientTree.setReportedObjectConsumer(this::addReportedObject);
clientTree.setPersistenceConsumer(
new ClientMap.PersistenceConsumer() {
@Override
public void onNodeAdded(ClientNode node) {
node.setPersistenceId(ClientMapDao.persistNode(node));
}

@Override
public void onNodeChanged(ClientNode node) {
ClientMapDao.updateNode(node);
}

@Override
public void onComponentAdded(ClientNode node, ClientSideComponent component) {
component.setPersistenceId(ClientMapDao.persistComponent(node, component));
}

@Override
public void onReportedObject(ReportedObject obj) {
addReportedObject(obj);
}
});
spiderScanController =
new SpiderScanController(
this,
Expand Down Expand Up @@ -416,6 +440,7 @@ public void postInit() {
@Override
public void postInstall() {
loadClientHistory();
loadClientMap();
}

private void loadClientHistory() {
Expand All @@ -425,6 +450,36 @@ private void loadClientHistory() {
}
}

private void loadClientMap() {
clientTree.setPublishEvents(false);
if (clientMapPanel != null) {
ThreadUtils.invokeAndWaitHandled(clientMapPanel::detachModel);
}
Comment thread
thc202 marked this conversation as resolved.
try {
Map<Long, String> nodeUrlById = new HashMap<>();
ClientMapDao.forEachNode(
entry -> {
ClientNode node =
clientTree.getOrAddNode(entry.getUrl(), entry.isVisited(), false);
node.getUserObject().setContentLoaded(entry.isContentLoaded());
node.getUserObject().setRedirect(entry.isRedirect());

node.setPersistenceId(entry.getId());
nodeUrlById.put(entry.getId(), entry.getUrl());
});
ClientMapDao.forEachComponent(
entry -> {
String url = nodeUrlById.get(entry.getNodeId());
clientTree.addComponent(url, ClientMapDao.toComponent(entry, url));
});
} finally {
clientTree.setPublishEvents(true);
if (clientMapPanel != null) {
ThreadUtils.invokeAndWaitHandled(clientMapPanel::attachModel);
}
}
}

public ClientOptions getClientParam() {
if (clientParam == null) {
clientParam = new ClientOptions();
Expand Down Expand Up @@ -536,6 +591,11 @@ public void clientNodeSelected(ClientNode node) {
}

public void deleteNodes(List<ClientNode> nodes) {
for (ClientNode node : nodes) {
if (!node.isRoot()) {
deleteNodeSubtreeFromDb(node);
}
}
this.clientTree.deleteNodes(nodes);
if (View.isInitialised()) {
String displayedUrl = this.getClientDetailsPanel().getCurrentUrl();
Expand All @@ -547,6 +607,15 @@ public void deleteNodes(List<ClientNode> nodes) {
}
}

private static void deleteNodeSubtreeFromDb(ClientNode node) {
if (!node.isStorage() && node.getPersistenceId() != -1) {
ClientMapDao.deleteNodeById(node.getPersistenceId());
}
for (int i = 0; i < node.getChildCount(); i++) {
deleteNodeSubtreeFromDb(node.getChildAt(i));
}
}

private ClientMapPanel getClientMapPanel() {
if (clientMapPanel == null) {
clientMapPanel = new ClientMapPanel(this, clientTree);
Expand Down Expand Up @@ -855,6 +924,7 @@ public void sessionChanged(final Session session) {
clientDetailsPanel.clear();
}
loadClientHistory();
loadClientMap();
spiderScanController.reset();

if (hasView()) {
Expand Down
120 changes: 82 additions & 38 deletions addOns/client/src/main/java/org/zaproxy/addon/client/internal/ClientMap.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,6 @@
import java.util.Locale;
import java.util.Map;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.function.Consumer;
import javax.swing.tree.DefaultTreeModel;
import javax.swing.tree.TreeNode;
import net.sf.json.JSONObject;
Expand Down Expand Up @@ -53,7 +52,8 @@ public class ClientMap extends SortedTreeModel implements EventPublisher {
private static final long serialVersionUID = 1L;
private static final Logger LOGGER = LogManager.getLogger(ClientMap.class);
private ClientNode root;
private Consumer<ReportedObject> reportedObjectConsumer;
private PersistenceConsumer persistenceConsumer;
private boolean publishEvents = true;
private final List<ClientMapListener> listeners = new CopyOnWriteArrayList<>();

public ClientMap(ClientNode root) {
Expand All @@ -75,6 +75,10 @@ public void removeListener(ClientMapListener listener) {
listeners.remove(listener);
}

public void setPublishEvents(boolean publishEvents) {
this.publishEvents = publishEvents;
}

public ClientNode getOrAddNode(String url, boolean visited, boolean storage) {
LOGGER.debug("getOrAddNode {}", url);
return this.getNode(url, visited, storage, true, true, 0);
Expand Down Expand Up @@ -119,7 +123,7 @@ private synchronized ClientNode getNode(
new ClientNode(
new ClientSideDetails(nodeName, url, visited, storage),
storage);
if (!storage && publishEvent) {
if (!storage && publishEvent && publishEvents) {
int depth = parent.getLevel() + 1;
int siblings = parent.getChildCount() + 1;
Map<String, String> map = new HashMap<>();
Expand Down Expand Up @@ -164,6 +168,9 @@ private synchronized ClientNode getNode(
new ClientNode(
new ClientSideDetails(nodeName, nodeUrl, false, false), false);
}
if (publishEvents && !storage && persistenceConsumer != null) {
persistenceConsumer.onNodeAdded(child);
}
this.insertNodeInto(child, parent);
this.nodeStructureChanged(parent);
}
Expand Down Expand Up @@ -222,15 +229,26 @@ private boolean addComponentToNode(ClientNode node, ClientSideComponent componen
if (!wasVisited || componentAdded) {
details.setVisited(true);

Map<String, String> map = new HashMap<>(component.getData());
map.put(DEPTH_KEY, Integer.toString(node.getLevel()));
map.put(SIBLINGS_KEY, Integer.toString(node.getChildCount()));
ZAP.getEventBus()
.publishSyncEvent(
this, new Event(this, MAP_COMPONENT_ADDED_EVENT, new Target(), map));
listeners.forEach(l -> l.componentAdded(map, source));
if (publishEvents) {
Map<String, String> map = new HashMap<>(component.getData());
map.put(DEPTH_KEY, Integer.toString(node.getLevel()));
map.put(SIBLINGS_KEY, Integer.toString(node.getChildCount()));
ZAP.getEventBus()
.publishSyncEvent(
this,
new Event(this, MAP_COMPONENT_ADDED_EVENT, new Target(), map));
listeners.forEach(l -> l.componentAdded(map, source));
}
notifyNodeChanged(node);
}
if (publishEvents && !node.isStorage() && persistenceConsumer != null) {
if (!wasVisited) {
persistenceConsumer.onNodeChanged(node);
}
if (componentAdded) {
persistenceConsumer.onComponentAdded(node, component);
}
}
return componentAdded;
}

Expand All @@ -239,19 +257,25 @@ public ClientNode setRedirect(String originalUrl, String redirectedUrl) {
if (node != null) {
node.getUserObject().setRedirect(true);
node.getUserObject().setVisited(true);
node.getUserObject()
.addComponent(
new ClientSideComponent(
Map.of(),
ClientSideComponent.REDIRECT,
null,
originalUrl,
redirectedUrl,
ClientSideComponent.REDIRECT,
ClientSideComponent.Type.REDIRECT,
null,
-1));
ClientSideComponent redirectComponent =
new ClientSideComponent(
Map.of(),
ClientSideComponent.REDIRECT,
null,
originalUrl,
redirectedUrl,
ClientSideComponent.REDIRECT,
ClientSideComponent.Type.REDIRECT,
null,
-1);
boolean componentAdded = node.getUserObject().addComponent(redirectComponent);
notifyNodeChanged(node);
if (publishEvents && persistenceConsumer != null) {
persistenceConsumer.onNodeChanged(node);
if (componentAdded) {
persistenceConsumer.onComponentAdded(node, redirectComponent);
}
}
return node;
}
LOGGER.debug("setRedirect, no node for URL {}", originalUrl);
Expand All @@ -263,6 +287,9 @@ public ClientNode setVisited(String url) {
if (node != null && !node.getUserObject().isVisited()) {
node.getUserObject().setVisited(true);
notifyNodeChanged(node);
if (publishEvents && persistenceConsumer != null) {
persistenceConsumer.onNodeChanged(node);
}
return node;
}
LOGGER.debug("setVisited, no node for URL or already visited {}", url);
Expand All @@ -276,24 +303,30 @@ public ClientNode setContentLoaded(String url) {
}

node.getUserObject().setContentLoaded(true);
node.getUserObject()
.addComponent(
new ClientSideComponent(
Map.of(),
ClientSideComponent.CONTENT_LOADED,
null,
null,
null,
ClientSideComponent.CONTENT_LOADED,
ClientSideComponent.Type.CONTENT_LOADED,
null,
-1));
ClientSideComponent contentLoadedComponent =
new ClientSideComponent(
Map.of(),
ClientSideComponent.CONTENT_LOADED,
null,
null,
null,
ClientSideComponent.CONTENT_LOADED,
ClientSideComponent.Type.CONTENT_LOADED,
null,
-1);
boolean componentAdded = node.getUserObject().addComponent(contentLoadedComponent);
notifyNodeChanged(node);
if (publishEvents && persistenceConsumer != null) {
persistenceConsumer.onNodeChanged(node);
if (componentAdded) {
persistenceConsumer.onComponentAdded(node, contentLoadedComponent);
}
}
return node;
}

public void setReportedObjectConsumer(Consumer<ReportedObject> consumer) {
this.reportedObjectConsumer = consumer;
public void setPersistenceConsumer(PersistenceConsumer consumer) {
this.persistenceConsumer = consumer;
}

public void handleReportObject(String jsonStr) {
Expand Down Expand Up @@ -339,14 +372,25 @@ private void notifyReportedObjectConsumer(ReportedObject reportObject) {
return;
}

if (reportedObjectConsumer != null) {
reportedObjectConsumer.accept(reportObject);
if (publishEvents && persistenceConsumer != null) {
persistenceConsumer.onReportedObject(reportObject);
}
}

private static boolean isApiUrl(String url) {
return url != null && (url.startsWith(API.API_URL) || url.startsWith(API.API_URL_S));
}

public interface PersistenceConsumer {

void onNodeAdded(ClientNode node);

void onNodeChanged(ClientNode node);

void onComponentAdded(ClientNode node, ClientSideComponent component);

void onReportedObject(ReportedObject reportedObject);
}
}

/**
Expand Down
9 changes: 9 additions & 0 deletions addOns/client/src/main/java/org/zaproxy/addon/client/internal/ClientNode.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@ public class ClientNode extends DefaultMutableTreeNode {
private static final long serialVersionUID = 1L;
private boolean storage;
private Session session;
private long persistenceId = -1;

public ClientNode(ClientSideDetails userObject, Session session) {
super(userObject);
Expand Down Expand Up @@ -70,6 +71,14 @@ public boolean isStorage() {
return storage;
}

public long getPersistenceId() {
return persistenceId;
}

public void setPersistenceId(long persistenceId) {
this.persistenceId = persistenceId;
}

public String getSite() {
ClientNode parent = this.getParent();
if (parent.isRoot()) {
Expand Down
Loading
Loading