code: // Licensed to the Software Freedom Conservancy (SFC) under one
  // or more contributor license agreements.  See the NOTICE file
  // distributed with this work for additional information
  // regarding copyright ownership.  The SFC licenses this file
  // to you under the Apache License, Version 2.0 (the
  // "License"); you may not use this file except in compliance
  // with the License.  You may obtain a copy of the License at
  //
  //   http://www.apache.org/licenses/LICENSE-2.0
  //
  // Unless required by applicable law or agreed to in writing,
  // software distributed under the License is distributed on an
  // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
  // KIND, either express or implied.  See the License for the
  // specific language governing permissions and limitations
  // under the License.
  
  package org.openqa.selenium.grid.node.config;
  
  import static org.openqa.selenium.remote.RemoteTags.CAPABILITIES;
  import static org.openqa.selenium.remote.RemoteTags.CAPABILITIES_EVENT;
  import static org.openqa.selenium.remote.tracing.Tags.EXCEPTION;
  
  import java.net.URI;
  import java.net.URISyntaxException;
  import java.net.URL;
  import java.time.Duration;
  import java.time.Instant;
  import java.util.HashMap;
  import java.util.Map;
  import java.util.Optional;
  import java.util.Set;
  import java.util.function.Function;
  import java.util.function.Predicate;
  import java.util.logging.Logger;
  import java.util.stream.Stream;
  import org.openqa.selenium.Capabilities;
  import org.openqa.selenium.ImmutableCapabilities;
  import org.openqa.selenium.MutableCapabilities;
  import org.openqa.selenium.PersistentCapabilities;
  import org.openqa.selenium.Platform;
  import org.openqa.selenium.SessionNotCreatedException;
  import org.openqa.selenium.WebDriverException;
  import org.openqa.selenium.devtools.CdpEndpointFinder;
  import org.openqa.selenium.grid.data.CreateSessionRequest;
  import org.openqa.selenium.grid.node.ActiveSession;
  import org.openqa.selenium.grid.node.DefaultActiveSession;
  import org.openqa.selenium.grid.node.SessionFactory;
  import org.openqa.selenium.internal.Either;
  import org.openqa.selenium.internal.Require;
  import org.openqa.selenium.net.HostIdentifier;
  import org.openqa.selenium.net.NetworkUtils;
  import org.openqa.selenium.remote.Command;
  import org.openqa.selenium.remote.Dialect;
  import org.openqa.selenium.remote.DriverCommand;
  import org.openqa.selenium.remote.ProtocolHandshake;
  import org.openqa.selenium.remote.Response;
  import org.openqa.selenium.remote.SessionId;
  import org.openqa.selenium.remote.http.ClientConfig;
  import org.openqa.selenium.remote.http.HttpClient;
  import org.openqa.selenium.remote.service.DriverService;
  import org.openqa.selenium.remote.tracing.AttributeKey;
  import org.openqa.selenium.remote.tracing.EventAttribute;
  import org.openqa.selenium.remote.tracing.EventAttributeValue;
  import org.openqa.selenium.remote.tracing.Span;
  import org.openqa.selenium.remote.tracing.Status;
  import org.openqa.selenium.remote.tracing.Tracer;
  
  public class DriverServiceSessionFactory implements SessionFactory {
  
    private static final Logger LOG = Logger.getLogger(DriverServiceSessionFactory.class.getName());
  
    private final Tracer tracer;
    private final HttpClient.Factory clientFactory;
    private final Duration sessionTimeout;
    private final Predicate<Capabilities> predicate;
    private final DriverService.Builder<?, ?> builder;
    private final Capabilities stereotype;
    private final SessionCapabilitiesMutator sessionCapabilitiesMutator;
  
    public DriverServiceSessionFactory(
        Tracer tracer,
        HttpClient.Factory clientFactory,
        Duration sessionTimeout,
        Capabilities stereotype,
        Predicate<Capabilities> predicate,
        DriverService.Builder<?, ?> builder) {
      this.tracer = Require.nonNull("Tracer", tracer);
      this.clientFactory = Require.nonNull("HTTP client factory", clientFactory);
      this.sessionTimeout = Require.nonNull("Session timeout", sessionTimeout);
      this.stereotype = ImmutableCapabilities.copyOf(Require.nonNull("Stereotype", stereotype));
      this.predicate = Require.nonNull("Accepted capabilities predicate", predicate);
      this.builder = Require.nonNull("Driver service builder", builder);
      this.sessionCapabilitiesMutator = new SessionCapabilitiesMutator(this.stereotype);
    }
  
    @Override
    public Capabilities getStereotype() {
      return stereotype;
    }
  
    @Override
    public boolean test(Capabilities capabilities) {
      return predicate.test(capabilities);
    }
  
    @Override
    public Either<WebDriverException, ActiveSession> apply(CreateSessionRequest sessionRequest) {
      if (sessionRequest.getDownstreamDialects().isEmpty()) {
        return Either.left(new SessionNotCreatedException("No downstream dialects were found."));
      }
  
      if (!test(sessionRequest.getDesiredCapabilities())) {
        return Either.left(
            new SessionNotCreatedException(
                "New session request capabilities do not " + "match the stereotype."));
      }
  
      Span span = tracer.getCurrentContext().createSpan("driver_service_factory.apply");
      Map<String, EventAttributeValue> attributeMap = new HashMap<>();
      try {
  
        Capabilities capabilities =
            sessionCapabilitiesMutator.apply(sessionRequest.getDesiredCapabilities());
  
        Optional<Platform> platformName = Optional.ofNullable(capabilities.getPlatformName());
        if (platformName.isPresent()) {
          capabilities = removePlatform(capabilities);
        }
  
        CAPABILITIES.accept(span, capabilities);
        CAPABILITIES_EVENT.accept(attributeMap, capabilities);
        attributeMap.put(
            AttributeKey.LOGGER_CLASS.getKey(), EventAttribute.setValue(this.getClass().getName()));
  
        DriverService service = builder.build();
        try {
          service.start();
  
          URL serviceURL = service.getUrl();
          attributeMap.put(
              AttributeKey.DRIVER_URL.getKey(), EventAttribute.setValue(serviceURL.toString()));
  
          ClientConfig clientConfig =
              ClientConfig.defaultConfig().readTimeout(sessionTimeout).baseUrl(serviceURL);
          HttpClient client = clientFactory.createClient(clientConfig);
  
          Command command = new Command(null, DriverCommand.NEW_SESSION(capabilities));
  
          ProtocolHandshake.Result result = new ProtocolHandshake().createSession(client, command);
  
          Set<Dialect> downstreamDialects = sessionRequest.getDownstreamDialects();
          Dialect upstream = result.getDialect();
          Dialect downstream =
              downstreamDialects.contains(result.getDialect())
                  ? result.getDialect()
                  : downstreamDialects.iterator().next();
  
          Response response = result.createResponse();
  
          attributeMap.put(
              AttributeKey.UPSTREAM_DIALECT.getKey(), EventAttribute.setValue(upstream.toString()));
          attributeMap.put(
              AttributeKey.DOWNSTREAM_DIALECT.getKey(),
              EventAttribute.setValue(downstream.toString()));
          attributeMap.put(
              AttributeKey.DRIVER_RESPONSE.getKey(), EventAttribute.setValue(response.toString()));
  
          Capabilities caps = new ImmutableCapabilities((Map<?, ?>) response.getValue());
          if (platformName.isPresent()) {
            caps = setInitialPlatform(caps, platformName.get());
          }
  
          caps = readDevToolsEndpointAndVersion(caps);
          caps = readBiDiEndpoint(caps);
          caps = readVncEndpoint(capabilities, caps);
  
          span.addEvent("Driver service created session", attributeMap);
          return Either.right(
              new DefaultActiveSession(
                  tracer,
                  client,
                  new SessionId(response.getSessionId()),
                  service.getUrl(),
                  downstream,
                  upstream,
                  stereotype,
                  caps,
                  Instant.now()) {
                @Override
                public void stop() {
                  service.stop();
                  client.close();
                }
              });
        } catch (Exception e) {
          span.setAttribute(AttributeKey.ERROR.getKey(), true);
          span.setStatus(Status.CANCELLED);
          EXCEPTION.accept(attributeMap, e);
          String errorMessage =
              "Error while creating session with the driver service. "
                  + "Stopping driver service: "
                  + e.getMessage();
          LOG.warning(errorMessage);
  
          attributeMap.put(
              AttributeKey.EXCEPTION_MESSAGE.getKey(), EventAttribute.setValue(errorMessage));
          span.addEvent(AttributeKey.EXCEPTION_EVENT.getKey(), attributeMap);
          service.stop();
          return Either.left(new SessionNotCreatedException(errorMessage));
        }
      } catch (Exception e) {
        span.setAttribute(AttributeKey.ERROR.getKey(), true);
        span.setStatus(Status.CANCELLED);
        EXCEPTION.accept(attributeMap, e);
        String errorMessage =
            "Error while creating session with the driver service. " + e.getMessage();
        LOG.warning(errorMessage);
  
        attributeMap.put(
            AttributeKey.EXCEPTION_MESSAGE.getKey(), EventAttribute.setValue(errorMessage));
        span.addEvent(AttributeKey.EXCEPTION_EVENT.getKey(), attributeMap);
  
        return Either.left(new SessionNotCreatedException(e.getMessage()));
      } finally {
        span.close();
      }
    }
  
    private Capabilities readDevToolsEndpointAndVersion(Capabilities caps) {
      class DevToolsInfo {
        public final URI cdpEndpoint;
        public final String version;
  
        public DevToolsInfo(URI cdpEndpoint, String version) {
          this.cdpEndpoint = cdpEndpoint;
          this.version = version;
        }
      }
  
      Function<Capabilities, Optional<DevToolsInfo>> chrome =
          c ->
              CdpEndpointFinder.getReportedUri("goog:chromeOptions", c)
                  .map(uri -> new DevToolsInfo(uri, c.getBrowserVersion()));
  
      Function<Capabilities, Optional<DevToolsInfo>> edge =
          c ->
              CdpEndpointFinder.getReportedUri("ms:edgeOptions", c)
                  .map(uri -> new DevToolsInfo(uri, c.getBrowserVersion()));
  
      Function<Capabilities, Optional<DevToolsInfo>> firefox =
          c ->
              CdpEndpointFinder.getReportedUri("moz:debuggerAddress", c)
                  .map(uri -> new DevToolsInfo(uri, "85.0"));
  
      Optional<DevToolsInfo> maybeInfo =
          Stream.of(chrome, edge, firefox)
              .map(finder -> finder.apply(caps))
              .filter(Optional::isPresent)
              .map(Optional::get)
              .findFirst();
  
      if (maybeInfo.isPresent()) {
        DevToolsInfo info = maybeInfo.get();
        return new PersistentCapabilities(caps)
            .setCapability("se:cdp", info.cdpEndpoint)
            .setCapability("se:cdpVersion", info.version);
      }
      return caps;
    }
  
    private Capabilities readBiDiEndpoint(Capabilities caps) {
  
      Optional<String> webSocketUrl =
          Optional.ofNullable((String) caps.getCapability("webSocketUrl"));
  
      Optional<URI> websocketUri =
          webSocketUrl.map(
              uri -> {
                try {
                  return new URI(uri);
                } catch (URISyntaxException e) {
                  LOG.warning(e.getMessage());
                }
                return null;
              });
  
      if (websocketUri.isPresent()) {
        return new PersistentCapabilities(caps).setCapability("se:bidi", websocketUri.get());
      }
  
      return caps;
    }
  
    private Capabilities readVncEndpoint(Capabilities requestedCaps, Capabilities returnedCaps) {
      String seVncEnabledCap = "se:vncEnabled";
      String seNoVncPortCap = "se:noVncPort";
      String seVncEnabled = String.valueOf(requestedCaps.getCapability(seVncEnabledCap));
      boolean vncLocalAddressSet = requestedCaps.getCapabilityNames().contains("se:vncLocalAddress");
      if (Boolean.parseBoolean(seVncEnabled) && !vncLocalAddressSet) {
        String seNoVncPort = String.valueOf(requestedCaps.getCapability(seNoVncPortCap));
        String vncLocalAddress = String.format("ws://%s:%s", getHost(), seNoVncPort);
        returnedCaps =
            new PersistentCapabilities(returnedCaps)
                .setCapability("se:vncLocalAddress", vncLocalAddress)
                .setCapability(seVncEnabledCap, true);
      }
      return returnedCaps;
    }
  
    // We remove the platform before sending the caps to the driver because some drivers will
    // reject session requests when they cannot parse the platform.
    private Capabilities removePlatform(Capabilities caps) {
      MutableCapabilities noPlatformName = new MutableCapabilities(new HashMap<>(caps.asMap()));
      noPlatformName.setCapability("platformName", (String) null);
      return new PersistentCapabilities(noPlatformName);
    }
  
    private Capabilities setInitialPlatform(Capabilities caps, Platform platform) {
      return new PersistentCapabilities(caps).setCapability("platformName", platform);
    }
  
    private String getHost() {
      try {
        return new NetworkUtils().getNonLoopbackAddressOfThisMachine();
      } catch (WebDriverException e) {
        return HostIdentifier.getHostName();
      }
    }
  }
  