authhelper: Session Detection rule performance fixes

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,54 @@
  */
 package org.zaproxy.addon.authhelper;
 
+import java.net.URLEncoder;
+import java.nio.charset.StandardCharsets;
+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;
 
     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 server;
+
     private ExtensionHistory extHist;
 
     private ExtensionHistory getExtHistory() {
@@ -61,21 +89,66 @@ 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) {
+            server = true;
+            if (Model.getSingleton().getDb().getDatabaseServer()
+                    instanceof ParosDatabaseServer pdbs) {
+                pds = pdbs;
+            } else {
+                LOGGER.warn("Unsupported Database Server.");
+            }
+        }
+        if (pds == null) {
+            return List.of();
+        }
+
+        try (PreparedStatement psGetHistory =
+                pds.getSingletonConnection().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));
+
+            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());
+            }
+            LOGGER.debug("Found: {} candidate messages for {}", msgIds.size(), value);
+            return msgIds;
+        } catch (SQLException e) {
+            LOGGER.warn("Failed to prepare query.", e);
+            return List.of();
+        }
+    }
+
     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 +172,25 @@ public SessionManagementRequestDetails findSessionTokenSource(String token, int
         }
         return null;
     }
+
+    @Override
+    public void sessionChanged(Session session) {
+        pds = null;
+        server = 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

@@ -60,7 +60,7 @@ public final class ClientSideHandler implements HttpMessageHandler {
     private AuthRequestDetails authReq;
     private int firstHrefId;
 
-    @Setter private HistoryProvider historyProvider = new HistoryProvider();
+    @Setter private HistoryProvider historyProvider = ExtensionAuthhelper.getHistoryProvider();
 
     public ClientSideHandler(User user) {
         this.user = user;

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

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

@@ -31,7 +31,6 @@
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.withSettings;
 
-import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
 import org.junit.jupiter.api.AfterEach;
@@ -43,9 +42,7 @@
 import org.parosproxy.paros.Constant;
 import org.parosproxy.paros.control.Control;
 import org.parosproxy.paros.core.scanner.Alert;
-import org.parosproxy.paros.db.DatabaseException;
 import org.parosproxy.paros.extension.ExtensionLoader;
-import org.parosproxy.paros.model.HistoryReference;
 import org.parosproxy.paros.model.Model;
 import org.parosproxy.paros.model.Session;
 import org.parosproxy.paros.network.HttpHeader;
@@ -67,7 +64,6 @@
 /** Unit test for {@link SessionDetectionScanRule}. */
 class SessionDetectionScanRuleUnitTest extends PassiveScannerTest<SessionDetectionScanRule> {
 
-    private List<HttpMessage> history;
     private HistoryProvider historyProvider;
 
     @Override
@@ -107,7 +103,6 @@ void shouldSetHeaderBasedSessionManagment() throws Exception {
         given(session.getContextsForUrl(anyString())).willReturn(Arrays.asList(context));
         given(model.getSession()).willReturn(session);
 
-        history = new ArrayList<>();
         historyProvider = new TestHistoryProvider();
         AuthUtils.setHistoryProvider(historyProvider);
 
@@ -204,7 +199,6 @@ void shouldDetectBodgeitSession() throws Exception {
                 DiagnosticDataLoader.loadTestData(
                         this.getResourcePath("internal/bodgeit.diags").toFile());
 
-        history = new ArrayList<>();
         historyProvider = new TestHistoryProvider();
         AuthUtils.setHistoryProvider(historyProvider);
         // When
@@ -228,7 +222,6 @@ void shouldDetectBodgeitSession() throws Exception {
     @Test
     void shouldDetectCtflearnSession() throws Exception {
         // Given
-        history = new ArrayList<>();
         historyProvider = new TestHistoryProvider();
         AuthUtils.setHistoryProvider(historyProvider);
 
@@ -267,7 +260,6 @@ void shouldDetectCtflearnSession() throws Exception {
     @Test
     void shouldDetectDefthewebSession() throws Exception {
         // Given
-        history = new ArrayList<>();
         historyProvider = new TestHistoryProvider();
         AuthUtils.setHistoryProvider(historyProvider);
 
@@ -295,7 +287,6 @@ void shouldDetectDefthewebSession() throws Exception {
     @Test
     void shouldDetectGinnjuiceSession() throws Exception {
         // Given
-        history = new ArrayList<>();
         historyProvider = new TestHistoryProvider();
         AuthUtils.setHistoryProvider(historyProvider);
 
@@ -325,7 +316,6 @@ void shouldDetectGinnjuiceSession() throws Exception {
     @Test
     void shouldDetectInfosecexSession() throws Exception {
         // Given
-        history = new ArrayList<>();
         historyProvider = new TestHistoryProvider();
         AuthUtils.setHistoryProvider(historyProvider);
 
@@ -372,7 +362,6 @@ void shouldFindTokenWhenOneIsPreviouslyUnknown() throws Exception {
         extensionLoader =
                 mock(ExtensionLoader.class, withSettings().strictness(Strictness.LENIENT));
 
-        history = new ArrayList<>();
         historyProvider = new TestHistoryProvider();
         AuthUtils.setHistoryProvider(historyProvider);
 
@@ -470,27 +459,4 @@ void shouldIgnoreUnknownOrUnwantedContentTypes(String contentType)
         // Then
         assertThat(alertsRaised, hasSize(0));
     }
-
-    class TestHistoryProvider extends HistoryProvider {
-        @Override
-        public void addAuthMessageToHistory(HttpMessage msg) {
-            history.add(msg);
-            int id = history.size() - 1;
-            HistoryReference href =
-                    mock(HistoryReference.class, withSettings().strictness(Strictness.LENIENT));
-            given(href.getHistoryId()).willReturn(id);
-            msg.setHistoryRef(href);
-        }
-
-        @Override
-        public HttpMessage getHttpMessage(int historyId)
-                throws HttpMalformedHeaderException, DatabaseException {
-            return history.get(historyId);
-        }
-
-        @Override
-        public int getLastHistoryId() {
-            return history.size() - 1;
-        }
-    }
 }