[java] Tapping the Node session when there is WebSocket activity

Fixes SeleniumHQ#12223

java/src/org/openqa/selenium/grid/node/ProxyNodeWebsockets.java

@@ -33,6 +33,7 @@
 import org.openqa.selenium.Capabilities;
 import org.openqa.selenium.devtools.CdpEndpointFinder;
 import org.openqa.selenium.grid.data.Session;
+import org.openqa.selenium.internal.Require;
 import org.openqa.selenium.remote.SessionId;
 import org.openqa.selenium.remote.http.BinaryMessage;
 import org.openqa.selenium.remote.http.ClientConfig;
@@ -75,57 +76,68 @@ public Optional<Consumer<Message>> apply(String uri, Consumer<Message> downstrea
  return Optional.empty();
  }
 
- String sessionId =
- Stream.of(fwdMatch, cdpMatch, bidiMatch, vncMatch)
- .filter(Objects::nonNull)
- .findFirst()
- .get()
- .getParameters()
- .get("sessionId");
+ Optional<UrlTemplate.Match> firstMatch =
+ Stream.of(fwdMatch, cdpMatch, bidiMatch, vncMatch).filter(Objects::nonNull).findFirst();
+
+ if (firstMatch.isEmpty()) {
+ LOG.warning("No session id found in uri " + uri);
+ return Optional.empty();
+ }
+
+ String sessionId = firstMatch.get().getParameters().get("sessionId");
 
  LOG.fine("Matching websockets for session id: " + sessionId);
  SessionId id = new SessionId(sessionId);
 
  if (!node.isSessionOwner(id)) {
- LOG.info("Not owner of " + id);
+ LOG.warning("Not owner of " + id);
  return Optional.empty();
  }
 
  Session session = node.getSession(id);
  Capabilities caps = session.getCapabilities();
  LOG.fine("Scanning for endpoint: " + caps);
 
+ // Used by the ForwardingListener to notify the node that the session is still active
+ Consumer<SessionId> sessionConsumer = node::isSessionOwner;
+
  if (bidiMatch != null) {
- return findBiDiEndpoint(downstream, caps);
+ return findBiDiEndpoint(downstream, caps, sessionConsumer, id);
  }
 
  if (vncMatch != null) {
- return findVncEndpoint(downstream, caps);
+ // Passing a fake consumer to the ForwardingListener to avoid sending a session notification
+ // when VNC is used.
+ sessionConsumer = fakeConsumer -> {};
+ return findVncEndpoint(downstream, caps, sessionConsumer, id);
  }
 
  // This match happens when a user wants to do CDP over Dynamic Grid
  if (fwdMatch != null) {
  LOG.info("Matched endpoint where CDP connection is being forwarded");
- return findCdpEndpoint(downstream, caps);
+ return findCdpEndpoint(downstream, caps, sessionConsumer, id);
  }
  if (caps.getCapabilityNames().contains("se:forwardCdp")) {
  LOG.info("Found endpoint where CDP connection needs to be forwarded");
- return findForwardCdpEndpoint(downstream, caps);
+ return findForwardCdpEndpoint(downstream, caps, sessionConsumer, id);
  }
