Skip to content

Commit

Permalink
Chrome-based browsers can do CDP-based script pinning (#13125)
Browse files Browse the repository at this point in the history
  • Loading branch information
shs96c authored Nov 9, 2023
1 parent 10adfe8 commit b5822ce
Show file tree
Hide file tree
Showing 4 changed files with 88 additions and 84 deletions.
4 changes: 4 additions & 0 deletions java/src/org/openqa/selenium/ScriptKey.java
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,10 @@ public ScriptKey(String identifier) {
this.identifier = Require.nonNull("Script ID", identifier);
}

public String getIdentifier() {
return identifier;
}

@Override
public boolean equals(Object o) {
if (!(o instanceof ScriptKey)) {
Expand Down
75 changes: 75 additions & 0 deletions java/src/org/openqa/selenium/chromium/ChromiumDriver.java
Original file line number Diff line number Diff line change
Expand Up @@ -23,9 +23,11 @@

import java.net.URI;
import java.net.URISyntaxException;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.function.Predicate;
import java.util.function.Supplier;
import java.util.logging.Level;
Expand All @@ -35,7 +37,9 @@
import org.openqa.selenium.Credentials;
import org.openqa.selenium.HasAuthentication;
import org.openqa.selenium.ImmutableCapabilities;
import org.openqa.selenium.JavascriptException;
import org.openqa.selenium.PersistentCapabilities;
import org.openqa.selenium.ScriptKey;
import org.openqa.selenium.WebDriver;
import org.openqa.selenium.WebDriverException;
import org.openqa.selenium.bidi.BiDi;
Expand All @@ -54,6 +58,7 @@
import org.openqa.selenium.html5.SessionStorage;
import org.openqa.selenium.html5.WebStorage;
import org.openqa.selenium.internal.Require;
import org.openqa.selenium.json.TypeToken;
import org.openqa.selenium.logging.EventType;
import org.openqa.selenium.logging.HasLogEvents;
import org.openqa.selenium.mobile.NetworkConnection;
Expand Down Expand Up @@ -102,6 +107,7 @@ public class ChromiumDriver extends RemoteWebDriver
private final Optional<BiDi> biDi;
protected HasCasting casting;
protected HasCdp cdp;
private final Map<Integer, ScriptKey> scriptKeys = new HashMap<>();

protected ChromiumDriver(
CommandExecutor commandExecutor, Capabilities capabilities, String capabilityKey) {
Expand Down Expand Up @@ -194,6 +200,75 @@ public Capabilities getCapabilities() {
return capabilities;
}

@Override
public ScriptKey pin(String script) {
Require.nonNull("Script to pin", script);

ScriptKey existingKey = scriptKeys.get(script.hashCode());
if (existingKey != null) {
return existingKey;
}

// Create the actual script we're going to use.
String scriptToUse =
String.format(
"window.seleniumPinnedScript%s = function(){%s}", Math.abs(script.hashCode()), script);

DevTools devTools = getDevTools();
devTools.createSessionIfThereIsNotOne();
devTools.send(new org.openqa.selenium.devtools.Command<>("Page.enable", Map.of()));
devTools.send(
new org.openqa.selenium.devtools.Command<>(
"Runtime.evaluate", Map.of("expression", scriptToUse)));
Map<String, Object> result =
devTools.send(
new org.openqa.selenium.devtools.Command<>(
"Page.addScriptToEvaluateOnNewDocument",
Map.of("source", scriptToUse),
new TypeToken<Map<String, Object>>() {}.getType()));

ScriptKey key = new ScriptKey((String) result.get("identifier"));
scriptKeys.put(script.hashCode(), key);
return key;
}

@Override
public Set<ScriptKey> getPinnedScripts() {
return Set.copyOf(scriptKeys.values());
}

@Override
public void unpin(ScriptKey key) {
int hashCode = getScriptId(key);

executeScript(String.format("window.seleniumPinnedScript%s = undefined", Math.abs(hashCode)));
scriptKeys.remove(hashCode);

DevTools devTools = getDevTools();
devTools.send(
new org.openqa.selenium.devtools.Command(
"Page.removeScriptToEvaluateOnNewDocument", Map.of("identifier", key.getIdentifier())));
}

@Override
public Object executeScript(ScriptKey key, Object... args) {
int hashCode = getScriptId(key);

String scriptToUse =
String.format(
"return window.seleniumPinnedScript%s.apply(window, arguments)", Math.abs(hashCode));

return this.executeScript(scriptToUse, args);
}

private int getScriptId(ScriptKey key) {
return scriptKeys.entrySet().stream()
.filter(entry -> key.equals(entry.getValue()))
.map(Map.Entry::getKey)
.findAny()
.orElseThrow(() -> new JavascriptException("Unable to find script key matching " + key));
}

@Override
public void setFileDetector(FileDetector detector) {
throw new WebDriverException(
Expand Down
84 changes: 0 additions & 84 deletions java/src/org/openqa/selenium/remote/RemoteWebDriver.java
Original file line number Diff line number Diff line change
Expand Up @@ -55,7 +55,6 @@
import org.openqa.selenium.HasCapabilities;
import org.openqa.selenium.HasDownloads;
import org.openqa.selenium.ImmutableCapabilities;
import org.openqa.selenium.JavascriptException;
import org.openqa.selenium.JavascriptExecutor;
import org.openqa.selenium.MutableCapabilities;
import org.openqa.selenium.NoAlertPresentException;
Expand All @@ -66,11 +65,9 @@
import org.openqa.selenium.Platform;
import org.openqa.selenium.Point;
import org.openqa.selenium.PrintsPage;
import org.openqa.selenium.ScriptKey;
import org.openqa.selenium.SearchContext;
import org.openqa.selenium.SessionNotCreatedException;
import org.openqa.selenium.TakesScreenshot;
import org.openqa.selenium.UnpinnedScriptKey;
import org.openqa.selenium.WebDriver;
import org.openqa.selenium.WebDriverException;
import org.openqa.selenium.WebElement;
Expand All @@ -86,7 +83,6 @@
import org.openqa.selenium.internal.Debug;
import org.openqa.selenium.internal.Require;
import org.openqa.selenium.io.Zip;
import org.openqa.selenium.json.TypeToken;
import org.openqa.selenium.logging.LocalLogs;
import org.openqa.selenium.logging.LoggingHandler;
import org.openqa.selenium.logging.Logs;
Expand Down Expand Up @@ -480,86 +476,6 @@ public Object executeAsyncScript(String script, Object... args) {
return execute(DriverCommand.EXECUTE_ASYNC_SCRIPT(script, convertedArgs)).getValue();
}

@Override
public ScriptKey pin(String script) {
UnpinnedScriptKey key = (UnpinnedScriptKey) JavascriptExecutor.super.pin(script);
String browserName = getCapabilities().getBrowserName().toLowerCase();
if ((browserName.equals("chrome")
|| browserName.equals("msedge")
|| browserName.equals("microsoftedge"))
&& this instanceof HasDevTools) {

((HasDevTools) this)
.maybeGetDevTools()
.ifPresent(
devTools -> {
devTools.createSessionIfThereIsNotOne();
devTools.send(
new org.openqa.selenium.devtools.Command<>("Page.enable", ImmutableMap.of()));
devTools.send(
new org.openqa.selenium.devtools.Command<>(
"Runtime.evaluate", ImmutableMap.of("expression", key.creationScript())));
Map<String, Object> result =
devTools.send(
new org.openqa.selenium.devtools.Command<>(
"Page.addScriptToEvaluateOnNewDocument",
ImmutableMap.of("source", key.creationScript()),
new TypeToken<Map<String, Object>>() {}.getType()));
key.setScriptId((String) result.get("identifier"));
});
}
return key;
}

@Override
public void unpin(ScriptKey scriptKey) {
UnpinnedScriptKey key = (UnpinnedScriptKey) scriptKey;

JavascriptExecutor.super.unpin(key);

String browserName = getCapabilities().getBrowserName().toLowerCase();
if ((browserName.equals("chrome")
|| browserName.equals("msedge")
|| browserName.equals("microsoftedge"))
&& this instanceof HasDevTools) {
((HasDevTools) this)
.maybeGetDevTools()
.ifPresent(
devTools -> {
devTools.send(
new org.openqa.selenium.devtools.Command<>("Page.enable", ImmutableMap.of()));
devTools.send(
new org.openqa.selenium.devtools.Command<>(
"Runtime.evaluate", ImmutableMap.of("expression", key.removalScript())));
devTools.send(
new org.openqa.selenium.devtools.Command<>(
"Page.removeScriptToEvaluateOnLoad",
ImmutableMap.of("identifier", key.getScriptId())));
});
}
}

@Override
public Object executeScript(ScriptKey key, Object... args) {
Require.stateCondition(
key instanceof UnpinnedScriptKey, "Script key should have been generated by this driver");

if (!getPinnedScripts().contains(key)) {
throw new JavascriptException("Script is unpinned");
}

String browserName = getCapabilities().getBrowserName().toLowerCase();

if ((browserName.equals("chrome")
|| browserName.equals("msedge")
|| browserName.equals("microsoftedge"))
&& this instanceof HasDevTools) {
return executeScript(((UnpinnedScriptKey) key).executionScript(), args);
}

return executeScript(((UnpinnedScriptKey) key).getScript(), args);
}

@Override
public TargetLocator switchTo() {
return new RemoteTargetLocator();
Expand Down
9 changes: 9 additions & 0 deletions java/test/org/openqa/selenium/ScriptPinningTest.java
Original file line number Diff line number Diff line change
Expand Up @@ -87,4 +87,13 @@ void callingAnUnpinnedScriptIsAnError() {
assertThatExceptionOfType(JavascriptException.class)
.isThrownBy(() -> executor.executeScript(cheese));
}

@Test
void afterPinningScriptShouldBeAvailableOnEveryPage() {
ScriptKey cheese = executor.pin("return 'havarti'");

driver.get(pages.xhtmlTestPage);

assertThat(executor.executeScript(cheese)).isEqualTo("havarti");
}
}

0 comments on commit b5822ce

Please sign in to comment.