WebSocketConnection.cs

// <copyright file="WebSocketConnection.cs" company="Selenium Committers">
// Licensed to the Software Freedom Conservancy (SFC) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The SFC licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.
// </copyright>

using System;
using System.Net.WebSockets;
using System.Text;
using System.Threading;
using System.Threading.Tasks;

#nullable enable

namespace OpenQA.Selenium.DevTools
{
    /// <summary>
    /// Represents a connection to a WebDriver Bidi remote end.
    /// </summary>
    public class WebSocketConnection
    {
        private static readonly TimeSpan DefaultTimeout = TimeSpan.FromSeconds(10);
        private readonly CancellationTokenSource clientTokenSource = new CancellationTokenSource();
        private readonly TimeSpan startupTimeout;
        private readonly TimeSpan shutdownTimeout;
        private Task? dataReceiveTask;
        private ClientWebSocket client = new ClientWebSocket();
        private readonly SemaphoreSlim sendMethodSemaphore = new SemaphoreSlim(1, 1);

        /// <summary>
        /// Initializes a new instance of the <see cref="WebSocketConnection" /> class.
        /// </summary>
        public WebSocketConnection()
            : this(DefaultTimeout)
        {
        }

        /// <summary>
        /// Initializes a new instance of the <see cref="WebSocketConnection" /> class with a given startup timeout.
        /// </summary>
        /// <param name="startupTimeout">The timeout before throwing an error when starting up the connection.</param>
        public WebSocketConnection(TimeSpan startupTimeout)
            : this(startupTimeout, DefaultTimeout)
        {
        }

        /// <summary>
        /// Initializes a new instance of the <see cref="WebSocketConnection" /> class with a given startup and shutdown timeout.
        /// </summary>
        /// <param name="startupTimeout">The timeout before throwing an error when starting up the connection.</param>
        /// <param name="shutdownTimeout">The timeout before throwing an error when shutting down the connection.</param>
        public WebSocketConnection(TimeSpan startupTimeout, TimeSpan shutdownTimeout)
        {
            this.startupTimeout = startupTimeout;
            this.shutdownTimeout = shutdownTimeout;
        }

        /// <summary>
        /// Occurs when data is received from this connection.
        /// </summary>
        public event EventHandler<WebSocketConnectionDataReceivedEventArgs>? DataReceived;

        /// <summary>
        /// Occurs when a log message is emitted from this connection.
        /// </summary>
        public event EventHandler<DevToolsSessionLogMessageEventArgs>? LogMessage;

        /// <summary>
        /// Gets a value indicating whether this connection is active.
        /// </summary>
        public bool IsActive { get; private set; } = false;

        /// <summary>
        /// Gets the buffer size for communication used by this connection.
        /// </summary>
        public int BufferSize { get; } = 4096;

        /// <summary>
        /// Asynchronously starts communication with the remote end of this connection.
        /// </summary>
        /// <param name="url">The URL used to connect to the remote end.</param>
        /// <returns>The task object representing the asynchronous operation.</returns>
        /// <exception cref="TimeoutException">Thrown when the connection is not established within the startup timeout.</exception>
        public virtual async Task Start(string url)
        {
            if (url is null)
            {
                throw new ArgumentNullException(nameof(url));
            }

            this.Log($"Opening connection to URL {url}", DevToolsSessionLogLevel.Trace);
            bool connected = false;
            DateTime timeout = DateTime.Now.Add(this.startupTimeout);
            while (!connected && DateTime.Now <= timeout)
            {
                try
                {
                    await this.client.ConnectAsync(new Uri(url), this.clientTokenSource.Token).ConfigureAwait(false);
                    connected = true;
                }
                catch (WebSocketException)
                {
                    // If the server-side socket is not yet ready, it leaves the client socket in a closed state,
                    // which sees the object as disposed, so we must create a new one to try again
                    await Task.Delay(TimeSpan.FromMilliseconds(500)).ConfigureAwait(false);
                    this.client = new ClientWebSocket();
                }
            }

            if (!connected)
            {
                throw new TimeoutException($"Could not connect to browser within {this.startupTimeout.TotalSeconds} seconds");
            }

            this.dataReceiveTask = Task.Run(async () => await this.ReceiveData());
            this.IsActive = true;
            this.Log($"Connection opened", DevToolsSessionLogLevel.Trace);
        }

        /// <summary>
        /// Asynchronously stops communication with the remote end of this connection.
        /// </summary>
        /// <returns>The task object representing the asynchronous operation.</returns>
        public virtual async Task Stop()
        {
            this.Log($"Closing connection", DevToolsSessionLogLevel.Trace);
            if (this.client.State != WebSocketState.Open)
            {
                this.Log($"Socket already closed (Socket state: {this.client.State})");
            }
            else
            {
                await this.CloseClientWebSocket().ConfigureAwait(false);
            }

            // Whether we closed the socket or timed out, we cancel the token causing ReceiveAsync to abort the socket.
            // The finally block at the end of the processing loop will dispose of the ClientWebSocket object.
            this.clientTokenSource.Cancel();
            if (this.dataReceiveTask != null)
            {
                await this.dataReceiveTask.ConfigureAwait(false);
            }

            this.client.Dispose();
        }

