Make ActiveSession track the protocol dialect

This allows us to correctly format responses when
necessary.

Closes #4119

Author: shs96c

java/server/src/org/openqa/selenium/remote/server/ActiveSession.java

@@ -17,6 +17,7 @@
 
 package org.openqa.selenium.remote.server;
 
+import org.openqa.selenium.remote.Dialect;
 import org.openqa.selenium.remote.SessionId;
 
 import java.util.Map;
@@ -25,6 +26,10 @@ interface ActiveSession extends CommandHandler {
 
   SessionId getId();
 
+  Dialect getUpstreamDialect();
+
+  Dialect getDownstreamDialect();
+
   /**
    * Desribe the current webdriver session's capabilities.
    */

java/server/src/org/openqa/selenium/remote/server/BeginSession.java

@@ -39,8 +39,6 @@
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.net.MediaType;
-import com.google.gson.Gson;
-import com.google.gson.GsonBuilder;
 import com.google.gson.stream.JsonReader;
 import com.google.gson.stream.JsonToken;
 
@@ -73,23 +71,17 @@
 class BeginSession implements CommandHandler {
 
   private final static Logger LOG = Logger.getLogger(BeginSession.class.getName());
-  private final static Gson gson = new GsonBuilder().serializeNulls().create();
 
   private final Cache<SessionId, ActiveSession> allSessions;
-  private final DriverSessions legacySessions;
   private final Map<String, SessionFactory> factories;
-  private final Function<Path, ActiveSession> defaultFactory;
 
   public BeginSession(Cache<SessionId, ActiveSession> allSessions, DriverSessions legacySessions) {
     this.allSessions = allSessions;
-    this.legacySessions = legacySessions;
 
     this.factories = ImmutableMap.of(
         chrome().getBrowserName(), new ServicedSession.Factory("org.openqa.selenium.chrome.ChromeDriverService"),
         firefox().getBrowserName(), new ServicedSession.Factory("org.openqa.selenium.firefox.GeckoDriverService"),
         htmlUnit().getBrowserName(), new InMemorySession.Factory(legacySessions));
-
-    defaultFactory = null;
   }
 
   @Override
@@ -133,18 +125,24 @@ public void execute(HttpRequest req, HttpResponse resp) throws IOException {
       allSessions.put(session.getId(), session);
 
       Object toConvert;
-      if (firstMatch.isEmpty() && alwaysMatch.isEmpty()) {
-        // JSON Wire Protocol response
-        toConvert = ImmutableMap.of(
-            "status", 0,
-            "sessionId", session.getId().toString(),
-            "value", session.getCapabilities());
-      } else {
-        // W3C protocol response
-        toConvert = ImmutableMap.of(
-            "value", ImmutableMap.of(
-                "sessionId", session.getId().toString(),
-                "capabilities", session.getCapabilities()));
+      switch (session.getDownstreamDialect()) {
+        case OSS:
+          toConvert = ImmutableMap.of(
+              "status", 0,
+              "sessionId", session.getId().toString(),
+              "value", session.getCapabilities());
+          break;
+
+        case W3C:
+          toConvert = ImmutableMap.of(
+              "value", ImmutableMap.of(
+                  "sessionId", session.getId().toString(),
+                  "capabilities", session.getCapabilities()));
+          break;
+
+          default:
+            throw new SessionNotCreatedException(
+                "Unrecognized downstream dialect: " + session.getDownstreamDialect());
       }
 
       byte[] payload = new BeanToJsonConverter().convert(toConvert).getBytes(UTF_8);

java/server/src/org/openqa/selenium/remote/server/InMemorySession.java

@@ -18,8 +18,10 @@
 package org.openqa.selenium.remote.server;
 
 import static java.nio.charset.StandardCharsets.UTF_8;
+import static org.openqa.selenium.remote.Dialect.OSS;
 
 import com.google.common.base.Preconditions;
+import com.google.common.collect.Iterables;
 import com.google.common.io.CharStreams;
 
 import org.openqa.selenium.Capabilities;
@@ -54,18 +56,34 @@ class InMemorySession implements ActiveSession {
   private final SessionId id;
   private final Session session;
   private JsonHttpCommandHandler commandHandler;
+  private Dialect downstreamDialect;
 
-  public InMemorySession(SessionId id, Session session, JsonHttpCommandHandler commandHandler) {
+  public InMemorySession(
+      SessionId id,
+      Session session,
+      JsonHttpCommandHandler commandHandler,
+      Dialect downstreamDialect) {
     this.id = id;
     this.session = session;
     this.commandHandler = commandHandler;
+    this.downstreamDialect = downstreamDialect;
   }
 
   @Override
   public SessionId getId() {
     return id;
   }
 
+  @Override
+  public Dialect getUpstreamDialect() {
+    return OSS;
+  }
+
+  @Override
+  public Dialect getDownstreamDialect() {
+    return downstreamDialect;
+  }
+
   @Override
   public Map<String, Object> getCapabilities() {
     return session.getCapabilities().asMap().entrySet().stream()
@@ -108,7 +126,12 @@ public ActiveSession apply(Path path, Set<Dialect> downstreamDialects) {
         SessionId sessionId = legacySessions.newSession(caps);
         Session session = legacySessions.get(sessionId);
 
-        return new InMemorySession(sessionId, session, jsonHttpCommandHandler);
+        // Force OSS dialect if downstream speaks it
+        Dialect downstream = downstreamDialects.contains(OSS) ?
+                             OSS :
+                             Iterables.getFirst(downstreamDialects, null);
+
+        return new InMemorySession(sessionId, session, jsonHttpCommandHandler, downstream);
       } catch (Exception e) {
         throw new SessionNotCreatedException("Unable to create session", e);
       }

java/server/src/org/openqa/selenium/remote/server/Passthrough.java

@@ -24,6 +24,7 @@
 import com.google.common.io.CharStreams;
 import com.google.common.net.MediaType;
 
+import org.openqa.selenium.remote.Dialect;
 import org.openqa.selenium.remote.http.HttpRequest;
 import org.openqa.selenium.remote.http.HttpResponse;
 
@@ -58,9 +59,11 @@ class Passthrough implements SessionCodec {
       .build();
 
   private final URL upstream;
+  private Dialect dialect;
 
   public Passthrough(URL upstream) {
     this.upstream = upstream;
+    this.dialect = dialect;
   }
 
   @Override

java/server/src/org/openqa/selenium/remote/server/ServicedSession.java

@@ -60,15 +60,21 @@ class ServicedSession implements ActiveSession {
 
   private final DriverService service;
   private final SessionId id;
+  private Dialect downstream;
+  private Dialect upstream;
   private final SessionCodec codec;
   private final Map<String, Object> capabilities;
 
   public ServicedSession(
       DriverService service,
+      Dialect downstream,
+      Dialect upstream,
       SessionCodec codec,
       SessionId id,
       Map<String, Object> capabilities) {
     this.service = service;
+    this.downstream = downstream;
+    this.upstream = upstream;
     this.codec = codec;
     this.id = id;
     this.capabilities = capabilities;
@@ -85,6 +91,16 @@ public SessionId getId() {
     return id;
   }
 
+  @Override
+  public Dialect getUpstreamDialect() {
+    return upstream;
+  }
+
+  @Override
+  public Dialect getDownstreamDialect() {
+    return downstream;
+  }
+
   @Override
   public Map<String, Object> getCapabilities() {
     return capabilities;
@@ -138,24 +154,28 @@ public ActiveSession apply(Path path, Set<Dialect> downstreamDialects) {
             .orElseThrow(() -> new SessionNotCreatedException("Unable to create session"));
 
         SessionCodec codec;
+        Dialect upstream = result.getDialect();
+        Dialect downstream;
         if (downstreamDialects.contains(result.getDialect())) {
           codec = new Passthrough(url);
+          downstream = upstream;
         } else {
-
-          Dialect dialact = downstreamDialects.iterator().next();
+          downstream = downstreamDialects.iterator().next();
 
           codec = new ProtocolConverter(
               url,
-              getCommandCodec(dialact),
-              getResponseCodec(dialact),
-              getCommandCodec(result.getDialect()),
-              getResponseCodec(result.getDialect()));
+              getCommandCodec(downstream),
+              getResponseCodec(downstream),
+              getCommandCodec(upstream),
+              getResponseCodec(upstream));
         }
 
         Response response = result.createResponse();
         //noinspection unchecked
         return new ServicedSession(
             service,
+            downstream,
+            upstream,
             codec,
             new SessionId(response.getSessionId()),
             (Map<String, Object>) response.getValue());