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.relay;
  
  import static org.openqa.selenium.remote.RemoteTags.CAPABILITIES;
  import static org.openqa.selenium.remote.RemoteTags.CAPABILITIES_EVENT;
  import static org.openqa.selenium.remote.tracing.AttributeKey.DOWNSTREAM_DIALECT;
  import static org.openqa.selenium.remote.tracing.AttributeKey.DRIVER_RESPONSE;
  import static org.openqa.selenium.remote.tracing.AttributeKey.DRIVER_URL;
  import static org.openqa.selenium.remote.tracing.AttributeKey.EXCEPTION_EVENT;
  import static org.openqa.selenium.remote.tracing.AttributeKey.EXCEPTION_MESSAGE;
  import static org.openqa.selenium.remote.tracing.AttributeKey.LOGGER_CLASS;
  import static org.openqa.selenium.remote.tracing.AttributeKey.UPSTREAM_DIALECT;
  import static org.openqa.selenium.remote.tracing.EventAttribute.setValue;
  import static org.openqa.selenium.remote.tracing.Tags.EXCEPTION;
  
  import java.net.MalformedURLException;
  import java.net.URI;
  import java.net.URL;
  import java.time.Duration;
  import java.time.Instant;
  import java.util.HashMap;
  import java.util.Map;
  import java.util.Objects;
  import java.util.Set;
  import java.util.logging.Level;
  import java.util.logging.Logger;
  import org.openqa.selenium.Capabilities;
  import org.openqa.selenium.ImmutableCapabilities;
  import org.openqa.selenium.SessionNotCreatedException;
  import org.openqa.selenium.WebDriverException;
  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.Debug;
  import org.openqa.selenium.internal.Either;
  import org.openqa.selenium.internal.Require;
  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.Contents;
  import org.openqa.selenium.remote.http.HttpClient;
  import org.openqa.selenium.remote.http.HttpMethod;
  import org.openqa.selenium.remote.http.HttpRequest;
  import org.openqa.selenium.remote.http.HttpResponse;
  import org.openqa.selenium.remote.tracing.AttributeKey;
  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 RelaySessionFactory implements SessionFactory {
  
    private static final Logger LOG = Logger.getLogger(RelaySessionFactory.class.getName());
  
    private final Tracer tracer;
    private final HttpClient.Factory clientFactory;
    private final Duration sessionTimeout;
    private final URL serviceUrl;
    private final URL serviceStatusUrl;
    private final Capabilities stereotype;
  
    public RelaySessionFactory(
        Tracer tracer,
        HttpClient.Factory clientFactory,
        Duration sessionTimeout,
        URI serviceUri,
        URI serviceStatusUri,
        Capabilities stereotype) {
      this.tracer = Require.nonNull("Tracer", tracer);
      this.clientFactory = Require.nonNull("HTTP client", clientFactory);
      this.sessionTimeout = Require.nonNull("Session timeout", sessionTimeout);
      this.serviceUrl = createUrlFromUri(Require.nonNull("Service URL", serviceUri));
      this.serviceStatusUrl = createUrlFromUri(serviceStatusUri);
      this.stereotype = ImmutableCapabilities.copyOf(Require.nonNull("Stereotype", stereotype));
    }
  
    @Override
    public Capabilities getStereotype() {
      return stereotype;
    }
  
    @Override
    public boolean test(Capabilities capabilities) {
      // If a request reaches this point is because the basic match of W3C caps has already been done.
  
      // Custom matching in case a platformVersion is requested
      boolean platformVersionMatch =
          capabilities.getCapabilityNames().stream()
              .filter(name -> name.contains("platformVersion"))
              .map(
                  platformVersionCapName ->
                      Objects.equals(
                          stereotype.getCapability(platformVersionCapName),
                          capabilities.getCapability(platformVersionCapName)))
              .reduce(Boolean::logicalAnd)
              .orElse(true);
  
      return platformVersionMatch
          && stereotype.getCapabilityNames().stream()
              .filter(name -> capabilities.asMap().containsKey(name))
              .map(
                  name -> {
                    if (capabilities.getCapability(name) instanceof String) {
                      return stereotype
                          .getCapability(name)
                          .toString()
                          .equalsIgnoreCase(capabilities.getCapability(name).toString());
                    } else {
                      return capabilities.getCapability(name) == null
                          || Objects.equals(
                              stereotype.getCapability(name), capabilities.getCapability(name));
                    }
                  })
              .reduce(Boolean::logicalAnd)
              .orElse(false);
    }
  
    @Override
    public Either<WebDriverException, ActiveSession> apply(CreateSessionRequest sessionRequest) {
      Capabilities capabilities = sessionRequest.getDesiredCapabilities();
      if (!test(capabilities)) {
        return Either.left(
            new SessionNotCreatedException(
                "New session request capabilities do not " + "match the stereotype."));
      }
  
      LOG.info("Starting session for " + capabilities);
  
      try (Span span = tracer.getCurrentContext().createSpan("relay_session_factory.apply")) {
  
        Map<String, EventAttributeValue> attributeMap = new HashMap<>();
        CAPABILITIES.accept(span, capabilities);
        CAPABILITIES_EVENT.accept(attributeMap, capabilities);
        attributeMap.put(LOGGER_CLASS.getKey(), setValue(this.getClass().getName()));
        attributeMap.put(DRIVER_URL.getKey(), 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));
        try {
          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(UPSTREAM_DIALECT.getKey(), setValue(upstream.toString()));
          attributeMap.put(DOWNSTREAM_DIALECT.getKey(), setValue(downstream.toString()));
          attributeMap.put(DRIVER_RESPONSE.getKey(), setValue(response.toString()));
  
          Capabilities responseCaps = new ImmutableCapabilities((Map<?, ?>) response.getValue());
          Capabilities mergedCapabilities = capabilities.merge(responseCaps);
  
          span.addEvent("Relay service created session", attributeMap);
          LOG.fine(String.format("Created session: %s - %s", response.getSessionId(), capabilities));
          return Either.right(
              new DefaultActiveSession(
                  tracer,
                  client,
                  new SessionId(response.getSessionId()),
                  serviceUrl,
                  downstream,
                  upstream,
                  stereotype,
                  mergedCapabilities,
                  Instant.now()) {
                @Override
                public void stop() {
                  // no-op
                }
              });
        } catch (Exception e) {
          span.setAttribute(AttributeKey.ERROR.getKey(), true);
          span.setStatus(Status.CANCELLED);
          EXCEPTION.accept(attributeMap, e);
          String errorMessage =
              String.format(
                  "Error while creating session with the service %s. %s", serviceUrl, e.getMessage());
          attributeMap.put(EXCEPTION_MESSAGE.getKey(), setValue(errorMessage));
          span.addEvent(EXCEPTION_EVENT.getKey(), attributeMap);
          return Either.left(new SessionNotCreatedException(errorMessage));
        }
      } catch (Exception e) {
        return Either.left(new SessionNotCreatedException(e.getMessage()));
      }
    }
  
    public boolean isServiceUp() {
      if (serviceStatusUrl == null) {
        // If no status endpoint was configured, we assume the server is up.
        return true;
      }
      try {
        HttpClient client = clientFactory.createClient(serviceStatusUrl);
        HttpResponse response =
            client.execute(new HttpRequest(HttpMethod.GET, serviceStatusUrl.toString()));
        LOG.log(Debug.getDebugLogLevel(), Contents.string(response));
        return response.getStatus() == 200;
      } catch (Exception e) {
        LOG.log(
            Level.WARNING,
            () ->
                String.format(
                    "Error checking service status %s. %s", serviceStatusUrl, e.getMessage()));
        LOG.log(Debug.getDebugLogLevel(), "Error checking service status " + serviceStatusUrl, e);
      }
      return false;
    }
  
    private URL createUrlFromUri(URI uri) {
      if (uri == null) {
        return null;
      }
      try {
        return uri.toURL();
      } catch (MalformedURLException e) {
        throw new RuntimeException(e.getMessage(), e);
      }
    }
  }
  