        /// <summary>
        /// Asynchronously sends data to the remote end of this connection.
        /// </summary>
        /// <param name="data">The data to be sent to the remote end of this connection.</param>
        /// <returns>The task object representing the asynchronous operation.</returns>
        public virtual async Task SendData(string data)
        {
            if (data is null)
            {
                throw new ArgumentNullException(nameof(data));
            }

            ArraySegment<byte> messageBuffer = new ArraySegment<byte>(Encoding.UTF8.GetBytes(data));
            this.Log($"SEND >>> {data}");

            await sendMethodSemaphore.WaitAsync().ConfigureAwait(false);

            try
            {
                await this.client.SendAsync(messageBuffer, WebSocketMessageType.Text, endOfMessage: true, CancellationToken.None).ConfigureAwait(false);
            }
            finally
            {
                sendMethodSemaphore.Release();
            }
        }

        /// <summary>
        /// Asynchronously closes the client WebSocket.
        /// </summary>
        /// <returns>The task object representing the asynchronous operation.</returns>
        protected virtual async Task CloseClientWebSocket()
        {
            // Close the socket first, because ReceiveAsync leaves an invalid socket (state = aborted) when the token is cancelled
            CancellationTokenSource timeout = new CancellationTokenSource(this.shutdownTimeout);
            try
            {
                // After this, the socket state which change to CloseSent
                await this.client.CloseOutputAsync(WebSocketCloseStatus.NormalClosure, "Closing", timeout.Token).ConfigureAwait(false);

                // Now we wait for the server response, which will close the socket
                while (this.client.State != WebSocketState.Closed && this.client.State != WebSocketState.Aborted && !timeout.Token.IsCancellationRequested)
                {
                    // The loop may be too tight for the cancellation token to get triggered, so add a small delay
                    await Task.Delay(TimeSpan.FromMilliseconds(10)).ConfigureAwait(false);
                }

                this.Log($"Client state is {this.client.State}", DevToolsSessionLogLevel.Trace);
            }
            catch (OperationCanceledException)
            {
                // An OperationCanceledException is normal upon task/token cancellation, so disregard it
            }
            catch (WebSocketException e)
            {
                this.Log($"Unexpected error during attempt at close: {e.Message}", DevToolsSessionLogLevel.Error);
            }
        }

        /// <summary>
        /// Raises the DataReceived event.
        /// </summary>
        /// <param name="e">The event args used when raising the event.</param>
        protected virtual void OnDataReceived(WebSocketConnectionDataReceivedEventArgs e)
        {
            if (this.DataReceived != null)
            {
                this.DataReceived(this, e);
            }
        }

        /// <summary>
        /// Raises the LogMessage event.
        /// </summary>
        /// <param name="e">The event args used when raising the event.</param>
        protected virtual void OnLogMessage(DevToolsSessionLogMessageEventArgs e)
        {
            if (this.LogMessage != null)
            {
                this.LogMessage(this, e);
            }
        }

        private async Task ReceiveData()
        {
            CancellationToken cancellationToken = this.clientTokenSource.Token;
            try
            {
                StringBuilder messageBuilder = new StringBuilder();
                ArraySegment<byte> buffer = WebSocket.CreateClientBuffer(this.BufferSize, this.BufferSize);
                while (this.client.State != WebSocketState.Closed && !cancellationToken.IsCancellationRequested)
                {
                    WebSocketReceiveResult receiveResult = await this.client.ReceiveAsync(buffer, cancellationToken).ConfigureAwait(false);

                    // If the token is cancelled while ReceiveAsync is blocking, the socket state changes to aborted and it can't be used
                    if (!cancellationToken.IsCancellationRequested)
                    {
                        // The server is notifying us that the connection will close, and we did
                        // not initiate the close; send acknowledgement
                        if (receiveResult.MessageType == WebSocketMessageType.Close && this.client.State != WebSocketState.Closed && this.client.State != WebSocketState.CloseSent)
                        {
                            this.Log($"Acknowledging Close frame received from server (client state: {this.client.State})", DevToolsSessionLogLevel.Trace);
                            await this.client.CloseOutputAsync(WebSocketCloseStatus.NormalClosure, "Acknowledge Close frame", CancellationToken.None).ConfigureAwait(false);
                        }

                        // Display text or binary data
                        if (this.client.State == WebSocketState.Open && receiveResult.MessageType != WebSocketMessageType.Close)
                        {
                            messageBuilder.Append(Encoding.UTF8.GetString(buffer.Array!, 0, receiveResult.Count));
                            if (receiveResult.EndOfMessage)
                            {
                                string message = messageBuilder.ToString();
                                messageBuilder = new StringBuilder();
                                if (message.Length > 0)
                                {
                                    this.Log($"RECV <<< {message}");
                                    this.OnDataReceived(new WebSocketConnectionDataReceivedEventArgs(message));
                                }
                            }
                        }
                    }
                }

                this.Log($"Ending processing loop in state {this.client.State}", DevToolsSessionLogLevel.Trace);
            }
            catch (OperationCanceledException)
            {
                // An OperationCanceledException is normal upon task/token cancellation, so disregard it
            }
            catch (WebSocketException e)
            {
                this.Log($"Unexpected error during receive of data: {e.Message}", DevToolsSessionLogLevel.Error);
            }
            finally
            {
                this.IsActive = false;
            }
        }

        private void Log(string message)
        {
            this.Log(message, DevToolsSessionLogLevel.Trace);
        }

        private void Log(string message, DevToolsSessionLogLevel level)
        {
            this.OnLogMessage(new DevToolsSessionLogMessageEventArgs(level, "[{0}] {1}", "Connection", message));
        }
    }
}