public ProtocolConverter(
  Tracer tracer,
  HttpClient client,
  Dialect downstream,
  Dialect upstream) {
  this.tracer = Require.nonNull("Tracer", tracer);
  this.client = Require.nonNull("HTTP client", client);

  this.downstream = getCommandCodec(Require.nonNull("Downstream dialect", downstream));
  this.downstreamResponse = getResponseCodec(downstream);

  this.upstream = getCommandCodec(Require.nonNull("Upstream dialect", upstream));
  this.upstreamResponse = getResponseCodec(upstream);

  converter = new JsonToWebElementConverter(null);

  newSessionConverter = downstream == W3C ? this::createW3CNewSessionResponse : this::createJwpNewSessionResponse;
}
 

@Override
protected ServicedSession newActiveSession(
    DriverService service,
    Dialect downstream,
    Dialect upstream,
    HttpHandler codec,
    SessionId id,
    Map<String, Object> capabilities) {
  return new ServicedSession(
      service,
      downstream,
      upstream,
      codec,
      id,
      capabilities);
}
 

protected RemoteSession(
    Dialect downstream,
    Dialect upstream,
    HttpHandler codec,
    SessionId id,
    Map<String, Object> capabilities) {
  this.downstream = Require.nonNull("Downstream dialect", downstream);
  this.upstream = Require.nonNull("Upstream dialect", upstream);
  this.codec = Require.nonNull("Codec", codec);
  this.id = Require.nonNull("Session id", id);
  this.capabilities = Require.nonNull("Capabilities", capabilities);

  File tempRoot = new File(StandardSystemProperty.JAVA_IO_TMPDIR.value(), id.toString());
  Require.stateCondition(tempRoot.mkdirs(), "Could not create directory %s", tempRoot);
  this.filesystem = TemporaryFilesystem.getTmpFsBasedOn(tempRoot);

  CommandExecutor executor = new ActiveSessionCommandExecutor(this);
  this.driver = new Augmenter().augment(new RemoteWebDriver(
      executor,
      new ImmutableCapabilities(getCapabilities())));
}
 

protected ProtocolConvertingSession(
    Tracer tracer,
    HttpClient client,
    SessionId id,
    URL url,
    Dialect downstream,
    Dialect upstream,
    Capabilities capabilities) {
  super(id, url, downstream, upstream, capabilities);

  Require.nonNull("HTTP client", client);

  if (downstream.equals(upstream)) {
    this.handler = new ReverseProxyHandler(tracer, client);
  } else {
    this.handler = new ProtocolConverter(tracer, client, downstream, upstream);
  }

  this.killUrl = "/session/" + id;
}
 

/**
 * Create a UTF-8 encoded response for a given dialect for use with the New Session command.
 */
private static byte[] encodeAsResponse(
    Dialect dialect,
    SessionId id,
    Capabilities capabilities,
    Map<String, Object> metadata) {

  Map<String, Object> toEncode;

  switch (dialect) {
    case OSS:
      toEncode = encodeJsonWireProtocol(id, capabilities, metadata);
      break;

    case W3C:
      toEncode = encodeW3C(id, capabilities, metadata);
      break;

    default:
      throw new IllegalArgumentException("Unknown dialect: " + dialect);
  }

  return JSON.toJson(toEncode).getBytes(UTF_8);
}
 

protected BaseActiveSession(
    SessionId id,
    URL url,
    Dialect downstream,
    Dialect upstream,
    Capabilities capabilities) {
  URI uri = null;
  try {
    uri = Require.nonNull("URL", url).toURI();
  } catch (URISyntaxException e) {
    throw new IllegalArgumentException(e);
  }

  this.session = new Session(
      Require.nonNull("Session id", id),
      uri,
      ImmutableCapabilities.copyOf(Require.nonNull("Capabilities", capabilities)));

  this.downstream = Require.nonNull("Downstream dialect", downstream);
  this.upstream = Require.nonNull("Upstream dialect", upstream);
}
 

