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

@@ -42,7 +42,6 @@
 import java.util.function.Supplier;
 import java.util.regex.Pattern;
 import java.util.stream.Stream;
-import lombok.Setter;
 import net.htmlparser.jericho.Element;
 import net.htmlparser.jericho.Source;
 import net.sf.json.JSONArray;
@@ -199,7 +198,7 @@ public void notifyMessageReceived(HttpMessage message) {
 
     private static long timeToWaitMs = TimeUnit.SECONDS.toMillis(5);
 
-    @Setter private static HistoryProvider historyProvider = new HistoryProvider();
+    private static HistoryProvider historyProvider = ExtensionAuthhelper.getHistoryProvider();
 
     /**
      * These are session tokens that have been seen in responses but not yet seen in use. When they
@@ -1551,4 +1550,9 @@ private static void setMinWaitFor(ZestStatement stmt, int minWaitForMsec) {
             }
         }
     }
+
+    /** For testing purposes only, not part of the public API */
+    public static void setHistoryProvider(HistoryProvider hp) {
+        AuthUtils.historyProvider = hp;
+    }
 }

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,28 +19,58 @@
  */
 package org.zaproxy.addon.authhelper;
 
+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.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
+            -- )
+            ORDER BY HISTORYID DESC
+            """;
+
+    private static PreparedStatement psGetHistory;
+
+    private ParosDatabaseServer pds;
+    private boolean warnedNonParosDb;
+
     private ExtensionHistory extHist;
 
+    HistoryProvider() {
+        getParaosDataBaseServer();
+    }
+
     private ExtensionHistory getExtHistory() {
         if (extHist == null) {
             extHist = AuthUtils.getExtension(ExtensionHistory.class);
@@ -61,6 +91,64 @@ 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) {
+        Connection conn = getDbConnection();
+        if (conn == null) {
+            return List.of();
+        }
+        PreparedStatement query = getHistoryQuery(first, last, value, conn);
+        if (query == null) {
+            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.");
+        }
+        LOGGER.debug("Found: {} candidate messages for {}", msgIds.size(), value);
+        LOGGER.info("{} IDs", msgIds.size());
+        return msgIds;
+    }
+
+    private static PreparedStatement getHistoryQuery(
+            int first, int last, String value, Connection conn) {
+        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);
+            }
+        } catch (SQLException e) {
+            LOGGER.warn("Failed to prepare query.", e);
+        }
+        return psGetHistory;
+    }
+
+    private Connection getDbConnection() {
+        if (pds == null) {
+            LOGGER.info("PDS was null");
+            return null;
+        }
+        Connection conn = null;
+        try {
+            conn = pds.getSingletonConnection();
+        } catch (SQLException | NullPointerException e) {
+            LOGGER.warn("Failed to get DB connection.", e);
+        }
+        return conn;
+    }
+
     public int getLastHistoryId() {
         return getExtHistory().getLastHistoryId();
     }
@@ -73,9 +161,9 @@ public SessionManagementRequestDetails findSessionTokenSource(String token, int
 
         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 +187,45 @@ public SessionManagementRequestDetails findSessionTokenSource(String token, int
         }
         return null;
     }
+
+    private ParosDatabaseServer getParaosDataBaseServer() {
+        if (Model.getSingleton().getDb().getDatabaseServer() instanceof ParosDatabaseServer pdbs) {
+            pds = pdbs;
+            LOGGER.info("PDS ? {}", pds != null);
+        } else {
+            if (pds == null && !warnedNonParosDb) {
+                LOGGER.warn("Unexpected Database Server.");
+                warnedNonParosDb = true;
+            }
+        }
+        return pds;
+    }
+
+    @Override
+    public void sessionChanged(Session session) {
+        pds = null;
+        warnedNonParosDb = false;
+        getParaosDataBaseServer();
+    }
+
+    @Override
+    public void sessionAboutToChange(Session session) {
+        try {
+            if (psGetHistory != null) {
+                psGetHistory.close();
+            }
+        } catch (SQLException e) {
+            // 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,16 @@ public List<SessionToken> getTokens() {
     public int getConfidence() {
         return confidence;
     }
+
+    @Override
+    public String toString() {
+        return "MsgID: "
+                + msg.getHistoryRef().getHistoryId()
+                + " 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/SessionDetectionScanRuleUnitTest.java

@@ -372,17 +372,6 @@ void shouldFindTokenWhenOneIsPreviouslyUnknown() throws Exception {
         extensionLoader =
                 mock(ExtensionLoader.class, withSettings().strictness(Strictness.LENIENT));
 
-        history = new ArrayList<>();
-        historyProvider = new TestHistoryProvider();
-        AuthUtils.setHistoryProvider(historyProvider);
-
-        Control.initSingletonForTesting(model, extensionLoader);
-        Model.setSingletonForTesting(model);
-
-        Session session = mock(Session.class);
-        given(session.getContextsForUrl(anyString())).willReturn(Arrays.asList());
-        given(model.getSession()).willReturn(session);
-
         String cookie = "67890123456789012345";
         String jwtValue = "bearer 677890123456789012345-677890123456789012345";
         String jwt = "{\"jwt\":\"%s\"}".formatted(jwtValue);
@@ -416,6 +405,17 @@ void shouldFindTokenWhenOneIsPreviouslyUnknown() throws Exception {
                         new HttpResponseHeader("HTTP/1.1 200 OK\r\n"),
                         new HttpResponseBody("<html></html>"));
 
+        history = new ArrayList<>();
+        historyProvider = new SuccessHistoryProvider(msg2);
+        AuthUtils.setHistoryProvider(historyProvider);
+
+        Control.initSingletonForTesting(model, extensionLoader);
+        Model.setSingletonForTesting(model);
+
+        Session session = mock(Session.class);
+        given(session.getContextsForUrl(anyString())).willReturn(Arrays.asList());
+        given(model.getSession()).willReturn(session);
+
         List<HttpMessage> msgs = List.of(msg, msg2);
         historyProvider.addAuthMessageToHistory(msg);
         historyProvider.addAuthMessageToHistory(msg2);
@@ -493,4 +493,21 @@ public int getLastHistoryId() {
             return history.size() - 1;
         }
     }
+
+    class SuccessHistoryProvider extends TestHistoryProvider {
+
+        private HttpMessage msg;
+
+        SuccessHistoryProvider(HttpMessage msg) {
+            this.msg = msg;
+        }
+
+        @Override
+        public SessionManagementRequestDetails findSessionTokenSource(String token, int firstId) {
+            SessionToken st =
+                    new SessionToken(
+                            "json", "jwt", "bearer 677890123456789012345-677890123456789012345");
+            return new SessionManagementRequestDetails(msg, List.of(st), 3);
+        }
+    }
 }