authhelper: Session Detection rule performance fixes

Signed-off-by: kingthorin <[email protected]>

Author: kingthorin

addOns/authhelper/CHANGELOG.md

@@ -11,6 +11,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 - Fail the Microsoft login if not able to perform all the expected steps.
 - Track GWT headers.
 - Handle additional exceptions when processing JSON authentication components.
+- Improved performance of the Session Detection scan rule.
 
 ### Fixed
 - Do not include known authentication providers in context.

addOns/authhelper/src/main/java/org/zaproxy/addon/authhelper/AuthUtils.java

@@ -199,7 +199,8 @@ public void notifyMessageReceived(HttpMessage message) {
 
     private static long timeToWaitMs = TimeUnit.SECONDS.toMillis(5);
 
-    @Setter private static HistoryProvider historyProvider = new HistoryProvider();
+    @Setter
+    private static HistoryProvider historyProvider = ExtensionAuthhelper.getHistoryProvider();
 
     /**
      * These are session tokens that have been seen in responses but not yet seen in use. When they

addOns/authhelper/src/main/java/org/zaproxy/addon/authhelper/ExtensionAuthhelper.java

@@ -33,6 +33,7 @@
 import java.util.Set;
 import java.util.stream.Stream;
 import javax.swing.ImageIcon;
+import lombok.Getter;
 import org.apache.commons.configuration.Configuration;
 import org.apache.commons.lang3.StringUtils;
 import org.apache.commons.text.StringEscapeUtils;
@@ -108,6 +109,8 @@ public class ExtensionAuthhelper extends ExtensionAdaptor {
 
     public static final Set<Integer> HISTORY_TYPES_SET = Set.of(HISTORY_TYPES);
 
+    @Getter private static HistoryProvider historyProvider = new HistoryProvider();
+
     private ZapMenuItem authTesterMenu;
     private AuthTestDialog authTestDialog;
 
@@ -212,6 +215,7 @@ public void destroy() {
     @Override
     public void hook(ExtensionHook extensionHook) {
         extensionHook.addSessionListener(new AuthSessionChangedListener());
+        extensionHook.addSessionListener(historyProvider);
         extensionHook.addHttpSenderListener(authHeaderTracker);
         extensionHook.addOptionsParamSet(getParam());
         if (hasView()) {

addOns/authhelper/src/main/java/org/zaproxy/addon/authhelper/HistoryProvider.java

@@ -19,26 +19,56 @@
  */
 package org.zaproxy.addon.authhelper;
 
+import java.net.URLEncoder;
+import java.nio.charset.StandardCharsets;
+import java.sql.Connection;
+import java.sql.PreparedStatement;
+import java.sql.ResultSet;
+import java.sql.SQLException;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Optional;
+import org.apache.commons.text.StringEscapeUtils;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
+import org.parosproxy.paros.control.Control.Mode;
 import org.parosproxy.paros.core.scanner.Alert;
 import org.parosproxy.paros.db.DatabaseException;
+import org.parosproxy.paros.db.paros.ParosDatabaseServer;
+import org.parosproxy.paros.extension.SessionChangedListener;
 import org.parosproxy.paros.extension.history.ExtensionHistory;
 import org.parosproxy.paros.model.HistoryReference;
+import org.parosproxy.paros.model.Model;
+import org.parosproxy.paros.model.Session;
 import org.parosproxy.paros.network.HttpMalformedHeaderException;
 import org.parosproxy.paros.network.HttpMessage;
 import org.zaproxy.zap.authentication.AuthenticationHelper;
 
 /** A very thin layer on top of the History functionality, to make testing easier. */
-public class HistoryProvider {
+public class HistoryProvider implements SessionChangedListener {
 
-    private static final int MAX_NUM_RECORDS_TO_CHECK = 200;
+    protected static final int MAX_NUM_RECORDS_TO_CHECK = 200;
 
     private static final Logger LOGGER = LogManager.getLogger(HistoryProvider.class);
 
+    private static final String QUERY_SESS_MGMT_TOKEN_MSG_IDS =
+            """
+            SELECT HISTORYID FROM HISTORY
+            WHERE HISTORYID BETWEEN ? AND ?
+            AND (
+                POSITION(? IN RESHEADER) > 0
+                OR POSITION(? IN RESBODY) > 0
+                OR POSITION(? IN REQHEADER) > 0
+                OR POSITION(? IN REQHEADER) > 0 -- URLEncoded
+                OR POSITION(? IN RESBODY) > 0 -- JSONEscaped
+            )
+            ORDER BY HISTORYID DESC
+            """;
+
+    private ParosDatabaseServer pds;
+    private boolean warnedNonParosDb;
+    private boolean server;
+
     private ExtensionHistory extHist;
 
     private ExtensionHistory getExtHistory() {
@@ -61,21 +91,77 @@ public HttpMessage getHttpMessage(int historyId)
         return null;
     }
 
+    /**
+     * The query is ordered DESCending so the List and subsequent processing should be newest
+     * message first.
+     */
+    List<Integer> getMessageIds(int first, int last, String value) {
+        if (!server) {
+            getParosDataBaseServer();
+        }
+        if (pds == null) {
+            return null;
+        }
+
+        Connection conn = null;
+        try {
+            conn = pds.getSingletonConnection();
+        } catch (SQLException | NullPointerException e) {
+            LOGGER.warn("Failed to get DB connection.", e);
+            return List.of();
+        }
+
+        PreparedStatement psGetHistory = null;
+        try {
+            if (psGetHistory == null || psGetHistory.isClosed()) {
+                psGetHistory = conn.prepareStatement(QUERY_SESS_MGMT_TOKEN_MSG_IDS);
+            }
+            psGetHistory.setInt(1, first);
+            psGetHistory.setInt(2, last);
+            psGetHistory.setString(3, value);
+            psGetHistory.setBytes(4, value.getBytes(StandardCharsets.UTF_8));
+            psGetHistory.setString(5, value);
+            psGetHistory.setString(6, URLEncoder.encode(value, StandardCharsets.UTF_8));
+            psGetHistory.setBytes(
+                    7, StringEscapeUtils.escapeJson(value).getBytes(StandardCharsets.UTF_8));
+        } catch (SQLException e) {
+            LOGGER.warn("Failed to prepare query.", e);
+            return List.of();
+        }
+
+        List<Integer> msgIds = new ArrayList<>();
+        try (ResultSet rs = psGetHistory.executeQuery()) {
+            while (rs.next()) {
+                msgIds.add(rs.getInt("HISTORYID"));
+            }
+        } catch (SQLException e) {
+            LOGGER.warn("Failed to process result set. {}", e.getMessage());
+        } finally {
+            try {
+                psGetHistory.close();
+            } catch (SQLException e) {
+                // Nothing to do
+            }
+        }
+        LOGGER.debug("Found: {} candidate messages for {}", msgIds.size(), value);
+        return msgIds;
+    }
+
     public int getLastHistoryId() {
         return getExtHistory().getLastHistoryId();
     }
 
     public SessionManagementRequestDetails findSessionTokenSource(String token, int firstId) {
         int lastId = getLastHistoryId();
         if (firstId == -1) {
-            firstId = Math.max(0, lastId - MAX_NUM_RECORDS_TO_CHECK);
+            firstId = Math.max(1, lastId - MAX_NUM_RECORDS_TO_CHECK);
         }
 
         LOGGER.debug("Searching for session token from {} down to {} ", lastId, firstId);
 
-        for (int i = lastId; i >= firstId; i--) {
+        for (int id : getMessageIds(firstId, lastId, token)) {
             try {
-                HttpMessage msg = getHttpMessage(i);
+                HttpMessage msg = getHttpMessage(id);
                 if (msg == null) {
                     continue;
                 }
@@ -99,4 +185,42 @@ public SessionManagementRequestDetails findSessionTokenSource(String token, int
         }
         return null;
     }
+
+    private ParosDatabaseServer getParosDataBaseServer() {
+        if (Model.getSingleton().getDb() == null) {
+            return null; // This should only happen in unit tests
+        }
+        if (Model.getSingleton().getDb().getDatabaseServer() instanceof ParosDatabaseServer pdbs) {
+            pds = pdbs;
+            server = true;
+        } else {
+            if (pds == null && !warnedNonParosDb) {
+                LOGGER.warn("Unsupported Database Server.");
+                warnedNonParosDb = true;
+            }
+        }
+        return pds;
+    }
+
+    @Override
+    public void sessionChanged(Session session) {
+        pds = null;
+        server = false;
+        warnedNonParosDb = false;
+    }
+
+    @Override
+    public void sessionAboutToChange(Session session) {
+        // Nothing to do
+    }
+
+    @Override
+    public void sessionScopeChanged(Session session) {
+        // Nothing to do
+    }
+
+    @Override
+    public void sessionModeChanged(Mode mode) {
+        // Nothing to do
+    }
 }

addOns/authhelper/src/main/java/org/zaproxy/addon/authhelper/SessionManagementRequestDetails.java

@@ -46,4 +46,20 @@ public List<SessionToken> getTokens() {
     public int getConfidence() {
         return confidence;
     }
+
+    @Override
+    public String toString() {
+        String historyIdStr =
+                msg.getHistoryRef() != null
+                        ? String.valueOf(msg.getHistoryRef().getHistoryId())
+                        : "N/A";
+        return "MsgID: "
+                + historyIdStr
+                + " tokens ("
+                + tokens.size()
+                + "): "
+                + tokens
+                + " confidence: "
+                + confidence;
+    }
 }

addOns/authhelper/src/main/java/org/zaproxy/addon/authhelper/SessionToken.java

@@ -122,4 +122,9 @@ private static int compareStrings(String string, String otherString) {
         }
         return string.compareTo(otherString);
     }
+
+    @Override
+    public String toString() {
+        return "Source: " + source + " key: " + key + " value: " + value;
+    }
 }

addOns/authhelper/src/main/java/org/zaproxy/addon/authhelper/internal/ClientSideHandler.java

@@ -27,7 +27,6 @@
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 import lombok.Getter;
-import lombok.Setter;
 import org.apache.commons.httpclient.URIException;
 import org.apache.commons.lang3.StringUtils;
 import org.apache.logging.log4j.LogManager;
@@ -60,7 +59,7 @@ public final class ClientSideHandler implements HttpMessageHandler {
     private AuthRequestDetails authReq;
     private int firstHrefId;
 
-    @Setter private HistoryProvider historyProvider = new HistoryProvider();
+    private HistoryProvider historyProvider = ExtensionAuthhelper.getHistoryProvider();
 
     public ClientSideHandler(User user) {
         this.user = user;
@@ -224,6 +223,11 @@ protected static int messageTokenCount(HttpMessage msg, List<Pair<String, String
         return count;
     }
 
+    /** For testing purposes only, not part of the public API */
+    void setHistoryProvider(HistoryProvider hp) {
+        this.historyProvider = hp;
+    }
+
     @Getter
     class AuthRequestDetails {
         private HttpMessage msg;

addOns/authhelper/src/test/java/org/zaproxy/addon/authhelper/AuthUtilsUnitTest.java

@@ -106,6 +106,7 @@ void setUp() throws Exception {
         setUpZap();
 
         mockMessages(new ExtensionAuthhelper());
+        AuthUtils.setHistoryProvider(new TestHistoryProvider());
     }
 
     @AfterEach