Skip to content

Commit

Permalink
[Grid] Support auto downloads in Grid
Browse files Browse the repository at this point in the history
Fixes #11656 #11658

Following has been done:

* Specify the default base directory into which 
all downloads at a node will go into via the flag 
“—-base-dir-downloads”. If this flag does not have 
a value then we default to user’s home.
* Turn ON managing download folders via the flag 
“-—enable-manage-downloads”
* Enabled support for Chrome|Edge|Firefox browsers.
* File downloads will be done only in a session aware 
directory for a given web driver session. After session
is killed, the directory gets cleaned up as well.
  • Loading branch information
krmahadevan committed Mar 7, 2023
1 parent 4aee416 commit b646a06
Show file tree
Hide file tree
Showing 14 changed files with 625 additions and 132 deletions.
20 changes: 20 additions & 0 deletions java/src/org/openqa/selenium/grid/data/DefaultSlotMatcher.java
Original file line number Diff line number Diff line change
Expand Up @@ -64,6 +64,10 @@ public boolean matches(Capabilities stereotype, Capabilities capabilities) {
return false;
}

if (!autoDownloadsMatch(stereotype, capabilities)) {
return false;
}

if (!platformVersionMatch(stereotype, capabilities)) {
return false;
}
Expand Down Expand Up @@ -106,6 +110,22 @@ private Boolean initialMatch(Capabilities stereotype, Capabilities capabilities)
.orElse(true);
}

private Boolean autoDownloadsMatch(Capabilities stereotype, Capabilities capabilities) {
//First lets check if user wanted auto downloads
Object raw = capabilities.getCapability("se:enableDownloads");
if (raw == null || !Boolean.parseBoolean(raw.toString())) {
//User didn't ask. So lets move on to the next matching criteria
return true;
}
//User wants auto downloads to be done on this browser flavor.
raw = stereotype.getCapability("se:enableDownloads");
if (raw == null || !Boolean.parseBoolean(raw.toString())) {
//User wants it, we don't have it. So no match
return false;
}
return true;
}

