Skip to content

Commit 9b18b17

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

File tree

12 files changed

+288
-64
lines changed

12 files changed

+288
-64
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+
historyProvider = hp;
1557+
}
15541558
}

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

Lines changed: 11 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;
113+
111114
private ZapMenuItem authTesterMenu;
112115
private AuthTestDialog authTestDialog;
113116

@@ -120,6 +123,13 @@ public ExtensionAuthhelper() {
120123
super();
121124
this.setI18nPrefix("authhelper");
122125
authHeaderTracker = new AuthHeaderTracker();
126+
ExtensionAuthhelper.historyProvider = new HistoryProvider();
127+
}
128+
129+
/** For testing purposes only, not part of the public API. */
130+
public ExtensionAuthhelper(HistoryProvider hp) {
131+
this();
132+
ExtensionAuthhelper.historyProvider = hp;
123133
}
124134

125135
@Override
@@ -212,6 +222,7 @@ public void destroy() {
212222
@Override
213223
public void hook(ExtensionHook extensionHook) {
214224
extensionHook.addSessionListener(new AuthSessionChangedListener());
225+
extensionHook.addSessionListener(historyProvider);
215226
extensionHook.addHttpSenderListener(authHeaderTracker);
216227
extensionHook.addOptionsParamSet(getParam());
217228
if (hasView()) {

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

Lines changed: 151 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -19,28 +19,63 @@
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

3850
private 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 RESBODY) > 0 -- URLEncoded
63+
OR POSITION(? IN RESBODY) > 0 -- JSONEscaped
64+
)
65+
ORDER BY HISTORYID DESC
66+
""";
67+
68+
private PreparedStatement psGetHistory;
69+
70+
private ParosDatabaseServer pds;
71+
private boolean warnedNonParosDb;
72+
4273
private ExtensionHistory extHist;
4374

75+
protected HistoryProvider() {
76+
getParosDataBaseServer();
77+
}
78+
4479
private ExtensionHistory getExtHistory() {
4580
if (extHist == null) {
4681
extHist = AuthUtils.getExtension(ExtensionHistory.class);
@@ -61,21 +96,90 @@ public HttpMessage getHttpMessage(int historyId)
6196
return null;
6297
}
6398

99+
/**
100+
* The query is ordered DESCending so the List and subsequent processing should be newest
101+
* message first.
102+
*/
103+
List<Integer> getMessageIds(int first, int last, String value) {
104+
Connection conn = getDbConnection();
105+
if (conn == null) {
106+
return List.of();
107+
}
108+
PreparedStatement query = getHistoryQuery(first, last, value, conn);
109+
if (query == null) {
110+
return List.of();
111+
}
112+
List<Integer> msgIds = new ArrayList<>();
113+
114+
try (ResultSet rs = query.executeQuery()) {
115+
while (rs.next()) {
116+
msgIds.add(rs.getInt("HISTORYID"));
117+
}
118+
} catch (SQLException e) {
119+
LOGGER.warn("Failed to process result set.");
120+
LOGGER.debug(e.getMessage(), e);
121+
} finally {
122+
try {
123+
query.close();
124+
} catch (SQLException e) {
125+
// Nothing to do
126+
}
127+
}
128+
LOGGER.debug("Found: {} candidate messages for {}", msgIds.size(), value);
129+
return msgIds;
130+
}
131+
132+
private PreparedStatement getHistoryQuery(int first, int last, String value, Connection conn) {
133+
try {
134+
if (psGetHistory == null || psGetHistory.isClosed()) {
135+
psGetHistory = conn.prepareStatement(QUERY_SESS_MGMT_TOKEN_MSG_IDS);
136+
}
137+
psGetHistory.setInt(1, first);
138+
psGetHistory.setInt(2, last);
139+
psGetHistory.setString(3, value);
140+
psGetHistory.setBytes(4, value.getBytes(StandardCharsets.UTF_8));
141+
psGetHistory.setString(5, value);
142+
psGetHistory.setBytes(
143+
6,
144+
URLEncoder.encode(value, StandardCharsets.UTF_8)
145+
.getBytes(StandardCharsets.UTF_8));
146+
psGetHistory.setBytes(
147+
7, StringEscapeUtils.escapeJson(value).getBytes(StandardCharsets.UTF_8));
148+
} catch (SQLException e) {
149+
LOGGER.warn("Failed to prepare query.", e);
150+
}
151+
return psGetHistory;
152+
}
153+
154+
private Connection getDbConnection() {
155+
if (pds == null) {
156+
LOGGER.info("PDS was null");
157+
return null;
158+
}
159+
Connection conn = null;
160+
try {
161+
conn = pds.getSingletonConnection();
162+
} catch (SQLException | NullPointerException e) {
163+
LOGGER.warn("Failed to get DB connection.", e);
164+
}
165+
return conn;
166+
}
167+
64168
public int getLastHistoryId() {
65169
return getExtHistory().getLastHistoryId();
66170
}
67171

68172
public SessionManagementRequestDetails findSessionTokenSource(String token, int firstId) {
69173
int lastId = getLastHistoryId();
70174
if (firstId == -1) {
71-
firstId = Math.max(0, lastId - MAX_NUM_RECORDS_TO_CHECK);
175+
firstId = Math.max(1, lastId - MAX_NUM_RECORDS_TO_CHECK);
72176
}
73177

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

76-
for (int i = lastId; i >= firstId; i--) {
180+
for (int id : getMessageIds(firstId, lastId, token)) {
77181
try {
78-
HttpMessage msg = getHttpMessage(i);
182+
HttpMessage msg = getHttpMessage(id);
79183
if (msg == null) {
80184
continue;
81185
}
@@ -99,4 +203,47 @@ public SessionManagementRequestDetails findSessionTokenSource(String token, int
99203
}
100204
return null;
101205
}
206+
207+
private ParosDatabaseServer getParosDataBaseServer() {
208+
if (Model.getSingleton().getDb() == null) {
209+
return null; // This should only happen in unit tests
210+
}
211+
if (Model.getSingleton().getDb().getDatabaseServer() instanceof ParosDatabaseServer pdbs) {
212+
pds = pdbs;
213+
} else {
214+
if (pds == null && !warnedNonParosDb) {
215+
LOGGER.warn("Unexpected Database Server.");
216+
warnedNonParosDb = true;
217+
}
218+
}
219+
return pds;
220+
}
221+
222+
@Override
223+
public void sessionChanged(Session session) {
224+
pds = null;
225+
warnedNonParosDb = false;
226+
getParosDataBaseServer();
227+
}
228+
229+
@Override
230+
public void sessionAboutToChange(Session session) {
231+
try {
232+
if (psGetHistory != null) {
233+
psGetHistory.close();
234+
}
235+
} catch (SQLException e) {
236+
// Nothing to do
237+
}
238+
}
239+
240+
@Override
241+
public void sessionScopeChanged(Session session) {
242+
// Nothing to do
243+
}
244+
245+
@Override
246+
public void sessionModeChanged(Mode mode) {
247+
// Nothing to do
248+
}
102249
}

‎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: 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/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

0 commit comments

Comments
 (0)