- return findCdpEndpoint(downstream, caps);
+ return findCdpEndpoint(downstream, caps, sessionConsumer, id);
  }
 
  private Optional<Consumer<Message>> findCdpEndpoint(
- Consumer<Message> downstream, Capabilities caps) {
- // Using strings here to avoid Node depending upon specific drivers.
+ Consumer<Message> downstream,
+ Capabilities caps,
+ Consumer<SessionId> sessionConsumer,
+ SessionId sessionId) {
+
  for (String cdpEndpointCap : CDP_ENDPOINT_CAPS) {
  Optional<URI> reportedUri = CdpEndpointFinder.getReportedUri(cdpEndpointCap, caps);
  Optional<HttpClient> client =
  reportedUri.map(uri -> CdpEndpointFinder.getHttpClient(clientFactory, uri));
  Optional<URI> cdpUri;
 
  try {
- cdpUri = client.flatMap(httpClient -> CdpEndpointFinder.getCdpEndPoint(httpClient));
+ cdpUri = client.flatMap(CdpEndpointFinder::getCdpEndPoint);
  } catch (Exception e) {
  try {
  client.ifPresent(HttpClient::close);
@@ -137,7 +149,7 @@ private Optional<Consumer<Message>> findCdpEndpoint(
 
  if (cdpUri.isPresent()) {
  LOG.log(getDebugLogLevel(), String.format("Endpoint found in %s", cdpEndpointCap));
- return cdpUri.map(cdp -> createWsEndPoint(cdp, downstream));
+ return cdpUri.map(cdp -> createWsEndPoint(cdp, downstream, sessionConsumer, sessionId));
  } else {
  try {
  client.ifPresent(HttpClient::close);
@@ -154,30 +166,41 @@ private Optional<Consumer<Message>> findCdpEndpoint(
  }
 
  private Optional<Consumer<Message>> findBiDiEndpoint(
- Consumer<Message> downstream, Capabilities caps) {
+ Consumer<Message> downstream,
+ Capabilities caps,
+ Consumer<SessionId> sessionConsumer,
+ SessionId sessionId) {
  try {
  URI uri = new URI(String.valueOf(caps.getCapability("webSocketUrl")));
- return Optional.of(uri).map(bidi -> createWsEndPoint(bidi, downstream));
+ return Optional.of(uri)
+ .map(bidi -> createWsEndPoint(bidi, downstream, sessionConsumer, sessionId));
  } catch (URISyntaxException e) {
  LOG.warning("Unable to create URI from: " + caps.getCapability("webSocketUrl"));
  return Optional.empty();
  }
  }
 
  private Optional<Consumer<Message>> findForwardCdpEndpoint(
- Consumer<Message> downstream, Capabilities caps) {
+ Consumer<Message> downstream,
+ Capabilities caps,
+ Consumer<SessionId> sessionConsumer,
+ SessionId sessionId) {
  // When using Dynamic Grid, we need to connect to a container before using the debuggerAddress
  try {
  URI uri = new URI(String.valueOf(caps.getCapability("se:forwardCdp")));
- return Optional.of(uri).map(cdp -> createWsEndPoint(cdp, downstream));
+ return Optional.of(uri)
+ .map(cdp -> createWsEndPoint(cdp, downstream, sessionConsumer, sessionId));
  } catch (URISyntaxException e) {
  LOG.warning("Unable to create URI from: " + caps.getCapability("se:forwardCdp"));
  return Optional.empty();
  }
  }
 
  private Optional<Consumer<Message>> findVncEndpoint(
- Consumer<Message> downstream, Capabilities caps) {
+ Consumer<Message> downstream,
+ Capabilities caps,
+ Consumer<SessionId> sessionConsumer,
+ SessionId sessionId) {
  String vncLocalAddress = (String) caps.getCapability("se:vncLocalAddress");
  Optional<URI> vncUri;
  try {
@@ -187,40 +210,57 @@ private Optional<Consumer<Message>> findVncEndpoint(
  return Optional.empty();
  }
  LOG.log(getDebugLogLevel(), String.format("Endpoint found in %s", "se:vncLocalAddress"));
- return vncUri.map(vnc -> createWsEndPoint(vnc, downstream));
+ return vncUri.map(vnc -> createWsEndPoint(vnc, downstream, sessionConsumer, sessionId));
  }
 
- private Consumer<Message> createWsEndPoint(URI uri, Consumer<Message> downstream) {
- Objects.requireNonNull(uri);
+ private Consumer<Message> createWsEndPoint(
+ URI uri,
+ Consumer<Message> downstream,
+ Consumer<SessionId> sessionConsumer,
+ SessionId sessionId) {
+ Require.nonNull("downstream", downstream);
+ Require.nonNull("uri", uri);
+ Require.nonNull("sessionConsumer", sessionConsumer);
+ Require.nonNull("sessionId", sessionId);
 
  LOG.info("Establishing connection to " + uri);
 
  HttpClient client = clientFactory.createClient(ClientConfig.defaultConfig().baseUri(uri));
  WebSocket upstream =
- client.openSocket(new HttpRequest(GET, uri.toString()), new ForwardingListener(downstream));
+ client.openSocket(
+ new HttpRequest(GET, uri.toString()),
+ new ForwardingListener(downstream, sessionConsumer, sessionId));
  return upstream::send;
  }
 
  private static class ForwardingListener implements WebSocket.Listener {
  private final Consumer<Message> downstream;
+ private final Consumer<SessionId> sessionConsumer;
+ private final SessionId sessionId;
 
- public ForwardingListener(Consumer<Message> downstream) {
+ public ForwardingListener(
+ Consumer<Message> downstream, Consumer<SessionId> sessionConsumer, SessionId sessionId) {
  this.downstream = Objects.requireNonNull(downstream);
+ this.sessionConsumer = Objects.requireNonNull(sessionConsumer);
+ this.sessionId = Objects.requireNonNull(sessionId);
  }
 
  @Override
  public void onBinary(byte[] data) {
  downstream.accept(new BinaryMessage(data));
+ sessionConsumer.accept(sessionId);
  }
 
  @Override
  public void onClose(int code, String reason) {
  downstream.accept(new CloseMessage(code, reason));
+ sessionConsumer.accept(sessionId);
  }
 
  @Override
  public void onText(CharSequence data) {
  downstream.accept(new TextMessage(data));
+ sessionConsumer.accept(sessionId);
  }
 
  @Override