diff --git a/src/StevesBot.Worker.Tests/Unit/DiscordGatewayClientOptionsTests.cs b/src/StevesBot.Worker.Tests/Unit/DiscordClientOptionsTests.cs similarity index 73% rename from src/StevesBot.Worker.Tests/Unit/DiscordGatewayClientOptionsTests.cs rename to src/StevesBot.Worker.Tests/Unit/DiscordClientOptionsTests.cs index 42cf692..cd2f5c5 100644 --- a/src/StevesBot.Worker.Tests/Unit/DiscordGatewayClientOptionsTests.cs +++ b/src/StevesBot.Worker.Tests/Unit/DiscordClientOptionsTests.cs @@ -1,14 +1,14 @@ namespace StevesBot.Worker.Tests.Unit; -public class DiscordGatewayClientOptionsTests +public class DiscordClientOptionsTests { [Fact] public void Constructor_WhenCalledWithoutParameters_ItShouldCreateInstance() { - var options = new DiscordGatewayClientOptions(); + var options = new DiscordClientOptions(); options.Should().NotBeNull(); - options.Should().BeOfType(); + options.Should().BeOfType(); options.ApiUrl.Should().Be(string.Empty); options.AppToken.Should().Be(string.Empty); options.Intents.Should().Be(0); @@ -21,7 +21,7 @@ public class DiscordGatewayClientOptionsTests var appToken = "test-token"; var intents = 123; - var options = new DiscordGatewayClientOptions + var options = new DiscordClientOptions { ApiUrl = apiUrl, AppToken = appToken, @@ -29,7 +29,7 @@ public class DiscordGatewayClientOptionsTests }; options.Should().NotBeNull(); - options.Should().BeOfType(); + options.Should().BeOfType(); options.ApiUrl.Should().Be(apiUrl); options.AppToken.Should().Be(appToken); options.Intents.Should().Be(intents); diff --git a/src/StevesBot.Worker.Tests/Unit/DiscordGatewayClientTests.cs b/src/StevesBot.Worker.Tests/Unit/DiscordGatewayClientTests.cs new file mode 100644 index 0000000..0a50e97 --- /dev/null +++ b/src/StevesBot.Worker.Tests/Unit/DiscordGatewayClientTests.cs @@ -0,0 +1,98 @@ +namespace StevesBot.Worker.Tests.Unit; + +public sealed class DiscordGatewayClientTests : IDisposable +{ + private readonly Mock _mockDiscordRestClient = new(); + private readonly Mock _mockWebSocketFactory = new(); + private readonly Mock> _mockLogger = new(); + private readonly DiscordClientOptions _options = new(); + private readonly DiscordGatewayClient _discordGatewayClient; + + public DiscordGatewayClientTests() + { + _discordGatewayClient = new DiscordGatewayClient( + _options, + _mockWebSocketFactory.Object, + _mockLogger.Object, + _mockDiscordRestClient.Object + ); + } + + [Fact] + public async Task ConnectAsync_WhenCalled_ItShouldConnectToGateway() + { + _mockDiscordRestClient + .Setup(static x => x.GetGatewayUrlAsync(It.IsAny())) + .ReturnsAsync("wss://gateway.discord.gg"); + + var mockWebSocket = new Mock(); + + _mockWebSocketFactory + .Setup(static x => x.Create()) + .Returns(mockWebSocket.Object); + + await _discordGatewayClient.ConnectAsync(CancellationToken.None); + + mockWebSocket + .Verify( + static x => x.ConnectAsync(It.IsAny(), It.IsAny()), + Times.Once + ); + } + + [Fact] + public async Task ConnectAsync_WhenCalledAndWebSocketIsClosed_ItShouldStopReceivingMessages() + { + _mockDiscordRestClient + .Setup(static x => x.GetGatewayUrlAsync(It.IsAny())) + .ReturnsAsync("wss://gateway.discord.gg"); + + var mockWebSocket = new Mock(); + + mockWebSocket + .Setup(static x => x.State) + .Returns(WebSocketState.Closed); + + _mockWebSocketFactory + .Setup(static x => x.Create()) + .Returns(mockWebSocket.Object); + + await _discordGatewayClient.ConnectAsync(CancellationToken.None); + + mockWebSocket.Verify( + static x => x.ReceiveAsync(It.IsAny>(), It.IsAny()), + Times.Never + ); + } + + [Fact] + public async Task ConnectAsync_WhenCalledAndCancellationIsRequested_ItShouldStopReceivingMessages() + { + _mockDiscordRestClient + .Setup(static x => x.GetGatewayUrlAsync(It.IsAny())) + .ReturnsAsync("wss://gateway.discord.gg"); + + var mockWebSocket = new Mock(); + + _mockWebSocketFactory + .Setup(static x => x.Create()) + .Returns(mockWebSocket.Object); + + using var cts = new CancellationTokenSource(); + + await _discordGatewayClient.ConnectAsync(cts.Token); + await Task.Delay(100); + await cts.CancelAsync(); + + mockWebSocket.Verify(static x => x.State, Times.AtMostOnce); + mockWebSocket.Verify( + static x => x.ReceiveAsync(It.IsAny>(), It.IsAny()), + Times.AtMostOnce + ); + } + + public void Dispose() + { + _discordGatewayClient.Dispose(); + } +} \ No newline at end of file diff --git a/src/StevesBot.Worker.Tests/Unit/DiscordRestClientExceptionTests.cs b/src/StevesBot.Worker.Tests/Unit/DiscordRestClientExceptionTests.cs new file mode 100644 index 0000000..5860b21 --- /dev/null +++ b/src/StevesBot.Worker.Tests/Unit/DiscordRestClientExceptionTests.cs @@ -0,0 +1,34 @@ +namespace StevesBot.Worker.Tests.Unit; + +public class DiscordRestClientExceptionTests +{ + [Fact] + public void Constructor_WhenCalledWithNoParameters_ItShouldCreateInstance() + { + var exception = new DiscordRestClientException(); + + exception.Should().NotBeNull(); + } + + [Fact] + public void Constructor_WhenCalledWithMessage_ItShouldCreateInstance() + { + var message = "Test message"; + var exception = new DiscordRestClientException(message); + + exception.Should().NotBeNull(); + exception.Message.Should().Be(message); + } + + [Fact] + public void Constructor_WhenCalledWithMessageAndInnerException_ItShouldCreateInstance() + { + var message = "Test message"; + var innerException = new Exception("Inner exception"); + var exception = new DiscordRestClientException(message, innerException); + + exception.Should().NotBeNull(); + exception.Message.Should().Be(message); + exception.InnerException.Should().Be(innerException); + } +} \ No newline at end of file diff --git a/src/StevesBot.Worker.Tests/Unit/DiscordRestClientTests.cs b/src/StevesBot.Worker.Tests/Unit/DiscordRestClientTests.cs index 4c7235e..6458516 100644 --- a/src/StevesBot.Worker.Tests/Unit/DiscordRestClientTests.cs +++ b/src/StevesBot.Worker.Tests/Unit/DiscordRestClientTests.cs @@ -1,17 +1,9 @@ -using System.Net; - -using Microsoft.Extensions.Logging; - -using Moq; - -using RichardSzalay.MockHttp; - namespace StevesBot.Worker.Tests.Unit; public sealed class DiscordRestClientTests : IDisposable { private const string BaseUrl = "https://discord.com/api/v10"; - private static string GatewayUrl => $"{BaseUrl}/gateway"; + private static string GatewayEndpoint => $"{BaseUrl}/gateway"; private readonly Mock> _loggerMock = new(); private readonly MockHttpMessageHandler _mockHttpMessageHandler; @@ -29,7 +21,7 @@ public sealed class DiscordRestClientTests : IDisposable public async Task GetGatewayUrlAsync_WhenRequestFails_ItShouldThrowException() { _mockHttpMessageHandler - .When(GatewayUrl) + .When(GatewayEndpoint) .Respond(HttpStatusCode.InternalServerError); var act = async () => await _discordRestClient.GetGatewayUrlAsync(CancellationToken.None); @@ -41,7 +33,7 @@ public sealed class DiscordRestClientTests : IDisposable public async Task GetGatewayUrlAsync_WhenResponseIsNull_ItShouldThrowException() { _mockHttpMessageHandler - .When(GatewayUrl) + .When(GatewayEndpoint) .Respond(HttpStatusCode.OK, "application/json", "null"); var act = async () => await _discordRestClient.GetGatewayUrlAsync(CancellationToken.None); @@ -58,7 +50,7 @@ public sealed class DiscordRestClientTests : IDisposable }}"; _mockHttpMessageHandler - .When(GatewayUrl) + .When(GatewayEndpoint) .Respond(HttpStatusCode.OK, "application/json", jsonResponse); var result = await _discordRestClient.GetGatewayUrlAsync(CancellationToken.None); diff --git a/src/StevesBot.Worker.Tests/Usings.cs b/src/StevesBot.Worker.Tests/Usings.cs index 6bc9191..fe8d80d 100644 --- a/src/StevesBot.Worker.Tests/Usings.cs +++ b/src/StevesBot.Worker.Tests/Usings.cs @@ -1,3 +1,4 @@ +global using System.Net; global using System.Net.WebSockets; global using System.Text; global using System.Text.Json; @@ -9,6 +10,11 @@ global using Microsoft.AspNetCore.Hosting.Server.Features; global using Microsoft.AspNetCore.Http.Features; global using Microsoft.Extensions.DependencyInjection; global using Microsoft.Extensions.Hosting; +global using Microsoft.Extensions.Logging; + +global using Moq; + +global using RichardSzalay.MockHttp; global using StevesBot.Worker.Discord; global using StevesBot.Worker.Discord.Events; diff --git a/src/StevesBot.Worker/Discord/DiscordGatewayClientOptions.cs b/src/StevesBot.Worker/Discord/DiscordClientOptions.cs similarity index 78% rename from src/StevesBot.Worker/Discord/DiscordGatewayClientOptions.cs rename to src/StevesBot.Worker/Discord/DiscordClientOptions.cs index 5dae7ab..3d3d11f 100644 --- a/src/StevesBot.Worker/Discord/DiscordGatewayClientOptions.cs +++ b/src/StevesBot.Worker/Discord/DiscordClientOptions.cs @@ -1,6 +1,6 @@ namespace StevesBot.Worker.Discord; -internal sealed class DiscordGatewayClientOptions +internal sealed class DiscordClientOptions { public string ApiUrl { get; init; } = string.Empty; public string AppToken { get; init; } = string.Empty; diff --git a/src/StevesBot.Worker/Discord/DiscordGatewayClient.cs b/src/StevesBot.Worker/Discord/DiscordGatewayClient.cs index 1ae3f27..b696506 100644 --- a/src/StevesBot.Worker/Discord/DiscordGatewayClient.cs +++ b/src/StevesBot.Worker/Discord/DiscordGatewayClient.cs @@ -1,34 +1,37 @@ +using System.Text; + namespace StevesBot.Worker.Discord; internal class DiscordGatewayClient : IDiscordGatewayClient { - private readonly DiscordGatewayClientOptions _options; + private readonly DiscordClientOptions _options; private readonly IWebSocketFactory _webSocketFactory; private readonly ILogger _logger; private readonly IDiscordRestClient _discordRestClient; - // private readonly JsonSerializerOptions _jsonSerializerOptions = new() - // { - // PropertyNamingPolicy = JsonNamingPolicy.CamelCase, - // ReferenceHandler = ReferenceHandler.IgnoreCycles, - // Converters = - // { - // new DiscordEventConverter(), - // }, - // }; + private readonly JsonSerializerOptions _jsonSerializerOptions = new() + { + PropertyNamingPolicy = JsonNamingPolicy.CamelCase, + ReferenceHandler = ReferenceHandler.IgnoreCycles, + Converters = + { + new DiscordEventConverter(), + }, + }; private readonly AsyncLock _lock = new(); private string _gatewayUrl = string.Empty; private IWebSocket? _webSocket; private Task? _receiveTask; - // private DateTime _timeLastHeartbeatSent = DateTime.MinValue; - // private DateTime _timeLastHeartbeatAcknowledged = DateTime.MinValue; - // private CancellationTokenSource? _heartbeatCts; - // private Task? _heartbeatTask; + private DateTime _timeLastHeartbeatSent = DateTime.MinValue; + private DateTime _timeLastHeartbeatAcknowledged = DateTime.MinValue; + private CancellationTokenSource? _heartbeatCts; + private CancellationTokenSource? _linkedCts; + private Task? _heartbeatTask; // private string _sessionId = string.Empty; // private string _resumeGatewayUrl = string.Empty; public DiscordGatewayClient( - DiscordGatewayClientOptions options, + DiscordClientOptions options, IWebSocketFactory webSocketFactory, ILogger logger, IDiscordRestClient discordRestClient @@ -60,16 +63,221 @@ internal class DiscordGatewayClient : IDiscordGatewayClient } } - private static Task ReceiveMessagesAsync(CancellationToken cancellationToken) + private async Task ReceiveMessagesAsync(CancellationToken cancellationToken) { - return Task.CompletedTask; + try + { + var messageBuffer = new byte[8192]; + + while (cancellationToken.IsCancellationRequested is false) + { + using (await _lock.LockAsync(cancellationToken)) + { + if (_webSocket?.State is not WebSocketState.Open) + { + _logger.LogWarning("WebSocket is not open. Cannot receive messages."); + return; + } + } + + // websocket message might be larger than the + // size of the buffer so we need to loop until + // we receive the end of the message and + // write the data we receive on each iteration + // to the memory stream + using var memoryStream = new MemoryStream(); + WebSocketReceiveResult result; + + do + { + result = await _webSocket.ReceiveAsync(new(messageBuffer), cancellationToken); + + if (result.MessageType is WebSocketMessageType.Close) + { + // TODO: If close status is 1000 or 1001 we cannot resume. + // if 1000 or 1001 we should close the connection and reconnect + // else we should close the connection and attempt to resume + + if (result.CloseStatus is WebSocketCloseStatus.NormalClosure or WebSocketCloseStatus.EndpointUnavailable) + { + return; + } + + return; + } + + if (result.MessageType is WebSocketMessageType.Text) + { +# pragma warning disable CA1849 + memoryStream.Write(messageBuffer, 0, result.Count); + } + + } while (result.EndOfMessage is false); + + memoryStream.Seek(0, SeekOrigin.Begin); + + var message = Encoding.UTF8.GetString(messageBuffer, 0, result.Count); + var e = await JsonSerializer.DeserializeAsync( + memoryStream, + _jsonSerializerOptions, + cancellationToken + ); + + if (e is null) + { + _logger.LogInformation("Received null event."); + continue; + } + + await HandleEventAsync(e, cancellationToken); + } + } + catch (WebSocketException ex) + { + _logger.LogError(ex, "WebSocket error"); + throw new DiscordGatewayClientException("WebSocket error.", ex); + } + catch (OperationCanceledException ex) + { + _logger.LogInformation(ex, "Receive messages operation canceled: {Message}", ex.Message); + } + catch (Exception ex) + { + _logger.LogError(ex, "Unexpected error while receiving messages"); + throw new DiscordGatewayClientException("Unexpected error while receiving messages.", ex); + } + } + + private async Task HandleEventAsync(DiscordEvent e, CancellationToken cancellationToken) + { + if (e is HelloDiscordEvent he) + { + await StartHeartbeatAsync(he, cancellationToken); + return; + } + + if (e is HeartbeatAckDiscordEvent hae) + { + using (await _lock.LockAsync(cancellationToken)) + { + _timeLastHeartbeatAcknowledged = DateTime.UtcNow; + } + + _logger.LogInformation("Heartbeat acknowledged at {Time}", _timeLastHeartbeatAcknowledged); + return; + } + } + + private async Task StartHeartbeatAsync(HelloDiscordEvent helloEvent, CancellationToken cancellationToken) + { + using (await _lock.LockAsync(cancellationToken)) + { + if (_heartbeatCts is not null) + { + await _heartbeatCts.CancelAsync(); + _heartbeatCts.Dispose(); + } + + _heartbeatCts = new CancellationTokenSource(); + _linkedCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, _heartbeatCts.Token); + + _heartbeatTask = Task.Run(async () => + { + _logger.LogInformation("Starting heartbeat task."); + + while (_linkedCts.Token.IsCancellationRequested is false) + { + using (await _lock.LockAsync(_linkedCts.Token)) + { + if (_webSocket?.State is not WebSocketState.Open) + { + _logger.LogWarning("WebSocket is not open. Cannot send heartbeat."); + break; + } + } + + try + { +# pragma warning disable CA5394 // Do not use insecure randomness + var jitter = Random.Shared.NextDouble(); +# pragma warning restore CA5394 // Do not use insecure randomness + await Task.Delay((int)(helloEvent.Data.HeartbeatInterval + jitter), _linkedCts.Token); + + using (await _lock.LockAsync(_linkedCts.Token)) + { + if (_timeLastHeartbeatAcknowledged < _timeLastHeartbeatSent) + { + if (_webSocket?.State is WebSocketState.Open) + { + _logger.LogWarning("Heartbeat not acknowledged. Closing WebSocket."); + await _webSocket.CloseAsync(WebSocketCloseStatus.NormalClosure, "Heartbeat not acknowledged", CancellationToken.None); + } + break; + } + + _timeLastHeartbeatSent = await SendHeartbeatAsync(helloEvent.Sequence, _linkedCts.Token); + } + + _logger.LogInformation("Heartbeat sent at {Time}", _timeLastHeartbeatSent); + } + catch (OperationCanceledException ex) + { + _logger.LogInformation(ex, "Heartbeat task canceled"); + } +# pragma warning disable CA1031 // Do not catch general exception types + catch (Exception ex) +# pragma warning restore CA1031 // Do not catch general exception types + { + _logger.LogError(ex, "Error in heartbeat task: {Message}", ex.Message); + // TODO: Attempt to reconnect + } + } + }, _linkedCts.Token); + } + } + + private async Task SendHeartbeatAsync(int? sequence, CancellationToken cancellationToken) + { + var heartbeat = new HeartbeatDiscordEvent(sequence); + await SendJsonAsync(heartbeat, cancellationToken); + return DateTime.UtcNow; + } + + private async Task SendJsonAsync(object data, CancellationToken cancellationToken) + { + using (await _lock.LockAsync(cancellationToken)) + { + if (_webSocket?.State is not WebSocketState.Open) + { + _logger.LogWarning("WebSocket is not open. Cannot send message."); + return; + } + } + + try + { + var json = JsonSerializer.Serialize(data, _jsonSerializerOptions); + var bytes = Encoding.UTF8.GetBytes(json); + var buffer = new ArraySegment(bytes); + await _webSocket.SendAsync(buffer, WebSocketMessageType.Text, true, cancellationToken); + } + catch (OperationCanceledException ex) + { + _logger.LogInformation(ex, "Send operation canceled"); + } + catch (Exception ex) + { + _logger.LogError(ex, "Failed to send message: {Message}", ex.Message); + throw new DiscordGatewayClientException("Failed to send message.", ex); + } } public void Dispose() { - // _heartbeatCts?.Cancel(); - // _heartbeatCts?.Dispose(); - // _heartbeatTask?.Dispose(); + _heartbeatCts?.Cancel(); + _heartbeatCts?.Dispose(); + _linkedCts?.Dispose(); + _heartbeatTask?.Dispose(); _receiveTask?.Dispose(); _webSocket?.Dispose(); _lock.Dispose();