30
30
import java .util .List ;
31
31
import java .util .Map ;
32
32
import java .util .Optional ;
33
+ import java .util .concurrent .CancellationException ;
33
34
import java .util .concurrent .CompletableFuture ;
34
35
import java .util .concurrent .ConcurrentHashMap ;
35
36
import java .util .concurrent .ExecutionException ;
45
46
import org .openqa .selenium .devtools .DevTools ;
46
47
import org .openqa .selenium .devtools .DevToolsException ;
47
48
import org .openqa .selenium .devtools .Event ;
49
+ import org .openqa .selenium .devtools .NetworkInterceptor ;
48
50
import org .openqa .selenium .internal .Either ;
49
51
import org .openqa .selenium .internal .Require ;
50
52
import org .openqa .selenium .remote .http .Contents ;
@@ -57,21 +59,33 @@ public abstract class Network<AUTHREQUIRED, REQUESTPAUSED> {
57
59
58
60
private static final Logger LOG = Logger .getLogger (Network .class .getName ());
59
61
62
+ private static final HttpResponse STOP_PROCESSING =
63
+ new HttpResponse ()
64
+ .addHeader ("Selenium-Interceptor" , "Stop" )
65
+ .setContent (Contents .utf8String ("Interception is stopped" ));
66
+
60
67
private final Map <Predicate <URI >, Supplier <Credentials >> authHandlers = new LinkedHashMap <>();
61
68
private final Filter defaultFilter = next -> next ::execute ;
62
69
private volatile Filter filter = defaultFilter ;
63
70
protected final DevTools devTools ;
64
71
65
72
private final AtomicBoolean fetchEnabled = new AtomicBoolean ();
73
+ private final Map <String , CompletableFuture <HttpResponse >> pendingResponses =
74
+ new ConcurrentHashMap <>();
66
75
67
76
public Network (DevTools devtools ) {
68
77
this .devTools = Require .nonNull ("DevTools" , devtools );
69
78
}
70
79
71
80
public void disable () {
72
81
fetchEnabled .set (false );
73
- devTools .send (disableFetch ());
74
- devTools .send (enableNetworkCaching ());
82
+ try {
83
+ devTools .send (disableFetch ());
84
+ devTools .send (enableNetworkCaching ());
85
+ } finally {
86
+ // we stopped the fetch we will not receive any pending responses
87
+ pendingResponses .values ().forEach (cf -> cf .cancel (false ));
88
+ }
75
89
76
90
synchronized (authHandlers ) {
77
91
authHandlers .clear ();
@@ -183,8 +197,6 @@ public void prepareToInterceptTraffic() {
183
197
devTools .send (cancelAuth (authRequired ));
184
198
});
185
199
186
- Map <String , CompletableFuture <HttpResponse >> responses = new ConcurrentHashMap <>();
187
-
188
200
devTools .addListener (
189
201
requestPausedEvent (),
190
202
pausedRequest -> {
@@ -194,7 +206,7 @@ public void prepareToInterceptTraffic() {
194
206
195
207
if (message .isRight ()) {
196
208
HttpResponse res = message .right ();
197
- CompletableFuture <HttpResponse > future = responses .remove (id );
209
+ CompletableFuture <HttpResponse > future = pendingResponses .remove (id );
198
210
199
211
if (future == null ) {
200
212
devTools .send (continueWithoutModification (pausedRequest ));
@@ -210,18 +222,22 @@ public void prepareToInterceptTraffic() {
210
222
.andFinally (
211
223
req -> {
212
224
// Convert the selenium request to a CDP one and fulfill.
213
-
214
- CompletableFuture <HttpResponse > res = new CompletableFuture <>();
215
- responses .put (id , res );
216
-
217
225
devTools .send (continueRequest (pausedRequest , req ));
226
+ CompletableFuture <HttpResponse > res = new CompletableFuture <>();
227
+ // Save the future after the browser accepted the continueRequest
228
+ pendingResponses .put (id , res );
218
229
219
230
// Wait for the CDP response and send that back.
220
231
try {
221
232
return res .get ();
222
233
} catch (InterruptedException e ) {
223
234
Thread .currentThread ().interrupt ();
224
235
throw new WebDriverException (e );
236
+ } catch (CancellationException e ) {
237
+ // The interception was intentionally stopped, network().disable() has
238
+ // been called
239
+ pendingResponses .remove (id );
240
+ return STOP_PROCESSING ;
225
241
} catch (ExecutionException e ) {
226
242
if (fetchEnabled .get ()) {
227
243
LOG .log (WARNING , e , () -> "Unable to process request" );
@@ -231,9 +247,12 @@ public void prepareToInterceptTraffic() {
231
247
})
232
248
.execute (message .left ());
233
249
234
- if ("Continue" . equals ( forBrowser . getHeader ( "Selenium-Interceptor" )) ) {
250
+ if (forBrowser == NetworkInterceptor . PROCEED_WITH_REQUEST ) {
235
251
devTools .send (continueWithoutModification (pausedRequest ));
236
252
return ;
253
+ } else if (forBrowser == STOP_PROCESSING ) {
254
+ // The interception was intentionally stopped, network().disable() has been called
255
+ return ;
237
256
}
238
257
239
258
devTools .send (fulfillRequest (pausedRequest , forBrowser ));
0 commit comments