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 8c9ee61e..c2685276 100644
--- a/LiveControlGateway/Websocket/FlatbuffersWebsocketBaseController.cs
+++ b/LiveControlGateway/Websocket/FlatbuffersWebsocketBaseController.cs
@@ -43,85 +43,60 @@ 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()
{
- while (!LinkedToken.IsCancellationRequested)
+ try
{
- try
+ while (!LinkedToken.IsCancellationRequested)
{
if (WebSocket is null or { State: WebSocketState.Aborted }) return;
var message =
await FlatbufferWebSocketUtils.ReceiveFullMessageAsyncNonAlloc(WebSocket,
_incomingSerializer, LinkedToken);
-
- // All is good, normal message, deserialize and handle
- if (message.IsT0)
- {
- try
+
+
+ bool ok = await message.Match(
+ Handle,
+ async _ =>
{
- var serverMessage = message.AsT0;
- await Handle(serverMessage);
- }
- catch (Exception e)
+ await WebSocket.CloseAsync(WebSocketCloseStatus.ProtocolError, "Invalid flatbuffers message", LinkedToken);
+ return false;
+ },
+ async _ =>
{
- Logger.LogError(e, "Error while handling device message");
- }
- }
-
- // Deserialization failed, log and continue
- else if (message.IsT1)
- {
- Logger.LogWarning(message.AsT1.Exception, "Deserialization failed for websocket message");
- }
+ if (WebSocket.State != WebSocketState.Open)
+ {
+ Logger.LogTrace("Client sent closure, but connection state is not open");
+ return false;
+ }
- // Device sent closure, close connection
- else if (message.IsT2)
- {
- if (WebSocket.State != WebSocketState.Open)
- {
- Logger.LogTrace("Client sent closure, but connection state is not open");
- break;
- }
+ await WebSocket.CloseAsync(WebSocketCloseStatus.NormalClosure, "Normal close", LinkedToken);
- try
- {
- await WebSocket.CloseAsync(WebSocketCloseStatus.NormalClosure, "Normal close",
- LinkedToken);
- }
- catch (OperationCanceledException e)
- {
- Logger.LogError(e, "Error during close handshake");
- }
+ Logger.LogInformation("Closing websocket connection");
+ return false;
+ });
- Logger.LogInformation("Closing websocket connection");
- break;
- }
- else
- {
- 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");
+ if (!ok) break;
}
- 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