private InMemorySession(WebDriver driver, Capabilities capabilities, Dialect downstream) {
  this.driver = Require.nonNull("Driver", driver);

  Capabilities caps;
  if (driver instanceof HasCapabilities) {
    caps = ((HasCapabilities) driver).getCapabilities();
  } else {
    caps = capabilities;
  }

  this.capabilities = caps.asMap().entrySet().stream()
      .filter(e -> e.getValue() != null)
      .collect(ImmutableMap.toImmutableMap(Map.Entry::getKey, Map.Entry::getValue));

  this.id = new SessionId(UUID.randomUUID().toString());
  this.downstream = Require.nonNull("Downstream dialect", downstream);

  File tempRoot = new File(StandardSystemProperty.JAVA_IO_TMPDIR.value(), id.toString());
  Require.stateCondition(tempRoot.mkdirs(), "Could not create directory %s", tempRoot);
  this.filesystem = TemporaryFilesystem.getTmpFsBasedOn(tempRoot);

  this.handler = new JsonHttpCommandHandler(
      new PretendDriverSessions(),
      LOG);
}
 

@Override
public Optional<ActiveSession> apply(CreateSessionRequest sessionRequest) {
  Require.nonNull("Session creation request", sessionRequest);

  // Assume the blob fits in the available memory.
  try {
    if (!provider.canCreateDriverInstanceFor(sessionRequest.getCapabilities())) {
      return Optional.empty();
    }

    WebDriver driver = provider.newInstance(sessionRequest.getCapabilities());

    // Prefer the OSS dialect.
    Set<Dialect> downstreamDialects = sessionRequest.getDownstreamDialects();
    Dialect downstream = downstreamDialects.contains(Dialect.OSS) || downstreamDialects.isEmpty() ?
                         Dialect.OSS :
                         downstreamDialects.iterator().next();
    return Optional.of(
        new InMemorySession(driver, sessionRequest.getCapabilities(), downstream));
  } catch (IllegalStateException e) {
    return Optional.empty();
  }
}
 

private Node createNode(Capabilities stereotype, int count, int currentLoad) {
  URI uri = createUri();
  LocalNode.Builder builder = LocalNode.builder(tracer, bus, uri, uri, null);
  for (int i = 0; i < count; i++) {
    builder.add(stereotype, new TestSessionFactory((id, caps) -> new HandledSession(uri, caps)));
  }

  LocalNode node = builder.build();
  for (int i = 0; i < currentLoad; i++) {
    // Ignore the session. We're just creating load.
    node.newSession(new CreateSessionRequest(
        ImmutableSet.copyOf(Dialect.values()),
        stereotype,
        ImmutableMap.of()));
  }

  return node;
}
 

@Test
public void shouldNotInstantiateSessionIfRemoteEndReturnsInvalidResponse() throws IOException {
  HttpClient httpClient = mock(HttpClient.class);
  when(httpClient.execute(any(HttpRequest.class))).thenReturn(
      new HttpResponse().setStatus(200).setContent(() -> new ByteArrayInputStream(
          "Hello, world!".getBytes())));
  when(clientFactory.createClient(any(URL.class))).thenReturn(httpClient);

  DriverServiceSessionFactory factory = factoryFor("chrome", builder);

  Optional<ActiveSession> session = factory.apply(new CreateSessionRequest(
      ImmutableSet.of(Dialect.W3C), toPayload("chrome"), ImmutableMap.of()));

  assertThat(session).isEmpty();

  verify(builder, times(1)).build();
  verifyNoMoreInteractions(builder);
  verify(driverService, times(1)).start();
  verify(driverService, atLeastOnce()).getUrl();
  verify(driverService, times(1)).stop();
  verifyNoMoreInteractions(driverService);
}
 

