Bug 1829337 - [bidi] Implement the "session.end" command. r=webdriver-reviewers,whimboo

Depends on D178621

Differential Revision: https://phabricator.services.mozilla.com/D179671

remote/cdp/CDPConnection.sys.mjs

@@ -192,14 +192,14 @@ export class CDPConnection extends WebSocketConnection {
   /**
    * Called by the `transport` when the connection is closed.
    */
-  onClosed() {
+  onConnectionClose() {
     // Cleanup all the registered sessions.
     for (const session of this.sessions.values()) {
       session.destructor();
     }
     this.sessions.clear();
 
-    super.onClosed();
+    super.onConnectionClose();
   }
 
   /**

remote/server/WebSocketTransport.sys.mjs

@@ -41,20 +41,64 @@ WebSocketTransport.prototype = {
     throw new Error("Bulk send is not supported by WebSocket transport");
   },
 
+  /**
+   * Force closing the active connection and WebSocket.
+   */
   close() {
     if (!this.socket) {
       return;
     }
     this.emit("close");
+
+    if (this.hooks) {
+      this.hooks.onConnectionClose();
+    }
+
+    this.active = false;
+
+    this.socket.removeEventListener("message", this);
+
+    // Remove the listener that is used when the connection
+    // is closed because we already emitted the 'close' event.
+    // Instead listen for the closing of the WebSocket.
+    this.socket.removeEventListener("close", this);
+    this.socket.addEventListener("close", this.onSocketClose.bind(this));
+
+    // Close socket with code `1000` for a normal closure.
+    this.socket.close(1000);
+  },
+
+  /**
+   * Callback for socket on close event,
+   * it is used in case websocket was closed by the client.
+   */
+  onClose() {
+    if (!this.socket) {
+      return;
+    }
+    this.emit("close");
+
+    if (this.hooks) {
+      this.hooks.onConnectionClose();
+      this.hooks.onSocketClose();
+      this.hooks = null;
+    }
+
     this.active = false;
 
     this.socket.removeEventListener("message", this);
     this.socket.removeEventListener("close", this);
-    this.socket.close();
+    this.socket = null;
+  },
+
+  /**
+   * Callback which is called when we can be sure that websocket is closed.
+   */
+  onSocketClose() {
     this.socket = null;
 
     if (this.hooks) {
-      this.hooks.onClosed();
+      this.hooks.onSocketClose();
       this.hooks = null;
     }
   },
@@ -65,7 +109,7 @@ WebSocketTransport.prototype = {
         this.onMessage(event);
         break;
       case "close":
-        this.close();
+        this.onClose();
         break;
     }
   },

remote/shared/WebSocketConnection.sys.mjs

@@ -62,11 +62,6 @@ export class WebSocketConnection {
    */
   close() {
     this.transport.close();
-
-    // In addition to the WebSocket transport, we also have to close the
-    // connection used internally within httpd.js. Otherwise the server doesn't
-    // shut down correctly, and keeps these Connection instances alive.
-    this.httpdConnection.close();
   }
 
   /**
@@ -128,10 +123,20 @@ export class WebSocketConnection {
   /**
    * Called by the `transport` when the connection is closed.
    */
-  onClosed(status) {
+  onConnectionClose(status) {
     lazy.logger.debug(`${this.constructor.name} ${this.id} closed`);
   }
 
+  /**
+   * Called when the socket is closed.
+   */
+  onSocketClose() {
+    // In addition to the WebSocket transport, we also have to close the
+    // connection used internally within httpd.js. Otherwise the server doesn't
+    // shut down correctly, and keeps these Connection instances alive.
+    this.httpdConnection.close();
+  }
+
   /**
    * Receive a packet from the WebSocket layer.
    *

remote/webdriver-bidi/WebDriverBiDiConnection.sys.mjs

@@ -127,10 +127,10 @@ export class WebDriverBiDiConnection extends WebSocketConnection {
   /**
    * Called by the `transport` when the connection is closed.
    */
-  onClosed() {
+  onConnectionClose() {
     this.unregisterSession();
 
-    super.onClosed();
+    super.onConnectionClose();
   }
 
   /**
@@ -194,6 +194,14 @@ export class WebDriverBiDiConnection extends WebSocketConnection {
       }
 
       this.sendResult(id, result);
+
+      // Session clean up.
+      if (module === "session" && command === "end") {
+        // TODO Bug 1838269. Implement session ending logic
+        // for the case of classic + bidi session.
+        // We currently only support one session, see Bug 1720707.
+        lazy.RemoteAgent.webDriverBiDi.deleteSession();
+      }
     } catch (e) {
       this.sendError(id, e);
     }

remote/webdriver-bidi/modules/root/session.sys.mjs

@@ -11,6 +11,7 @@ ChromeUtils.defineESModuleGetters(lazy, {
   ContextDescriptorType:
     "chrome://remote/content/shared/messagehandler/MessageHandler.sys.mjs",
   error: "chrome://remote/content/shared/webdriver/Errors.sys.mjs",
+  Marionette: "chrome://remote/content/components/Marionette.sys.mjs",
   RootMessageHandler:
     "chrome://remote/content/shared/messagehandler/RootMessageHandler.sys.mjs",
   TabManager: "chrome://remote/content/shared/TabManager.sys.mjs",
@@ -45,6 +46,19 @@ class SessionModule extends Module {
    * Commands
    */
 
+  /**
+   * End the current session.
+   *
+   * Session clean up will happen later in WebDriverBiDiConnection class.
+   */
+  async end() {
+    if (lazy.Marionette.running) {
+      throw new lazy.error.UnsupportedOperationError(
+        "Ending session which was started with Webdriver classic is not supported, use Webdriver classic delete command instead."
+      );
+    }
+  }
+
   /**
    * Enable certain events either globally, or for a list of browsing contexts.
    *