Skip to content

Commit de9f7e9

Browse files
kingthorinkingthorin
committed
authhelper: Session Detection rule performance fixes
Signed-off-by: kingthorin <[email protected]>
1 parent 942b206 commit de9f7e9

File tree

8 files changed

+194
-18
lines changed

8 files changed

+194
-18
lines changed

‎addOns/authhelper/CHANGELOG.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
1111
- Fail the Microsoft login if not able to perform all the expected steps.
1212
- Track GWT headers.
1313
- Handle additional exceptions when processing JSON authentication components.
14+
- Improved performance of the Session Detection scan rule.
1415

1516
### Fixed
1617
- Do not include known authentication providers in context.

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

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,6 @@
4242
import java.util.function.Supplier;
4343
import java.util.regex.Pattern;
4444
import java.util.stream.Stream;
45-
import lombok.Setter;
4645
import net.htmlparser.jericho.Element;
4746
import net.htmlparser.jericho.Source;
4847
import net.sf.json.JSONArray;
@@ -199,7 +198,7 @@ public void notifyMessageReceived(HttpMessage message) {
199198

200199
private static long timeToWaitMs = TimeUnit.SECONDS.toMillis(5);
201200

202-
@Setter private static HistoryProvider historyProvider = new HistoryProvider();
201+
private static HistoryProvider historyProvider = ExtensionAuthhelper.getHistoryProvider();
203202

204203
/**
205204
* 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) {
15511550
}
15521551
}
15531552
}
1553+
1554+
/** For testing purposes only, not part of the public API */
1555+
public static void setHistoryProvider(HistoryProvider hp) {
1556+
AuthUtils.historyProvider = hp;
1557+
}
15541558
}

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@
3333
import java.util.Set;
3434
import java.util.stream.Stream;
3535
import javax.swing.ImageIcon;
36+
import lombok.Getter;
3637
import org.apache.commons.configuration.Configuration;
3738
import org.apache.commons.lang3.StringUtils;
3839
import org.apache.commons.text.StringEscapeUtils;
@@ -108,6 +109,8 @@ public class ExtensionAuthhelper extends ExtensionAdaptor {
108109

109110
public static final Set<Integer> HISTORY_TYPES_SET = Set.of(HISTORY_TYPES);
110111

112+
@Getter private static HistoryProvider historyProvider = new HistoryProvider();
113+
111114
private ZapMenuItem authTesterMenu;
112115
private AuthTestDialog authTestDialog;
113116

@@ -212,6 +215,7 @@ public void destroy() {
212215
@Override
213216
public void hook(ExtensionHook extensionHook) {
214217
extensionHook.addSessionListener(new AuthSessionChangedListener());
218+
extensionHook.addSessionListener(historyProvider);
215219
extensionHook.addHttpSenderListener(authHeaderTracker);
216220
extensionHook.addOptionsParamSet(getParam());
217221
if (hasView()) {

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

Lines changed: 132 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -19,28 +19,58 @@
1919
*/
2020
package org.zaproxy.addon.authhelper;
2121

22+
import java.sql.Connection;
23+
import java.sql.PreparedStatement;
24+
import java.sql.ResultSet;
25+
import java.sql.SQLException;
2226
import java.util.ArrayList;
2327
import java.util.List;
2428
import java.util.Optional;
2529
import org.apache.logging.log4j.LogManager;
2630
import org.apache.logging.log4j.Logger;
31+
import org.parosproxy.paros.control.Control.Mode;
2732
import org.parosproxy.paros.core.scanner.Alert;
2833
import org.parosproxy.paros.db.DatabaseException;
34+
import org.parosproxy.paros.db.paros.ParosDatabaseServer;
35+
import org.parosproxy.paros.extension.SessionChangedListener;
2936
import org.parosproxy.paros.extension.history.ExtensionHistory;
3037
import org.parosproxy.paros.model.HistoryReference;
38+
import org.parosproxy.paros.model.Model;
39+
import org.parosproxy.paros.model.Session;
3140
import org.parosproxy.paros.network.HttpMalformedHeaderException;
3241
import org.parosproxy.paros.network.HttpMessage;
3342
import org.zaproxy.zap.authentication.AuthenticationHelper;
3443

3544
/** A very thin layer on top of the History functionality, to make testing easier. */
36-
public class HistoryProvider {
45+
public class HistoryProvider implements SessionChangedListener {
3746

3847
private static final int MAX_NUM_RECORDS_TO_CHECK = 200;
3948

4049
private static final Logger LOGGER = LogManager.getLogger(HistoryProvider.class);
4150

51+
private static final String QUERY_SESS_MGMT_TOKEN_MSG_IDS =
52+
"""
53+
SELECT HISTORYID FROM HISTORY
54+
WHERE HISTORYID BETWEEN ? AND ?
55+
-- AND (
56+
-- POSITION(? IN RESHEADER) > 0
57+
-- OR POSITION(? IN RESBODY) > 0
58+
-- OR POSITION(? IN REQHEADER) > 0
59+
-- )
60+
ORDER BY HISTORYID DESC
61+
""";
62+
63+
private static PreparedStatement psGetHistory;
64+
65+
private ParosDatabaseServer pds;
66+
private boolean warnedNonParosDb;
67+
4268
private ExtensionHistory extHist;
4369

70+
HistoryProvider() {
71+
getParaosDataBaseServer();
72+
}
73+
4474
private ExtensionHistory getExtHistory() {
4575
if (extHist == null) {
4676
extHist = AuthUtils.getExtension(ExtensionHistory.class);
@@ -61,6 +91,64 @@ public HttpMessage getHttpMessage(int historyId)
6191
return null;
6292
}
6393

94+
/**
95+
* The query is ordered DESCending so the List and subsequent processing should be newest
96+
* message first.
97+
*/
98+
List<Integer> getMessageIds(int first, int last, String value) {
99+
Connection conn = getDbConnection();
100+
if (conn == null) {
101+
return List.of();
102+
}
103+
PreparedStatement query = getHistoryQuery(first, last, value, conn);
104+
if (query == null) {
105+
return List.of();
106+
}
107+
List<Integer> msgIds = new ArrayList<>();
108+
109+
try (ResultSet rs = psGetHistory.executeQuery()) {
110+
while (rs.next()) {
111+
msgIds.add(rs.getInt("HISTORYID"));
112+
}
113+
} catch (SQLException e) {
114+
LOGGER.warn("Failed to process result set.");
115+
}
116+
LOGGER.debug("Found: {} candidate messages for {}", msgIds.size(), value);
117+
LOGGER.info("{} IDs", msgIds.size());
118+
return msgIds;
119+
}
120+
121+
private static PreparedStatement getHistoryQuery(
122+
int first, int last, String value, Connection conn) {
123+
try {
124+
if (psGetHistory == null || psGetHistory.isClosed()) {
125+
psGetHistory = conn.prepareStatement(QUERY_SESS_MGMT_TOKEN_MSG_IDS);
126+
psGetHistory.setInt(1, first);
127+
psGetHistory.setInt(2, last);
128+
// psGetHistory.setString(3, value);
129+
// psGetHistory.setBytes(4, value.getBytes(StandardCharsets.UTF_8));
130+
// psGetHistory.setString(5, value);
131+
}
132+
} catch (SQLException e) {
133+
LOGGER.warn("Failed to prepare query.", e);
134+
}
135+
return psGetHistory;
136+
}
137+
138+
private Connection getDbConnection() {
139+
if (pds == null) {
140+
LOGGER.info("PDS was null");
141+
return null;
142+
}
143+
Connection conn = null;
144+
try {
145+
conn = pds.getSingletonConnection();
146+
} catch (SQLException | NullPointerException e) {
147+
LOGGER.warn("Failed to get DB connection.", e);
148+
}
149+
return conn;
150+
}
151+
64152
public int getLastHistoryId() {
65153
return getExtHistory().getLastHistoryId();
66154
}
@@ -73,9 +161,9 @@ public SessionManagementRequestDetails findSessionTokenSource(String token, int
73161

74162
LOGGER.debug("Searching for session token from {} down to {} ", lastId, firstId);
75163

76-
for (int i = lastId; i >= firstId; i--) {
164+
for (int id : getMessageIds(firstId, lastId, token)) {
77165
try {
78-
HttpMessage msg = getHttpMessage(i);
166+
HttpMessage msg = getHttpMessage(id);
79167
if (msg == null) {
80168
continue;
81169
}
@@ -99,4 +187,45 @@ public SessionManagementRequestDetails findSessionTokenSource(String token, int
99187
}
100188
return null;
101189
}
190+
191+
private ParosDatabaseServer getParaosDataBaseServer() {
192+
if (Model.getSingleton().getDb().getDatabaseServer() instanceof ParosDatabaseServer pdbs) {
193+
pds = pdbs;
194+
LOGGER.info("PDS ? {}", pds != null);
195+
} else {
196+
if (pds == null && !warnedNonParosDb) {
197+
LOGGER.warn("Unexpected Database Server.");
198+
warnedNonParosDb = true;
199+
}
200+
}
201+
return pds;
202+
}
203+
204+
@Override
205+
public void sessionChanged(Session session) {
206+
pds = null;
207+
warnedNonParosDb = false;
208+
getParaosDataBaseServer();
209+
}
210+
211+
@Override
212+
public void sessionAboutToChange(Session session) {
213+
try {
214+
if (psGetHistory != null) {
215+
psGetHistory.close();
216+
}
217+
} catch (SQLException e) {
218+
// Nothing to do
219+
}
220+
}
221+
222+
@Override
223+
public void sessionScopeChanged(Session session) {
224+
// Nothing to do
225+
}
226+
227+
@Override
228+
public void sessionModeChanged(Mode mode) {
229+
// Nothing to do
230+
}
102231
}

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

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,4 +46,16 @@ public List<SessionToken> getTokens() {
4646
public int getConfidence() {
4747
return confidence;
4848
}
49+
50+
@Override
51+
public String toString() {
52+
return "MsgID: "
53+
+ msg.getHistoryRef().getHistoryId()
54+
+ " tokens ("
55+
+ tokens.size()
56+
+ "): "
57+
+ tokens
58+
+ " confidence: "
59+
+ confidence;
60+
}
4961
}

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -122,4 +122,9 @@ private static int compareStrings(String string, String otherString) {
122122
}
123123
return string.compareTo(otherString);
124124
}
125+
126+
@Override
127+
public String toString() {
128+
return "Source: " + source + " key: " + key + " value: " + value;
129+
}
125130
}

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

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,6 @@
2727
import java.util.regex.Matcher;
2828
import java.util.regex.Pattern;
2929
import lombok.Getter;
30-
import lombok.Setter;
3130
import org.apache.commons.httpclient.URIException;
3231
import org.apache.commons.lang3.StringUtils;
3332
import org.apache.logging.log4j.LogManager;
@@ -60,7 +59,7 @@ public final class ClientSideHandler implements HttpMessageHandler {
6059
private AuthRequestDetails authReq;
6160
private int firstHrefId;
6261

63-
@Setter private HistoryProvider historyProvider = new HistoryProvider();
62+
private HistoryProvider historyProvider = ExtensionAuthhelper.getHistoryProvider();
6463

6564
public ClientSideHandler(User user) {
6665
this.user = user;
@@ -224,6 +223,11 @@ protected static int messageTokenCount(HttpMessage msg, List<Pair<String, String
224223
return count;
225224
}
226225

226+
/** For testing purposes only, not part of the public API */
227+
void setHistoryProvider(HistoryProvider hp) {
228+
this.historyProvider = hp;
229+
}
230+
227231
@Getter
228232
class AuthRequestDetails {
229233
private HttpMessage msg;

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

Lines changed: 28 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -372,17 +372,6 @@ void shouldFindTokenWhenOneIsPreviouslyUnknown() throws Exception {
372372
extensionLoader =
373373
mock(ExtensionLoader.class, withSettings().strictness(Strictness.LENIENT));
374374

375-
history = new ArrayList<>();
376-
historyProvider = new TestHistoryProvider();
377-
AuthUtils.setHistoryProvider(historyProvider);
378-
379-
Control.initSingletonForTesting(model, extensionLoader);
380-
Model.setSingletonForTesting(model);
381-
382-
Session session = mock(Session.class);
383-
given(session.getContextsForUrl(anyString())).willReturn(Arrays.asList());
384-
given(model.getSession()).willReturn(session);
385-
386375
String cookie = "67890123456789012345";
387376
String jwtValue = "bearer 677890123456789012345-677890123456789012345";
388377
String jwt = "{\"jwt\":\"%s\"}".formatted(jwtValue);
@@ -416,6 +405,17 @@ void shouldFindTokenWhenOneIsPreviouslyUnknown() throws Exception {
416405
new HttpResponseHeader("HTTP/1.1 200 OK\r\n"),
417406
new HttpResponseBody("<html></html>"));
418407

408+
history = new ArrayList<>();
409+
historyProvider = new SuccessHistoryProvider(msg2);
410+
AuthUtils.setHistoryProvider(historyProvider);
411+
412+
Control.initSingletonForTesting(model, extensionLoader);
413+
Model.setSingletonForTesting(model);
414+
415+
Session session = mock(Session.class);
416+
given(session.getContextsForUrl(anyString())).willReturn(Arrays.asList());
417+
given(model.getSession()).willReturn(session);
418+
419419
List<HttpMessage> msgs = List.of(msg, msg2);
420420
historyProvider.addAuthMessageToHistory(msg);
421421
historyProvider.addAuthMessageToHistory(msg2);
@@ -493,4 +493,21 @@ public int getLastHistoryId() {
493493
return history.size() - 1;
494494
}
495495
}
496+
497+
class SuccessHistoryProvider extends TestHistoryProvider {
498+
499+
private HttpMessage msg;
500+
501+
SuccessHistoryProvider(HttpMessage msg) {
502+
this.msg = msg;
503+
}
504+
505+
@Override
506+
public SessionManagementRequestDetails findSessionTokenSource(String token, int firstId) {
507+
SessionToken st =
508+
new SessionToken(
509+
"json", "jwt", "bearer 677890123456789012345-677890123456789012345");
510+
return new SessionManagementRequestDetails(msg, List.of(st), 3);
511+
}
512+
}
496513
}

0 commit comments

Comments
 (0)