From af6c3aa9bcfc4ab3c7acf9f5f1d5a2ccab4d6fef Mon Sep 17 00:00:00 2001 From: hhvrc Date: Fri, 18 Apr 2025 23:05:12 +0200 Subject: [PATCH 1/2] Kick client immidiatley on exception --- .../FlatbuffersWebsocketBaseController.cs | 69 +++++++------------ 1 file changed, 26 insertions(+), 43 deletions(-) diff --git a/LiveControlGateway/Websocket/FlatbuffersWebsocketBaseController.cs b/LiveControlGateway/Websocket/FlatbuffersWebsocketBaseController.cs index 8c9ee61e..f0bad877 100644 --- a/LiveControlGateway/Websocket/FlatbuffersWebsocketBaseController.cs +++ b/LiveControlGateway/Websocket/FlatbuffersWebsocketBaseController.cs @@ -49,9 +49,9 @@ protected override Task SendWebSocketMessage(TOut message, WebSocket websocket, /// protected override async Task Logic() { - while (!LinkedToken.IsCancellationRequested) + try { - try + while (!LinkedToken.IsCancellationRequested) { if (WebSocket is null or { State: WebSocketState.Aborted }) return; var message = @@ -61,25 +61,20 @@ await FlatbufferWebSocketUtils.ReceiveFullMessageAsyncNonAlloc(WebSocket, // All is good, normal message, deserialize and handle if (message.IsT0) { - try - { - var serverMessage = message.AsT0; - await Handle(serverMessage); - } - catch (Exception e) - { - Logger.LogError(e, "Error while handling device message"); - } + var serverMessage = message.AsT0; + await Handle(serverMessage); + continue; } - // Deserialization failed, log and continue - else if (message.IsT1) + // Deserialization failed, log and exit + if (message.IsT1) { - Logger.LogWarning(message.AsT1.Exception, "Deserialization failed for websocket message"); + await WebSocket.CloseAsync(WebSocketCloseStatus.ProtocolError, "Invalid flatbuffers message", LinkedToken); + break; } // Device sent closure, close connection - else if (message.IsT2) + if (message.IsT2) { if (WebSocket.State != WebSocketState.Open) { @@ -87,41 +82,29 @@ await FlatbufferWebSocketUtils.ReceiveFullMessageAsyncNonAlloc(WebSocket, break; } - try - { - await WebSocket.CloseAsync(WebSocketCloseStatus.NormalClosure, "Normal close", - LinkedToken); - } - catch (OperationCanceledException e) - { - Logger.LogError(e, "Error during close handshake"); - } + await WebSocket.CloseAsync(WebSocketCloseStatus.NormalClosure, "Normal close", LinkedToken); Logger.LogInformation("Closing websocket connection"); break; } - else - { - throw new NotImplementedException(); // message.T? is not implemented - } + + throw new NotImplementedException(); // message.T? is not implemented } - catch (OperationCanceledException) - { - Logger.LogInformation("WebSocket connection terminated due to close or shutdown"); - break; - } - catch (WebSocketException e) - { - if (e.WebSocketErrorCode != WebSocketError.ConnectionClosedPrematurely) - Logger.LogError(e, "Error in receive loop, websocket exception"); - } - catch (Exception ex) + } + catch (OperationCanceledException) + { + Logger.LogInformation("WebSocket connection terminated due to close or shutdown"); + } + catch (WebSocketException e) + { + if (e.WebSocketErrorCode != WebSocketError.ConnectionClosedPrematurely) { - Logger.LogError(ex, "Exception while processing websocket request"); + Logger.LogError(e, "Error in receive loop, websocket exception"); } } - - await Close.CancelAsync(); + catch (Exception ex) + { + Logger.LogError(ex, "Exception while processing websocket request"); + } } - } \ No newline at end of file From 3fb08652ab5d243371c6a14b5f7692e4074f6466 Mon Sep 17 00:00:00 2001 From: hhvrc Date: Fri, 18 Apr 2025 23:19:30 +0200 Subject: [PATCH 2/2] More gentle error handling --- .../Controllers/HubControllerBase.cs | 7 +-- .../Controllers/HubV1Controller.cs | 13 +++-- .../Controllers/HubV2Controller.cs | 13 +++-- .../FlatbuffersWebsocketBaseController.cs | 54 ++++++++----------- 4 files changed, 45 insertions(+), 42 deletions(-) diff --git a/LiveControlGateway/Controllers/HubControllerBase.cs b/LiveControlGateway/Controllers/HubControllerBase.cs index 532037f1..fad00244 100644 --- a/LiveControlGateway/Controllers/HubControllerBase.cs +++ b/LiveControlGateway/Controllers/HubControllerBase.cs @@ -179,14 +179,13 @@ protected override async Task UnregisterConnection() /// /// Keep the hub online /// - protected async Task SelfOnline(ulong uptimeMs, ushort? latency = null, int? rssi = null) + protected async Task SelfOnline(ulong uptimeMs, ushort? latency = null, int? rssi = null) { var bootedAt = GetBootedAtFromUptimeMs(uptimeMs); if (!bootedAt.HasValue) { Logger.LogDebug("Client attempted to abuse reported boot time, uptime indicated that hub [{HubId}] booted prior to 2024", CurrentHub.Id); - await DisposeAsync(); - return; + return false; } Logger.LogDebug("Received keep alive from hub [{HubId}]", CurrentHub.Id); @@ -205,6 +204,8 @@ protected async Task SelfOnline(ulong uptimeMs, ushort? latency = null, int? rss LatencyMs = latency, Rssi = rssi }); + + return true; } /// diff --git a/LiveControlGateway/Controllers/HubV1Controller.cs b/LiveControlGateway/Controllers/HubV1Controller.cs index 8900f1e2..5ebc6822 100644 --- a/LiveControlGateway/Controllers/HubV1Controller.cs +++ b/LiveControlGateway/Controllers/HubV1Controller.cs @@ -55,9 +55,9 @@ ILogger logger private IUserHub HcOwner => _userHubContext.Clients.User(CurrentHub.Owner.ToString()); /// - protected override async Task Handle(HubToGatewayMessage data) + protected override async Task Handle(HubToGatewayMessage data) { - if(data.Payload == null) return; + if(!data.Payload.HasValue) return false; var payload = data.Payload.Value; await using var scope = ServiceProvider.CreateAsyncScope(); @@ -67,7 +67,10 @@ protected override async Task Handle(HubToGatewayMessage data) switch (payload.Kind) { case HubToGatewayMessagePayload.ItemKind.KeepAlive: - await SelfOnline(payload.KeepAlive.Uptime); + if (!await SelfOnline(payload.KeepAlive.Uptime)) + { + return false; + } break; case HubToGatewayMessagePayload.ItemKind.OtaInstallStarted: @@ -157,8 +160,10 @@ await otaService.Error(CurrentHub.Id, payload.BootStatus.OtaUpdateId, false, case HubToGatewayMessagePayload.ItemKind.NONE: default: Logger.LogWarning("Payload kind not defined [{Kind}]", payload.Kind); - break; + return false; } + + return true; } /// diff --git a/LiveControlGateway/Controllers/HubV2Controller.cs b/LiveControlGateway/Controllers/HubV2Controller.cs index 9a45960a..dd80dfe0 100644 --- a/LiveControlGateway/Controllers/HubV2Controller.cs +++ b/LiveControlGateway/Controllers/HubV2Controller.cs @@ -80,7 +80,7 @@ await QueueMessage(new GatewayToHubMessage private IUserHub HcOwner => _userHubContext.Clients.User(CurrentHub.Owner.ToString()); /// - protected override async Task Handle(HubToGatewayMessage data) + protected override async Task Handle(HubToGatewayMessage data) { var payload = data.Payload; @@ -96,12 +96,15 @@ protected override async Task Handle(HubToGatewayMessage data) if (_pingTimestamp == 0) { // TODO: Kick or warn client. - return; + return false; } _latencyMs = (ushort)Math.Min(Stopwatch.GetElapsedTime(_pingTimestamp).TotalMilliseconds, ushort.MaxValue); // If someone has a ping higher than 65 seconds, they are messing with us. Cap it to 65 seconds _pingTimestamp = 0; - await SelfOnline(payload.Pong.Uptime, _latencyMs, payload.Pong.Rssi); + if (!await SelfOnline(payload.Pong.Uptime, _latencyMs, payload.Pong.Rssi)) + { + return false; + } break; case HubToGatewayMessagePayload.ItemKind.OtaUpdateStarted: @@ -190,8 +193,10 @@ await otaService.Error(CurrentHub.Id, payload.BootStatus.OtaUpdateId, false, case HubToGatewayMessagePayload.ItemKind.NONE: default: Logger.LogWarning("Payload kind not defined [{Kind}]", payload.Kind); - break; + return false; } + + return true; } /// diff --git a/LiveControlGateway/Websocket/FlatbuffersWebsocketBaseController.cs b/LiveControlGateway/Websocket/FlatbuffersWebsocketBaseController.cs index f0bad877..c2685276 100644 --- a/LiveControlGateway/Websocket/FlatbuffersWebsocketBaseController.cs +++ b/LiveControlGateway/Websocket/FlatbuffersWebsocketBaseController.cs @@ -43,8 +43,8 @@ protected override Task SendWebSocketMessage(TOut message, WebSocket websocket, /// Handle the incoming message /// /// - /// - protected abstract Task Handle(TIn data); + /// false if message is invalid + protected abstract Task Handle(TIn data); /// protected override async Task Logic() @@ -57,38 +57,30 @@ protected override async Task Logic() var message = await FlatbufferWebSocketUtils.ReceiveFullMessageAsyncNonAlloc(WebSocket, _incomingSerializer, LinkedToken); - - // All is good, normal message, deserialize and handle - if (message.IsT0) - { - var serverMessage = message.AsT0; - await Handle(serverMessage); - continue; - } - - // Deserialization failed, log and exit - if (message.IsT1) - { - await WebSocket.CloseAsync(WebSocketCloseStatus.ProtocolError, "Invalid flatbuffers message", LinkedToken); - break; - } - - // Device sent closure, close connection - if (message.IsT2) - { - if (WebSocket.State != WebSocketState.Open) + + + bool ok = await message.Match( + Handle, + async _ => + { + await WebSocket.CloseAsync(WebSocketCloseStatus.ProtocolError, "Invalid flatbuffers message", LinkedToken); + return false; + }, + async _ => { - Logger.LogTrace("Client sent closure, but connection state is not open"); - break; - } + if (WebSocket.State != WebSocketState.Open) + { + Logger.LogTrace("Client sent closure, but connection state is not open"); + return false; + } - await WebSocket.CloseAsync(WebSocketCloseStatus.NormalClosure, "Normal close", LinkedToken); + await WebSocket.CloseAsync(WebSocketCloseStatus.NormalClosure, "Normal close", LinkedToken); - Logger.LogInformation("Closing websocket connection"); - break; - } - - throw new NotImplementedException(); // message.T? is not implemented + Logger.LogInformation("Closing websocket connection"); + return false; + }); + + if (!ok) break; } } catch (OperationCanceledException)