Solar/Swarm
2
1
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

🐛 Bug fixes and improvements

Browse Source
This commit is contained in:
LittleSheep
2025-05-01 00:47:26 +08:00
parent 758186f674
commit 84a88222bd
13 changed files with 111 additions and 59 deletions
Download Patch File Download Diff File Expand all files Collapse all files

31
DysonNetwork.Sphere/Connection/WebSocketController.cs
View File

@ -1,6 +1,5 @@
using System.Collections.Concurrent;
using System.Net.WebSockets;
using DysonNetwork.Sphere.Account;
using Microsoft.AspNetCore.Authorization;
using Microsoft.AspNetCore.Mvc;
using Swashbuckle.AspNetCore.Annotations;
@ -11,13 +10,10 @@ namespace DysonNetwork.Sphere.Connection;
[Route("/ws")]
public class WebSocketController : ControllerBase
{
// Concurrent dictionary to store active WebSocket connections.
// Key: Tuple (AccountId, DeviceId); Value: WebSocket and CancellationTokenSource
private static readonly ConcurrentDictionary<
(long AccountId, string DeviceId),
(WebSocket Socket, CancellationTokenSource Cts)
> ActiveConnections =
new ConcurrentDictionary<(long, string), (WebSocket, CancellationTokenSource)>();
> ActiveConnections = new();
[Route("/ws")]
[Authorize]
@ -26,18 +22,18 @@ public class WebSocketController : ControllerBase
{
if (HttpContext.WebSockets.IsWebSocketRequest)
{
// Get AccountId from HttpContext
if (
!HttpContext.Items.TryGetValue("CurrentUser", out var currentUserValue)
|| currentUserValue is not Account.Account currentUser
)
HttpContext.Items.TryGetValue("CurrentUser", out var currentUserValue);
HttpContext.Items.TryGetValue("CurrentSession", out var currentSessionValue);
if (currentUserValue is not Account.Account currentUser ||
currentSessionValue is not Auth.Session currentSession)
{
HttpContext.Response.StatusCode = StatusCodes.Status401Unauthorized;
return;
}
long accountId = currentUser.Id;
// Verify deviceId
var accountId = currentUser.Id;
var deviceId = currentSession.Challenge.DeviceId;
if (string.IsNullOrEmpty(deviceId))
{
HttpContext.Response.StatusCode = StatusCodes.Status400BadRequest;
@ -45,14 +41,11 @@ public class WebSocketController : ControllerBase
}
using var webSocket = await HttpContext.WebSockets.AcceptWebSocketAsync();
// Create a CancellationTokenSource for this connection
var cts = new CancellationTokenSource();
var connectionKey = (accountId, deviceId);
// Add the connection to the active connections dictionary
if (!ActiveConnections.TryAdd(connectionKey, (webSocket, cts)))
{
// Failed to add
await webSocket.CloseAsync(
WebSocketCloseStatus.InternalServerError,
"Failed to establish connection.",
@ -71,7 +64,6 @@ public class WebSocketController : ControllerBase
}
finally
{
// Connection is closed, remove it from the active connections dictionary
ActiveConnections.TryRemove(connectionKey, out _);
cts.Dispose();
}
@ -88,25 +80,21 @@ public class WebSocketController : ControllerBase
CancellationToken cancellationToken
)
{
// Buffer for receiving messages.
var buffer = new byte[1024 * 4];
try
{
// We don't handle receiving data, so we ignore the return.
var receiveResult = await webSocket.ReceiveAsync(
new ArraySegment<byte>(buffer),
cancellationToken
);
while (!receiveResult.CloseStatus.HasValue)
{
// Keep connection alive and wait for close requests
receiveResult = await webSocket.ReceiveAsync(
new ArraySegment<byte>(buffer),
cancellationToken
);
}
// Close connection
await webSocket.CloseAsync(
receiveResult.CloseStatus.Value,
receiveResult.CloseStatusDescription,
@ -144,5 +132,4 @@ public class WebSocketController : ControllerBase
);
}
}
}
}