private Boolean platformVersionMatch(Capabilities stereotype, Capabilities capabilities) {
/*
This platform version match is not W3C compliant but users can add Appium servers as
Expand Down
21 changes: 15 additions & 6 deletions java/src/org/openqa/selenium/grid/node/config/NodeFlags.java
Original file line number Diff line number Diff line change
Expand Up @@ -230,12 +230,21 @@ public class NodeFlags implements HasRoles {
private String nodeImplementation = DEFAULT_NODE_IMPLEMENTATION;

@Parameter(
names = {"--downloads-path"},
description = "The default location wherein all browser triggered file downloads would be "
+ "available to be retrieved from. This is usually the directory that you configure in "
+ "your browser as the default location for storing downloaded files.")
@ConfigValue(section = NODE_SECTION, name = "downloads-path", example = "")
private String downloadsPath = "";
names = {"--enable-manage-downloads"},
description = "When enabled, the Grid node will automatically do the following: \n" +
"\t1. Creates a directory named '$HOME/.cache/selenium/downloads/' which will now represent the directory "
+ "into which files downloaded by Chrome/Firefox/Edge browser will be under.\n" +
"\t2. For every new session, a sub-directory will be created/deleted so that all files that were "
+ "downloaded for a given session are stored in.")
@ConfigValue(section = NODE_SECTION, name = "enable-manage-downloads", example = "false")
public Boolean enableManageDownloads;

@Parameter(
names = {"--base-dir-downloads"},
description = "The base directory into which all downloads would be saved. If no value is specified "
+ "then this defaults to the user's home directory.")
@ConfigValue(section = NODE_SECTION, name = "enable-manage-downloads", example = "user.home")
public String baseDirectory;

@Override
public Set<Role> getRoles() {
Expand Down
18 changes: 15 additions & 3 deletions java/src/org/openqa/selenium/grid/node/config/NodeOptions.java
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,7 @@
import org.openqa.selenium.json.JsonOutput;
import org.openqa.selenium.net.NetworkUtils;
import org.openqa.selenium.net.Urls;
import org.openqa.selenium.remote.Browser;
import org.openqa.selenium.remote.service.DriverService;

import java.io.File;
Expand Down Expand Up @@ -149,8 +150,14 @@ public Optional<URI> getPublicGridUri() {
}
}

public Optional<String> getDownloadsPath() {
return config.get(NODE_SECTION, "downloads-path");
public boolean isAutoManageDownloadsFolder() {
return config.getBool(NODE_SECTION, "enable-manage-downloads")
.orElse(Boolean.FALSE);
}

public String baseDirForDownloads() {
return config.get(NODE_SECTION, "base-dir-downloads")
.orElse(System.getProperty("user.home"));
}

public Node getNode() {
Expand Down Expand Up @@ -216,8 +223,9 @@ public Map<Capabilities, Collection<SessionFactory>> getSessionFactories(
addDriverConfigs(factoryFactory, sessionFactories);
addSpecificDrivers(allDrivers, sessionFactories);
addDetectedDrivers(allDrivers, sessionFactories);
ImmutableMultimap<Capabilities, SessionFactory> built = sessionFactories.build();

return sessionFactories.build().asMap();
return built.asMap();
}

public int getMaxSessions() {
Expand Down Expand Up @@ -634,6 +642,10 @@ private Capabilities enhanceStereotype(Capabilities capabilities) {
.setCapability("se:vncEnabled", true)
.setCapability("se:noVncPort", noVncPort());
}
if (isAutoManageDownloadsFolder() && Browser.honoursSpecifiedDownloadsDir(capabilities)) {
capabilities = new PersistentCapabilities(capabilities)
.setCapability("se:enableDownloads", true);
}
return capabilities;
}

Expand Down
154 changes: 139 additions & 15 deletions java/src/org/openqa/selenium/grid/node/local/LocalNode.java
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,14 @@
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;

import java.io.Serializable;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.Arrays;
import java.util.Comparator;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Stream;
import org.openqa.selenium.Capabilities;
import org.openqa.selenium.ImmutableCapabilities;
import org.openqa.selenium.MutableCapabilities;
Expand Down Expand Up @@ -63,6 +70,7 @@
import org.openqa.selenium.io.TemporaryFilesystem;
import org.openqa.selenium.io.Zip;
import org.openqa.selenium.json.Json;
import org.openqa.selenium.remote.Browser;
import org.openqa.selenium.remote.SessionId;
import org.openqa.selenium.remote.http.HttpMethod;
import org.openqa.selenium.remote.http.HttpRequest;
Expand Down Expand Up @@ -116,6 +124,7 @@ public class LocalNode extends Node {

private static final Json JSON = new Json();
private static final Logger LOG = Logger.getLogger(LocalNode.class.getName());

private final EventBus bus;
private final URI externalUri;
private final URI gridUri;
Expand All @@ -124,7 +133,7 @@ public class LocalNode extends Node {
private final int maxSessionCount;
private final int configuredSessionCount;
private final boolean cdpEnabled;
private final String downloadsPath;
private final boolean enableManageDownloads;

private final boolean bidiEnabled;
private final AtomicBoolean drainAfterSessions = new AtomicBoolean();
Expand All @@ -133,6 +142,10 @@ public class LocalNode extends Node {
private final Cache<SessionId, TemporaryFilesystem> tempFileSystems;
private final AtomicInteger pendingSessions = new AtomicInteger();
private final AtomicInteger sessionCount = new AtomicInteger();
private final Map<SessionId,UUID> downloadFolderToSessionMapping = new ConcurrentHashMap<>();
private final String baseDir;

private File defaultDownloadsDir;

protected LocalNode(
Tracer tracer,
Expand All @@ -149,7 +162,8 @@ protected LocalNode(
Duration heartbeatPeriod,
List<SessionSlot> factories,
Secret registrationSecret,
String downloadsPath) {
String baseDir,
boolean enableManageDownloads) {
super(tracer, new NodeId(UUID.randomUUID()), uri, registrationSecret);

this.bus = Require.nonNull("Event bus", bus);
Expand All @@ -166,7 +180,8 @@ protected LocalNode(
this.sessionCount.set(drainAfterSessionCount);
this.cdpEnabled = cdpEnabled;
this.bidiEnabled = bidiEnabled;
this.downloadsPath = Optional.ofNullable(downloadsPath).orElse("");
this.enableManageDownloads = enableManageDownloads;
this.baseDir = baseDir;

this.healthCheck = healthCheck == null ?
() -> {
Expand Down Expand Up @@ -283,6 +298,11 @@ public int getCurrentSessionCount() {
return Math.toIntExact(currentSessions.size());
}

@VisibleForTesting
public String getDownloadsDirForSession(SessionId sessionId) {
return downloadFolderToSessionMapping.get(sessionId).toString();
}

@ManagedAttribute(name = "MaxSessions")
public int getMaxSessionCount() {
return maxSessionCount;
Expand Down Expand Up @@ -381,10 +401,19 @@ public Either<WebDriverException, CreateSessionResponse> newSession(CreateSessio
new RetrySessionRequestException("No slot matched the requested capabilities."));
}

UUID folderIdForPossibleSession = UUID.randomUUID();
Capabilities desiredCapabilities = sessionRequest.getDesiredCapabilities();
Capabilities enriched = enrichWithDownloadsPathInfo(folderIdForPossibleSession, desiredCapabilities);
enriched = desiredCapabilities.merge(enriched);
sessionRequest = new CreateSessionRequest(sessionRequest.getDownstreamDialects(), enriched, sessionRequest.getMetadata());

Either<WebDriverException, ActiveSession> possibleSession = slotToUse.apply(sessionRequest);

if (possibleSession.isRight()) {
ActiveSession session = possibleSession.right();
downloadFolderToSessionMapping.put(session.getId(), folderIdForPossibleSession);
LOG.info("Downloads pertaining to Session Id [" + session.getId().toString() + "] will be " +
"saved to " + defaultDownloadsBaseFolder().getAbsolutePath() + "/" + folderIdForPossibleSession);
currentSessions.put(session.getId(), slotToUse);

checkSessionCount();
Expand All @@ -407,7 +436,7 @@ public Either<WebDriverException, CreateSessionResponse> newSession(CreateSessio
externalUri,
slotToUse.isSupportingCdp(),
slotToUse.isSupportingBiDi(),
sessionRequest.getDesiredCapabilities());
desiredCapabilities);

String sessionCreatedMessage = "Session created by the Node";
LOG.info(String.format("%s. Id: %s, Caps: %s", sessionCreatedMessage, sessionId, caps));
Expand All @@ -425,6 +454,57 @@ public Either<WebDriverException, CreateSessionResponse> newSession(CreateSessio
}
}

private Capabilities enrichWithDownloadsPathInfo(UUID id, Capabilities caps) {
if (!enableManageDownloads || !Browser.honoursSpecifiedDownloadsDir(caps)) {
// Auto enable of downloads is not turned on. So don't bother fiddling around
// with the capabilities.
return caps;
}

File subDir = new File(defaultDownloadsBaseFolder(), id.toString());
boolean created = subDir.mkdirs();
if (created) {
LOG.fine("Created folder " + subDir.getAbsolutePath() + " for auto downloads");
} else {
//Should we error out here because if the downloads folder can't be created, then
//there's a good chance that the downloads may not work too.
LOG.warning("Could not create folder " + subDir.getAbsolutePath() + " for auto downloads");
}
if (Browser.CHROME.is(caps)) {
ImmutableMap<String, Serializable> map = ImmutableMap.of(
"download.prompt_for_download", false,
"download.default_directory", subDir.getAbsolutePath());
appendPrefs(caps, "goog:chromeOptions", map);
return caps;
}

if (Browser.EDGE.is(caps)) {
ImmutableMap<String, Serializable> map = ImmutableMap.of(
"download.prompt_for_download", false,
"download.default_directory", subDir.getAbsolutePath());
appendPrefs(caps, "ms:edgeOptions", map);
return caps;
}

if (Browser.FIREFOX.is(caps)) {
ImmutableMap<String, Serializable> map = ImmutableMap.of(
"browser.download.folderList", 2,
"browser.download.dir", subDir.getAbsolutePath());
appendPrefs(caps, "moz:firefoxOptions", map);
}
return caps;
}

@SuppressWarnings("unchecked")
private Capabilities appendPrefs(Capabilities caps, String optionsKey, Map<String, Serializable> map) {
Map<String, Object> currentOptions = (Map<String, Object>) Optional.ofNullable(
caps.getCapability(optionsKey))
.orElse(new HashMap<>());

((Map<String, Serializable>)currentOptions.computeIfAbsent("prefs",k -> new HashMap<>())).putAll(map);
return caps;
}

@Override
public boolean isSessionOwner(SessionId id) {
Require.nonNull("Session ID", id);
Expand Down Expand Up @@ -484,18 +564,19 @@ public HttpResponse downloadFile(HttpRequest req, SessionId id) {
if (slot != null && slot.getSession() instanceof DockerSession) {
return executeWebDriverCommand(req);
}
if (this.downloadsPath.isEmpty()) {
String msg = "Please specify the path wherein the files downloaded using the browser "
+ "would be available via the command line arg [--downloads-path] and restart the node";
if (!this.enableManageDownloads) {
String msg = "Please enable management of downloads via the command line arg "
+ "[--enable-manage-downloads] and restart the node";
throw new WebDriverException(msg);
}
File dir = new File(this.downloadsPath);
File dir = new File(defaultDownloadsBaseFolder(),
downloadFolderToSessionMapping.get(id).toString());
if (!dir.exists()) {
throw new WebDriverException(
String.format("Cannot locate downloads directory %s.", downloadsPath));
String.format("Cannot locate downloads directory %s.", dir));
}
if (!dir.isDirectory()) {
throw new WebDriverException(String.format("Invalid directory: %s.", downloadsPath));
throw new WebDriverException(String.format("Invalid directory: %s.", dir));
}
if (req.getMethod().equals(HttpMethod.GET)) {
//User wants to list files that can be downloaded
Expand Down Expand Up @@ -525,7 +606,7 @@ public HttpResponse downloadFile(HttpRequest req, SessionId id) {
).orElse(new File[]{});
if (allFiles.length == 0) {
throw new WebDriverException(
String.format("Cannot find file [%s] in directory %s.", filename, downloadsPath));
String.format("Cannot find file [%s] in directory %s.", filename, dir.getAbsolutePath()));
}
if (allFiles.length != 1) {
throw new WebDriverException(
Expand Down Expand Up @@ -590,6 +671,26 @@ public void stop(SessionId id) throws NoSuchSessionException {
}

currentSessions.invalidate(id);
deleteDownloadsFolderForSession(id);
}

private void deleteDownloadsFolderForSession(SessionId id) {
String childDir = downloadFolderToSessionMapping.get(id).toString();
Path toDelete = Paths.get(defaultDownloadsBaseFolder().getAbsolutePath(), childDir);
if (!toDelete.toFile().exists()) {
return;
}
try (Stream<Path> walk = Files.walk(toDelete)) {
walk
.sorted(Comparator.reverseOrder())
.map(Path::toFile)
.peek(each -> LOG.info("Deleting downloads folder for session " + id))
.forEach(File::delete);
downloadFolderToSessionMapping.remove(id);
} catch (IOException e) {
LOG.warning("Could not delete downloads directory for session " + id + ". Root cause: "
+ e.getMessage());
}
}

private void stopAllSessions() {
Expand Down Expand Up @@ -756,6 +857,22 @@ private Map<String, Object> toJson() {
.collect(Collectors.toSet()));
}

private File defaultDownloadsBaseFolder() {
if (defaultDownloadsDir != null) {
return defaultDownloadsDir;
}
String location = baseDir + File.separator + ".cache" +
File.separator + "selenium" + File.separator + "downloads";
defaultDownloadsDir = new File(location);
boolean created = defaultDownloadsDir.mkdirs();
if (created) {
LOG.info("All files downloaded by sessions on this node can be found under [" + location + "].");
} else {
LOG.info(location + " already exists. Using it for managing downloaded files.");
}
return defaultDownloadsDir;
}

public static class Builder {

private final Tracer tracer;
Expand All @@ -772,7 +889,8 @@ public static class Builder {
private Duration sessionTimeout = Duration.ofSeconds(NodeOptions.DEFAULT_SESSION_TIMEOUT);
private HealthCheck healthCheck;
private Duration heartbeatPeriod = Duration.ofSeconds(NodeOptions.DEFAULT_HEARTBEAT_PERIOD);
private String downloadsPath = "";
private boolean enableManageDownloads = false;
private String baseDirectory;

private Builder(
Tracer tracer,
Expand Down Expand Up @@ -827,8 +945,13 @@ public Builder heartbeatPeriod(Duration heartbeatPeriod) {
return this;
}

public Builder downloadsPath(String path) {
this.downloadsPath = path;
public Builder enableManageDownloads(boolean enable) {
this.enableManageDownloads = enable;
return this;
}

public Builder baseDirectory(String directory) {
this.baseDirectory = directory;
return this;
}

Expand All @@ -848,7 +971,8 @@ public LocalNode build() {
heartbeatPeriod,
factories.build(),
registrationSecret,
downloadsPath);
Optional.ofNullable(baseDirectory).orElse(System.getProperty("user.home")),
enableManageDownloads);
}

public Advanced advanced() {
Expand Down
Loading

0 comments on commit b646a06

Please sign in to comment.