Fixing the .NET WebSocket message handler

This change fixes how responses are mapped to the request and propagated back to the caller, as well as correcting the use of Try* methods to check for return values. Signed-off-by: Jim Evans <[email protected]>
SeleniumHQ · Jan 10, 2020 · 5f38a0d · 5f38a0d
1 parent 5d02494
commit 5f38a0d
Show file tree

Hide file tree

Showing 2 changed files with 63 additions and 29 deletions.
diff --git a/dotnet/src/webdriver/DevTools/AutoGenerated/DevToolsSession.cs b/dotnet/src/webdriver/DevTools/AutoGenerated/DevToolsSession.cs
@@ -1,3 +1,4 @@
+
 namespace OpenQA.Selenium.DevTools
 {
     using System;
@@ -159,31 +160,52 @@ public string EndpointAddress
             var responseWasReceived = await Task.Run(() => message.SyncEvent.Wait(millisecondsTimeout.Value, cancellationToken));
 
             if (!responseWasReceived && throwExceptionIfResponseNotReceived)
+            {
                 throw new InvalidOperationException($"A command response was not received: {commandName}");
+            }
 
-            DevToolsCommandData modified;
-            m_pendingCommands.TryRemove(message.CommandId, out modified);
-            if (modified.IsError)
+            if (m_pendingCommands.TryRemove(message.CommandId, out var modified))
             {
-                var errorMessage = modified.Result.Value<string>("message");
-                var errorData = modified.Result.Value<string>("data");
+                if (modified.IsError)
+                {
+                    var errorMessage = modified.Result.Value<string>("message");
+                    var errorData = modified.Result.Value<string>("data");
 
-                var exceptionMessage = $"{commandName}: {errorMessage}";
-                if (!String.IsNullOrWhiteSpace(errorData))
-                    exceptionMessage = $"{exceptionMessage} - {errorData}";
+                    var exceptionMessage = $"{commandName}: {errorMessage}";
+                    if (!String.IsNullOrWhiteSpace(errorData))
+                        exceptionMessage = $"{exceptionMessage} - {errorData}";
 
-                LogTrace("Recieved Error Response {0}: {1} {2}", modified.CommandId, message, errorData);
-                throw new CommandResponseException(exceptionMessage)
-                {
-                    Code = modified.Result.Value<long>("code")
-                };
+                    LogTrace("Recieved Error Response {0}: {1} {2}", modified.CommandId, message, errorData);
+                    throw new CommandResponseException(exceptionMessage)
+                    {
+                        Code = modified.Result.Value<long>("code")
+                    };
+                }
+
+                return modified.Result;
             }
 
-            return modified.Result;
+            return null;
         }
 
         private async Task OpenSessionConnection(CancellationToken cancellationToken)
         {
+            if (m_sessionSocket == null)
+            {
+                return;
+            }
+
+            // Try to prevent "System.InvalidOperationException: The WebSocket has already been started."
+            while (m_sessionSocket.State == WebSocketState.Connecting)
+            {
+                if (cancellationToken.IsCancellationRequested)
+                {
+                    return;
+                }
+
+                await Task.Delay(10);
+            }
+
             if (m_sessionSocket.State != WebSocketState.Open)
             {
                 await m_sessionSocket.ConnectAsync(new Uri(m_endpointAddress), cancellationToken);
@@ -247,26 +269,35 @@ private void ProcessIncomingMessage(string message)
         {
             var messageObject = JObject.Parse(message);
 
-            long commandId;
-            if (messageObject.TryGetValue("id", out JToken idProperty) && (commandId = idProperty.Value<long>()) == m_currentCommandId)
+            if (messageObject.TryGetValue("id", out var idProperty))
             {
+                var commandId = idProperty.Value<long>();
+
                 DevToolsCommandData commandInfo;
-                m_pendingCommands.TryGetValue(commandId, out commandInfo);
-                if (messageObject.TryGetValue("error", out JToken errorProperty))
+                if (m_pendingCommands.TryGetValue(commandId, out commandInfo))
                 {
-                    commandInfo.IsError = true;
-                    commandInfo.Result = errorProperty;
+                    if (messageObject.TryGetValue("error", out var errorProperty))
+                    {
+                        commandInfo.IsError = true;
+                        commandInfo.Result = errorProperty;
+                    }
+                    else
+                    {
+                        commandInfo.Result = messageObject["result"];
+                        LogTrace("Recieved Response {0}: {1}", commandId, commandInfo.Result.ToString());
+                    }
+
                     commandInfo.SyncEvent.Set();
-                    return;
+                }
+                else
+                {
+                    LogError("Recieved Unknown Response {0}: {1}", commandId, message);
                 }
 
-                commandInfo.Result = messageObject["result"];
-                LogTrace("Recieved Response {0}: {1}", commandId, commandInfo.Result.ToString());
-                commandInfo.SyncEvent.Set();
                 return;
             }
 
-            if (messageObject.TryGetValue("method", out JToken methodProperty))
+            if (messageObject.TryGetValue("method", out var methodProperty))
             {
                 var method = methodProperty.Value<string>();
                 var methodParts = method.Split(new char[] { '.' }, 2);

diff --git a/dotnet/src/webdriver/Remote/RemoteWebDriver.cs b/dotnet/src/webdriver/Remote/RemoteWebDriver.cs
@@ -1006,7 +1006,7 @@ protected void StartSession(ICapabilities desiredCapabilities)
                 string errorMessage = string.Format(CultureInfo.InvariantCulture, "The new session command returned a value ('{0}') that is not a valid JSON object.", response.Value);
                 throw new WebDriverException(errorMessage);
             }
-            
+
             ReturnedCapabilities returnedCapabilities = new ReturnedCapabilities(rawCapabilities);
             this.capabilities = returnedCapabilities;
             this.sessionId = new SessionId(response.SessionId);
@@ -1044,16 +1044,19 @@ protected virtual Response Execute(string driverCommandToExecute, Dictionary<str
         {
             Command commandToExecute = new Command(this.sessionId, driverCommandToExecute, parameters);
 
-            Response commandResponse = new Response();
+            Response commandResponse;
 
             try
             {
                 commandResponse = this.executor.Execute(commandToExecute);
             }
             catch (System.Net.WebException e)
             {
-                commandResponse.Status = WebDriverResult.UnhandledError;
-                commandResponse.Value = e;
+                commandResponse = new Response
+                {
+                    Status = WebDriverResult.UnhandledError,
+                    Value = e
+                };
             }
 
             if (commandResponse.Status != WebDriverResult.Success)