@Test
public void shouldInstantiateSessionIfEverythingIsOK() throws IOException {
  HttpClient httpClient = mock(HttpClient.class);
  when(httpClient.execute(any(HttpRequest.class))).thenReturn(
      new HttpResponse().setStatus(200).setContent(() -> new ByteArrayInputStream(
          "{ \"value\": { \"sessionId\": \"1\", \"capabilities\": {} } }".getBytes())));
  when(clientFactory.createClient(any(URL.class))).thenReturn(httpClient);

  DriverServiceSessionFactory factory = factoryFor("chrome", builder);

  Optional<ActiveSession> session = factory.apply(new CreateSessionRequest(
      ImmutableSet.of(Dialect.W3C), toPayload("chrome"), ImmutableMap.of()));

  assertThat(session).isNotEmpty();

  verify(builder, times(1)).build();
  verifyNoMoreInteractions(builder);
  verify(driverService, times(1)).start();
  verify(driverService, atLeastOnce()).getUrl();
  verifyNoMoreInteractions(driverService);
}
 

@Test
public void factoriesFoundViaServiceLoadersAreUsedFirst() {
  WebDriver driver = Mockito.mock(WebDriver.class);
  Capabilities caps = new ImmutableCapabilities("browserName", "chrome");
  DriverProvider provider = new StubbedProvider(caps, driver);

  ActiveSessionFactory sessionFactory = new ActiveSessionFactory(new NullTracer()) {
    @Override
    protected Iterable<DriverProvider> loadDriverProviders() {
      return ImmutableSet.of(provider);
    }
  };

  ActiveSession session = sessionFactory.apply(
      new CreateSessionRequest(ImmutableSet.of(Dialect.W3C), caps, ImmutableMap.of()))
      .get();
  assertEquals(driver, session.getWrappedDriver());
}
 

@Test
public void shouldWriteABase64EncodedZippedFileToDiskAndKeepName() throws Exception {
  ActiveSession session = mock(ActiveSession.class);
  when(session.getId()).thenReturn(new SessionId("1234567"));
  when(session.getFileSystem()).thenReturn(tempFs);
  when(session.getDownstreamDialect()).thenReturn(Dialect.OSS);

  File tempFile = touch(null, "foo");
  String encoded = Zip.zip(tempFile);

  UploadFile uploadFile = new UploadFile(new Json(), session);
  Map<String, Object> args = ImmutableMap.of("file", encoded);
  HttpRequest request = new HttpRequest(HttpMethod.POST, "/session/%d/se/file");
  request.setContent(asJson(args));
  HttpResponse response = uploadFile.execute(request);

  Response res = new Json().toType(string(response), Response.class);
  String path = (String) res.getValue();
  assertTrue(new File(path).exists());
  assertTrue(path.endsWith(tempFile.getName()));
}
 

private void setCodec() {
	try {
		CommandExecutor executor = getCommandExecutor();
		setField("commandCodec", executor, Dialect.W3C.getCommandCodec());
		setField("responseCodec", executor, Dialect.W3C.getResponseCodec());
	} catch (Throwable e) {
		logger.error("Unable to set W3C codec", e);
	}
}
 

@Test
public void convertsAnArrayWithAWebElement() {
  RemoteWebElement element = new RemoteWebElement();
  element.setId("abc123");

  Object value = CONVERTER.apply(new Object[] { element });
  assertContentsInOrder(new ArrayList<>((Collection<?>) value),
      ImmutableMap.of(
        Dialect.OSS.getEncodedElementKey(), "abc123",
        Dialect.W3C.getEncodedElementKey(), "abc123"));
}
 

private static void assertIsWebElementObject(Object value, String expectedKey) {
  assertThat(value).isInstanceOf(Map.class);

  Map<?, ?>  map = (Map<?, ?>) value;
  assertThat(map).hasSize(2);
  assertThat(map.containsKey(Dialect.OSS.getEncodedElementKey())).isTrue();
  assertThat(map.get(Dialect.OSS.getEncodedElementKey())).isEqualTo(expectedKey);
  assertThat(map.containsKey(Dialect.W3C.getEncodedElementKey())).isTrue();
  assertThat(map.get(Dialect.W3C.getEncodedElementKey())).isEqualTo(expectedKey);
}
 

