dotnet/src/webdriver/DevTools/WebSocketConnection.cs - external/github.com/SeleniumHQ/selenium - Git at Google

 // <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;

 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));
     }
 }
	// <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;

	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));
	}
	}