Skip to content

Commit c9b10de

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

File tree

11 files changed

+210
-70
lines changed

11 files changed

+210
-70
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: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -199,7 +199,8 @@ public void notifyMessageReceived(HttpMessage message) {
199199

200200
private static long timeToWaitMs = TimeUnit.SECONDS.toMillis(5);
201201

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

204205
/**
205206
* 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‎

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: 101 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -19,26 +19,55 @@
1919
*/
2020
package org.zaproxy.addon.authhelper;
2121

22+
import java.net.URLEncoder;
23+
import java.nio.charset.StandardCharsets;
24+
import java.sql.Connection;
25+
import java.sql.PreparedStatement;
26+
import java.sql.ResultSet;
27+
import java.sql.SQLException;
2228
import java.util.ArrayList;
2329
import java.util.List;
2430
import java.util.Optional;
31+
import org.apache.commons.text.StringEscapeUtils;
2532
import org.apache.logging.log4j.LogManager;
2633
import org.apache.logging.log4j.Logger;
34+
import org.parosproxy.paros.control.Control.Mode;
2735
import org.parosproxy.paros.core.scanner.Alert;
2836
import org.parosproxy.paros.db.DatabaseException;
37+
import org.parosproxy.paros.db.paros.ParosDatabaseServer;
38+
import org.parosproxy.paros.extension.SessionChangedListener;
2939
import org.parosproxy.paros.extension.history.ExtensionHistory;
3040
import org.parosproxy.paros.model.HistoryReference;
41+
import org.parosproxy.paros.model.Model;
42+
import org.parosproxy.paros.model.Session;
3143
import org.parosproxy.paros.network.HttpMalformedHeaderException;
3244
import org.parosproxy.paros.network.HttpMessage;
3345
import org.zaproxy.zap.authentication.AuthenticationHelper;
3446

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

38-
private static final int MAX_NUM_RECORDS_TO_CHECK = 200;
50+
protected static final int MAX_NUM_RECORDS_TO_CHECK = 200;
3951

4052
private static final Logger LOGGER = LogManager.getLogger(HistoryProvider.class);
4153

54+
private static final String QUERY_SESS_MGMT_TOKEN_MSG_IDS =
55+
"""
56+
SELECT HISTORYID FROM HISTORY
57+
WHERE HISTORYID BETWEEN ? AND ?
58+
AND (
59+
POSITION(? IN RESHEADER) > 0
60+
OR POSITION(? IN RESBODY) > 0
61+
OR POSITION(? IN REQHEADER) > 0
62+
OR POSITION(? IN REQHEADER) > 0 -- URLEncoded
63+
OR POSITION(? IN RESBODY) > 0 -- JSONEscaped
64+
)
65+
ORDER BY HISTORYID DESC
66+
""";
67+
68+
private ParosDatabaseServer pds;
69+
private boolean server;
70+
4271
private ExtensionHistory extHist;
4372

4473
private ExtensionHistory getExtHistory() {
@@ -61,21 +90,67 @@ public HttpMessage getHttpMessage(int historyId)
6190
return null;
6291
}
6392

93+
/**
94+
* The query is ordered DESCending so the List and subsequent processing should be newest
95+
* message first.
96+
*/
97+
List<Integer> getMessageIds(int first, int last, String value) {
98+
if (!server) {
99+
if (Model.getSingleton().getDb().getDatabaseServer()
100+
instanceof ParosDatabaseServer pdbs) {
101+
server = true;
102+
pds = pdbs;
103+
} else {
104+
LOGGER.warn("Unsupported Database Server.");
105+
}
106+
}
107+
if (pds == null) {
108+
return List.of();
109+
}
110+
111+
try (Connection conn = pds.getSingletonConnection();
112+
PreparedStatement psGetHistory =
113+
conn.prepareStatement(QUERY_SESS_MGMT_TOKEN_MSG_IDS)) {
114+
psGetHistory.setInt(1, first);
115+
psGetHistory.setInt(2, last);
116+
psGetHistory.setString(3, value);
117+
psGetHistory.setBytes(4, value.getBytes(StandardCharsets.UTF_8));
118+
psGetHistory.setString(5, value);
119+
psGetHistory.setString(6, URLEncoder.encode(value, StandardCharsets.UTF_8));
120+
psGetHistory.setBytes(
121+
7, StringEscapeUtils.escapeJson(value).getBytes(StandardCharsets.UTF_8));
122+
123+
List<Integer> msgIds = new ArrayList<>();
124+
try (ResultSet rs = psGetHistory.executeQuery()) {
125+
while (rs.next()) {
126+
msgIds.add(rs.getInt("HISTORYID"));
127+
}
128+
} catch (SQLException e) {
129+
LOGGER.warn("Failed to process result set. {}", e.getMessage());
130+
}
131+
LOGGER.debug("Found: {} candidate messages for {}", msgIds.size(), value);
132+
return msgIds;
133+
} catch (SQLException e) {
134+
LOGGER.warn("Failed to prepare query.", e);
135+
return List.of();
136+
}
137+
}
138+
64139
public int getLastHistoryId() {
65140
return getExtHistory().getLastHistoryId();
66141
}
67142

68143
public SessionManagementRequestDetails findSessionTokenSource(String token, int firstId) {
69144
int lastId = getLastHistoryId();
70145
if (firstId == -1) {
71-
firstId = Math.max(0, lastId - MAX_NUM_RECORDS_TO_CHECK);
146+
firstId = Math.max(1, lastId - MAX_NUM_RECORDS_TO_CHECK);
72147
}
73148

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

76-
for (int i = lastId; i >= firstId; i--) {
151+
for (int id : getMessageIds(firstId, lastId, token)) {
77152
try {
78-
HttpMessage msg = getHttpMessage(i);
153+
HttpMessage msg = getHttpMessage(id);
79154
if (msg == null) {
80155
continue;
81156
}
@@ -99,4 +174,25 @@ public SessionManagementRequestDetails findSessionTokenSource(String token, int
99174
}
100175
return null;
101176
}
177+
178+
@Override
179+
public void sessionChanged(Session session) {
180+
pds = null;
181+
server = false;
182+
}
183+
184+
@Override
185+
public void sessionAboutToChange(Session session) {
186+
// Nothing to do
187+
}
188+
189+
@Override
190+
public void sessionScopeChanged(Session session) {
191+
// Nothing to do
192+
}
193+
194+
@Override
195+
public void sessionModeChanged(Mode mode) {
196+
// Nothing to do
197+
}
102198
}

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

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,4 +46,20 @@ public List<SessionToken> getTokens() {
4646
public int getConfidence() {
4747
return confidence;
4848
}
49+
50+
@Override
51+
public String toString() {
52+
String historyIdStr =
53+
msg.getHistoryRef() != null
54+
? String.valueOf(msg.getHistoryRef().getHistoryId())
55+
: "N/A";
56+
return "MsgID: "
57+
+ historyIdStr
58+
+ " tokens ("
59+
+ tokens.size()
60+
+ "): "
61+
+ tokens
62+
+ " confidence: "
63+
+ confidence;
64+
}
4965
}

‎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: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -60,7 +60,7 @@ public final class ClientSideHandler implements HttpMessageHandler {
6060
private AuthRequestDetails authReq;
6161
private int firstHrefId;
6262

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

6565
public ClientSideHandler(User user) {
6666
this.user = user;

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -106,6 +106,7 @@ void setUp() throws Exception {
106106
setUpZap();
107107

108108
mockMessages(new ExtensionAuthhelper());
109+
AuthUtils.setHistoryProvider(new TestHistoryProvider());
109110
}
110111

111112
@AfterEach

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

Lines changed: 5 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -43,9 +43,7 @@
4343
import org.parosproxy.paros.Constant;
4444
import org.parosproxy.paros.control.Control;
4545
import org.parosproxy.paros.core.scanner.Alert;
46-
import org.parosproxy.paros.db.DatabaseException;
4746
import org.parosproxy.paros.extension.ExtensionLoader;
48-
import org.parosproxy.paros.model.HistoryReference;
4947
import org.parosproxy.paros.model.Model;
5048
import org.parosproxy.paros.model.Session;
5149
import org.parosproxy.paros.network.HttpHeader;
@@ -67,7 +65,6 @@
6765
/** Unit test for {@link SessionDetectionScanRule}. */
6866
class SessionDetectionScanRuleUnitTest extends PassiveScannerTest<SessionDetectionScanRule> {
6967

70-
private List<HttpMessage> history;
7168
private HistoryProvider historyProvider;
7269

7370
@Override
@@ -107,7 +104,6 @@ void shouldSetHeaderBasedSessionManagment() throws Exception {
107104
given(session.getContextsForUrl(anyString())).willReturn(Arrays.asList(context));
108105
given(model.getSession()).willReturn(session);
109106

110-
history = new ArrayList<>();
111107
historyProvider = new TestHistoryProvider();
112108
AuthUtils.setHistoryProvider(historyProvider);
113109

@@ -204,7 +200,6 @@ void shouldDetectBodgeitSession() throws Exception {
204200
DiagnosticDataLoader.loadTestData(
205201
this.getResourcePath("internal/bodgeit.diags").toFile());
206202

207-
history = new ArrayList<>();
208203
historyProvider = new TestHistoryProvider();
209204
AuthUtils.setHistoryProvider(historyProvider);
210205
// When
@@ -228,7 +223,6 @@ void shouldDetectBodgeitSession() throws Exception {
228223
@Test
229224
void shouldDetectCtflearnSession() throws Exception {
230225
// Given
231-
history = new ArrayList<>();
232226
historyProvider = new TestHistoryProvider();
233227
AuthUtils.setHistoryProvider(historyProvider);
234228

@@ -267,7 +261,6 @@ void shouldDetectCtflearnSession() throws Exception {
267261
@Test
268262
void shouldDetectDefthewebSession() throws Exception {
269263
// Given
270-
history = new ArrayList<>();
271264
historyProvider = new TestHistoryProvider();
272265
AuthUtils.setHistoryProvider(historyProvider);
273266

@@ -295,7 +288,6 @@ void shouldDetectDefthewebSession() throws Exception {
295288
@Test
296289
void shouldDetectGinnjuiceSession() throws Exception {
297290
// Given
298-
history = new ArrayList<>();
299291
historyProvider = new TestHistoryProvider();
300292
AuthUtils.setHistoryProvider(historyProvider);
301293

@@ -325,7 +317,6 @@ void shouldDetectGinnjuiceSession() throws Exception {
325317
@Test
326318
void shouldDetectInfosecexSession() throws Exception {
327319
// Given
328-
history = new ArrayList<>();
329320
historyProvider = new TestHistoryProvider();
330321
AuthUtils.setHistoryProvider(historyProvider);
331322

@@ -372,10 +363,6 @@ void shouldFindTokenWhenOneIsPreviouslyUnknown() throws Exception {
372363
extensionLoader =
373364
mock(ExtensionLoader.class, withSettings().strictness(Strictness.LENIENT));
374365

375-
history = new ArrayList<>();
376-
historyProvider = new TestHistoryProvider();
377-
AuthUtils.setHistoryProvider(historyProvider);
378-
379366
Control.initSingletonForTesting(model, extensionLoader);
380367
Model.setSingletonForTesting(model);
381368

@@ -416,7 +403,11 @@ void shouldFindTokenWhenOneIsPreviouslyUnknown() throws Exception {
416403
new HttpResponseHeader("HTTP/1.1 200 OK\r\n"),
417404
new HttpResponseBody("<html></html>"));
418405

419-
List<HttpMessage> msgs = List.of(msg, msg2);
406+
List<HttpMessage> msgs = new ArrayList<>();
407+
msgs.add(msg);
408+
msgs.add(msg2);
409+
historyProvider = new TestHistoryProvider(msgs);
410+
AuthUtils.setHistoryProvider(historyProvider);
420411
historyProvider.addAuthMessageToHistory(msg);
421412
historyProvider.addAuthMessageToHistory(msg2);
422413
AuthUtils.recordSessionToken(
@@ -470,27 +461,4 @@ void shouldIgnoreUnknownOrUnwantedContentTypes(String contentType)
470461
// Then
471462
assertThat(alertsRaised, hasSize(0));
472463
}
473-
474-
class TestHistoryProvider extends HistoryProvider {
475-
@Override
476-
public void addAuthMessageToHistory(HttpMessage msg) {
477-
history.add(msg);
478-
int id = history.size() - 1;
479-
HistoryReference href =
480-
mock(HistoryReference.class, withSettings().strictness(Strictness.LENIENT));
481-
given(href.getHistoryId()).willReturn(id);
482-
msg.setHistoryRef(href);
483-
}
484-
485-
@Override
486-
public HttpMessage getHttpMessage(int historyId)
487-
throws HttpMalformedHeaderException, DatabaseException {
488-
return history.get(historyId);
489-
}
490-
491-
@Override
492-
public int getLastHistoryId() {
493-
return history.size() - 1;
494-
}
495-
}
496464
}

0 commit comments

Comments
 (0)