@Test
public void shouldConvertElementReferenceToRemoteWebElement() {
  HttpResponse response = new HttpResponse();
  response.setStatus(HTTP_OK);
  response.setContent(asJson(ImmutableMap.of(
      "status", 0,
      "value", ImmutableMap.of(Dialect.OSS.getEncodedElementKey(), "345678"))));

  Response decoded = codec.decode(response);
  assertThat(((RemoteWebElement) decoded.getValue()).getId()).isEqualTo("345678");
}
 

DockerSession(
    Container container,
    Tracer tracer,
    HttpClient client,
    SessionId id,
    URL url,
    Capabilities capabilities,
    Dialect downstream,
    Dialect upstream) {
  super(tracer, client, id, url, downstream, upstream, capabilities);
  this.container = Require.nonNull("Container", container);
}
 

private CommandCodec<HttpRequest> getCommandCodec(Dialect dialect) {
  switch (dialect) {
    case OSS:
      return new JsonHttpCommandCodec();

    case W3C:
      return new W3CHttpCommandCodec();

    default:
      throw new IllegalStateException("Unknown dialect: " + dialect);
  }
}
 

private ResponseCodec<HttpResponse> getResponseCodec(Dialect dialect) {
  switch (dialect) {
    case OSS:
      return new JsonHttpResponseCodec();

    case W3C:
      return new W3CHttpResponseCodec();

    default:
      throw new IllegalStateException("Unknown dialect: " + dialect);
  }
}
 

public ServicedSession(
    DriverService service,
    Dialect downstream,
    Dialect upstream,
    HttpHandler codec,
    SessionId id,
    Map<String, Object> capabilities) {
  super(downstream, upstream, codec, id, capabilities);

  this.service = service;

  new JMXHelper().register(this);
}
 

protected abstract ActiveSession newActiveSession(
X additionalData,
Dialect downstream,
Dialect upstream,
HttpHandler codec,
SessionId id,
Map<String, Object> capabilities);
 

public CreateSessionRequest(
    Set<Dialect> downstreamDialects,
    Capabilities capabilities,
    Map<String, Object> metadata) {
  this.downstreamDialects = ImmutableSet.copyOf(
      Require.nonNull("Downstream dialects", downstreamDialects));
  this.capabilities = ImmutableCapabilities.copyOf(Require.nonNull("Capabilities", capabilities));
  this.metadata = ImmutableMap.copyOf(Require.nonNull("Metadata", metadata));
}
 

private static CreateSessionRequest fromJson(JsonInput input) {
  Set<Dialect> downstreamDialects = null;
  Capabilities capabilities = null;
  Map<String, Object> metadata = null;

  input.beginObject();
  while (input.hasNext()) {
    switch (input.nextName()) {
      case "capabilities":
        capabilities = input.read(Capabilities.class);
        break;

      case "downstreamDialects":
        downstreamDialects = input.read(new TypeToken<Set<Dialect>>(){}.getType());
        break;

      case "metadata":
        metadata = input.read(MAP_TYPE);
        break;

      default:
        input.skipValue();
    }
  }
  input.endObject();

  return new CreateSessionRequest(downstreamDialects, capabilities, metadata);
}
 

public static ResponseEncoder<Session, Map<String, Object>, byte[]> getEncoder(Dialect dialect) {
  switch (dialect) {
    case OSS:
      return JWP_ENCODER;

    case W3C:
      return W3C_ENCODER;

    default:
      throw new IllegalArgumentException("Unrecognised dialect: " + dialect);
  }
}
 

