Ensure that the protocol converter handles the new session responses …

…properly

The "shape" of the New Session response is different between the
JSON Wire Protocol and the W3C spec. Handle this difference in
the protocol converter.

java/client/src/org/openqa/selenium/remote/http/Contents.java

@@ -21,6 +21,7 @@
 
 import com.google.common.io.ByteStreams;
 import com.google.common.io.FileBackedOutputStream;
+import org.openqa.selenium.json.Json;
 
 import java.io.ByteArrayInputStream;
 import java.io.ByteArrayOutputStream;
@@ -35,6 +36,8 @@
 
 public class Contents {
 
+  private static final Json JSON = new Json();
+
   private Contents() {
     // Utility class
   }
@@ -106,6 +109,13 @@ public static Reader reader(HttpMessage<?> message) {
     return reader(message.getContent(), message.getContentEncoding());
   }
 
+  /**
+   * @return an {@link InputStream} containing the object converted to a UTF-8 JSON string.
+   */
+  public static Supplier<InputStream> asJson(Object obj) {
+    return utf8String(JSON.toJson(obj));
+  }
+
   public static Supplier<InputStream> memoize(Supplier<InputStream> delegate) {
     return new MemoizedSupplier(delegate);
   }

java/server/src/org/openqa/selenium/grid/web/ProtocolConverter.java

@@ -18,10 +18,14 @@
 package org.openqa.selenium.grid.web;
 
 import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Preconditions;
+import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
+import org.openqa.selenium.json.Json;
 import org.openqa.selenium.remote.Command;
 import org.openqa.selenium.remote.CommandCodec;
 import org.openqa.selenium.remote.Dialect;
+import org.openqa.selenium.remote.DriverCommand;
 import org.openqa.selenium.remote.JsonToWebElementConverter;
 import org.openqa.selenium.remote.Response;
 import org.openqa.selenium.remote.ResponseCodec;
@@ -35,11 +39,20 @@
 import org.openqa.selenium.remote.http.HttpResponse;
 
 import java.io.UncheckedIOException;
+import java.net.HttpURLConnection;
 import java.util.Map;
 import java.util.Objects;
+import java.util.function.Function;
+
+import static java.net.HttpURLConnection.HTTP_OK;
+import static org.openqa.selenium.json.Json.MAP_TYPE;
+import static org.openqa.selenium.remote.Dialect.W3C;
+import static org.openqa.selenium.remote.http.Contents.asJson;
+import static org.openqa.selenium.remote.http.Contents.string;
 
 public class ProtocolConverter implements HttpHandler {
 
+  private final static Json JSON = new Json();
   private final static ImmutableSet<String> IGNORED_REQ_HEADERS = ImmutableSet.<String>builder()
       .add("connection")
       .add("keep-alive")
@@ -58,6 +71,7 @@ public class ProtocolConverter implements HttpHandler {
   private final ResponseCodec<HttpResponse> downstreamResponse;
   private final ResponseCodec<HttpResponse> upstreamResponse;
   private final JsonToWebElementConverter converter;
+  private final Function<HttpResponse, HttpResponse> newSessionConverter;
 
   public ProtocolConverter(
       HttpClient client,
@@ -74,6 +88,8 @@ public ProtocolConverter(
     this.upstreamResponse = getResponseCodec(upstream);
 
     converter = new JsonToWebElementConverter(null);
+
+    newSessionConverter = downstream == W3C ? this::createW3CNewSessionResponse : this::createJwpNewSessionResponse;
   }
 
   @Override
@@ -91,10 +107,19 @@ public HttpResponse execute(HttpRequest req) throws UncheckedIOException {
 
     HttpResponse res = makeRequest(request);
 
-    Response decoded = upstreamResponse.decode(res);
-    HttpResponse toReturn = downstreamResponse.encode(HttpResponse::new, decoded);
+    HttpResponse toReturn;
+    if (DriverCommand.NEW_SESSION.equals(command.getName()) && res.getStatus() == HTTP_OK) {
+      toReturn = newSessionConverter.apply(res);
+    } else {
+      Response decoded = upstreamResponse.decode(res);
+      toReturn = downstreamResponse.encode(HttpResponse::new, decoded);
+    }
 
-    IGNORED_REQ_HEADERS.forEach(toReturn::removeHeader);
+    res.getHeaderNames().forEach(name -> {
+      if (!IGNORED_REQ_HEADERS.contains(name)) {
+        res.getHeaders(name).forEach(value -> toReturn.addHeader(name, value));
+      }
+    });
 
     return toReturn;
   }
@@ -130,4 +155,31 @@ private ResponseCodec<HttpResponse> getResponseCodec(Dialect dialect) {
     }
   }
 
+  private HttpResponse createW3CNewSessionResponse(HttpResponse response) {
+    Map<String, Object> value = JSON.toType(string(response), MAP_TYPE);
+
+    Preconditions.checkState(value.get("sessionId") != null);
+    Preconditions.checkState(value.get("value") instanceof Map);
+
+    return new HttpResponse()
+      .setContent(asJson(ImmutableMap.of(
+        "value", ImmutableMap.of(
+          "sessionId", value.get("sessionId"),
+          "capabilities", value.get("value")))));
+  }
+
+  private HttpResponse createJwpNewSessionResponse(HttpResponse response) {
+    Map<String, Object> value = Objects.requireNonNull(Values.get(response, MAP_TYPE));
+
+    // Check to see if the values we need are set
+    Preconditions.checkState(value.get("sessionId") != null);
+    Preconditions.checkState(value.get("capabilities") instanceof Map);
+
+    return new HttpResponse()
+      .setContent(asJson(ImmutableMap.of(
+        "status", 0,
+        "sessionId", value.get("sessionId"),
+        "value", value.get("capabilities"))));
+  }
+
 }

java/server/test/org/openqa/selenium/grid/web/BUILD.bazel

@@ -1,8 +1,11 @@
 load("//java:rules.bzl", "java_test_suite")
 
 java_test_suite(
-    name = "MediumTests",
-    size = "medium",
+    name = "SmallTests",
+    size = "small",
+    tags = [
+        "no-sandbox",
+    ],
     srcs = glob(["*Test.java"]),
     deps = [
         "//java/client/src/org/openqa/selenium/json",
@@ -13,5 +16,6 @@ java_test_suite(
         "//third_party/java/assertj",
         "//third_party/java/guava",
         "//third_party/java/junit",
+        "//third_party/java/mockito:mockito-core",
     ],
 )

java/server/test/org/openqa/selenium/grid/web/ProtocolConverterTest.java

@@ -22,6 +22,7 @@
 import com.google.common.collect.ImmutableMap;
 import com.google.common.net.MediaType;
 import org.junit.Test;
+import org.mockito.Mockito;
 import org.openqa.selenium.WebDriverException;
 import org.openqa.selenium.json.Json;
 import org.openqa.selenium.remote.Command;
@@ -41,15 +42,20 @@
 import java.util.Map;
 
 import static java.net.HttpURLConnection.HTTP_INTERNAL_ERROR;
+import static org.assertj.core.api.Assertions.assertThat;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.mock;
 import static org.openqa.selenium.json.Json.MAP_TYPE;
 import static org.openqa.selenium.remote.Dialect.OSS;
 import static org.openqa.selenium.remote.Dialect.W3C;
 import static org.openqa.selenium.remote.ErrorCodes.UNHANDLED_ERROR;
+import static org.openqa.selenium.remote.http.Contents.asJson;
 import static org.openqa.selenium.remote.http.Contents.string;
 import static org.openqa.selenium.remote.http.Contents.utf8String;
+import static org.openqa.selenium.remote.http.HttpMethod.POST;
 
 
 public class ProtocolConverterTest {
@@ -204,4 +210,53 @@ protected HttpResponse makeRequest(HttpRequest request) {
     assertTrue(((String) value.get("message")).startsWith("I love cheese and peas"));
   }
 
+  @Test
+  public void newJwpSessionResponseShouldBeCorrectlyConvertedToW3C() {
+    Map<String, Object> jwpNewSession = ImmutableMap.of("desiredCapabilities", ImmutableMap.of("cheese", "brie"));
+
+    Map<String, Object> w3cResponse = ImmutableMap.of(
+      "value", ImmutableMap.of(
+        "sessionId", "i like cheese very much",
+        "capabilities", ImmutableMap.of("cheese", "brie")));
+
+    HttpClient client = mock(HttpClient.class);
+    Mockito.when(client.execute(any())).thenReturn(new HttpResponse().setContent(asJson(w3cResponse)));
+
+    ProtocolConverter converter = new ProtocolConverter(client, OSS, W3C);
+
+    HttpResponse response = converter.execute(new HttpRequest(POST, "/session").setContent(asJson(jwpNewSession)));
+
+    Map<String, Object> convertedResponse = json.toType(string(response), MAP_TYPE);
+
+    assertThat(convertedResponse.get("sessionId")).isEqualTo("i like cheese very much");
+    assertThat(convertedResponse.get("status")).isEqualTo(0L);
+    assertThat(convertedResponse.get("value")).isEqualTo(ImmutableMap.of("cheese", "brie"));
+  }
+
+  @Test
+  public void newW3CSessionResponseShouldBeCorrectlyConvertedToJwp() {
+    Map<String, Object> w3cNewSession = ImmutableMap.of(
+      "capabilities", ImmutableMap.of());
+
+    Map<String, Object> jwpResponse = ImmutableMap.of(
+      "status", 0,
+      "sessionId", "i like cheese very much",
+      "value", ImmutableMap.of("cheese", "brie"));
+
+    HttpClient client = mock(HttpClient.class);
+    Mockito.when(client.execute(any())).thenReturn(new HttpResponse().setContent(asJson(jwpResponse)));
+
+    ProtocolConverter converter = new ProtocolConverter(client, W3C, OSS);
+
+    HttpResponse response = converter.execute(new HttpRequest(POST, "/session").setContent(asJson(w3cNewSession)));
+
+    Map<String, Object> convertedResponse = json.toType(string(response), MAP_TYPE);
+    Map<?, ?> value = (Map<?, ?>) convertedResponse.get("value");
+    assertThat(value.get("capabilities"))
+      .as("capabilities: " + convertedResponse)
+      .isEqualTo(jwpResponse.get("value"));
+    assertThat(value.get("sessionId"))
+      .as("session id: " + convertedResponse)
+      .isEqualTo(jwpResponse.get("sessionId"));
+  }
 }