@Test
public void shouldNotInstantiateSessionIfCapabilitiesDoNotMatch() {
  DriverServiceSessionFactory factory = factoryFor("chrome", builder);

  Optional<ActiveSession> session = factory.apply(new CreateSessionRequest(
      ImmutableSet.of(Dialect.W3C), toPayload("firefox"), ImmutableMap.of()));

  assertThat(session).isEmpty();
  verifyNoInteractions(builder);
}
 

@Test
public void shouldNotInstantiateSessionIfBuilderCanNotBuildService() {
  when(builder.build()).thenThrow(new WebDriverException());

  DriverServiceSessionFactory factory = factoryFor("chrome", builder);

  Optional<ActiveSession> session = factory.apply(new CreateSessionRequest(
      ImmutableSet.of(Dialect.W3C), toPayload("chrome"), ImmutableMap.of()));

  assertThat(session).isEmpty();
  verify(builder, times(1)).build();
  verifyNoMoreInteractions(builder);
}
 

@Override
public Optional<ActiveSession> apply(CreateSessionRequest sessionRequest) {
  SessionId id = new SessionId(UUID.randomUUID());
  Session session = sessionGenerator.apply(id, sessionRequest.getCapabilities());

  URL url = null;
  try {
    url = session.getUri().toURL();
  } catch (MalformedURLException e) {
    throw new UncheckedIOException(e);
  }

  Dialect downstream = sessionRequest.getDownstreamDialects().contains(W3C) ?
                       W3C :
                       sessionRequest.getDownstreamDialects().iterator().next();


  BaseActiveSession activeSession = new BaseActiveSession(
      session.getId(),
      url,
      downstream,
      W3C,
      session.getCapabilities()) {
    @Override
    public void stop() {
      // Do nothing
    }

    @Override
    public HttpResponse execute(HttpRequest req) throws UncheckedIOException {
      if (session instanceof HttpHandler) {
        return ((HttpHandler) session).execute(req);
      } else {
        // Do nothing.
        return new HttpResponse().setStatus(HTTP_NOT_FOUND).setContent(utf8String("No handler found for " + req));
      }
    }
  };
  return Optional.of(activeSession);
}
 

@Test
public void canBindNewFactoriesAtRunTime() {
  ActiveSession session = Mockito.mock(ActiveSession.class);

  ActiveSessionFactory sessionFactory = new ActiveSessionFactory(new NullTracer())
      .bind(caps ->
                "cheese".equals(caps.getBrowserName()),
            new SessionFactory() {
              @Override
              public boolean test(Capabilities capabilities) {
                return true;
              }

              @Override
              public Optional<ActiveSession> apply(CreateSessionRequest sessionRequest) {
                return Optional.of(session);
              }
            });

  ActiveSession created = sessionFactory.apply(
      new CreateSessionRequest(
          ImmutableSet.copyOf(Dialect.values()),
          toPayload("cheese"),
          ImmutableMap.of()))
      .get();

  assertSame(session, created);
}
 

@Test
public void shouldThrowAnExceptionIfMoreThanOneFileIsSent() throws Exception {
  ActiveSession session = mock(ActiveSession.class);
  when(session.getId()).thenReturn(new SessionId("1234567"));
  when(session.getFileSystem()).thenReturn(tempFs);
  when(session.getDownstreamDialect()).thenReturn(Dialect.OSS);

  File baseDir = Files.createTempDir();
  touch(baseDir, "example");
  touch(baseDir, "unwanted");
  String encoded = Zip.zip(baseDir);

  UploadFile uploadFile = new UploadFile(new Json(), session);
  Map<String, Object> args = ImmutableMap.of("file", encoded);
  HttpRequest request = new HttpRequest(HttpMethod.POST, "/session/%d/se/file");
  request.setContent(asJson(args));
  HttpResponse response = uploadFile.execute(request);

  try {
    new ErrorHandler(false).throwIfResponseFailed(
        new Json().toType(string(response), Response.class),
        100);
    fail("Should not get this far");
  } catch (WebDriverException ignored) {
    // Expected
  }
}