From 9023509b01ccfa66c38ecfb7e882e40dc994023e Mon Sep 17 00:00:00 2001 From: Assistant builder Date: Thu, 11 Jun 2026 21:55:33 +0000 Subject: [PATCH] DeviceHub: /device WS with bearer auth + hello/hello_ack + ping/pong --- backend.tests/DeviceHubHandshakeTests.cs | 91 ++++++++++++++++++ backend/DeviceHub/DeviceAuth.cs | 21 ++++ backend/DeviceHub/DeviceHubEndpoint.cs | 116 +++++++++++++++++++++++ backend/Program.cs | 5 + 4 files changed, 233 insertions(+) create mode 100644 backend.tests/DeviceHubHandshakeTests.cs create mode 100644 backend/DeviceHub/DeviceAuth.cs create mode 100644 backend/DeviceHub/DeviceHubEndpoint.cs diff --git a/backend.tests/DeviceHubHandshakeTests.cs b/backend.tests/DeviceHubHandshakeTests.cs new file mode 100644 index 0000000..4142b12 --- /dev/null +++ b/backend.tests/DeviceHubHandshakeTests.cs @@ -0,0 +1,91 @@ +using System.Net.Http.Json; +using System.Net.WebSockets; +using System.Text; +using System.Text.Json; +using FluentAssertions; +using Xunit; + +namespace backend.tests; + +public class DeviceHubHandshakeTests(ApiFactory factory) : IClassFixture +{ + private readonly ApiFactory _factory = factory; + + private async Task PairAndGetTokenAsync(string email) + { + var http = _factory.CreateClient(); + await http.PostAsJsonAsync("/api/auth/register", + new { email, password = "Passw0rd!" }); + await http.PostAsJsonAsync("/api/auth/login", + new { email, password = "Passw0rd!" }); + var codeRes = await http.PostAsJsonAsync("/api/pair-code", new { name = "ws-test" }); + var code = (await codeRes.Content.ReadFromJsonAsync>())!["code"].ToString()!; + + var pairRes = await _factory.CreateClient().PostAsJsonAsync("/api/pair", + new { code, hostname = "smoke", client_version = "0.1" }); + var pairBody = await pairRes.Content.ReadFromJsonAsync>(); + return pairBody!["device_token"].ToString()!; + } + + [Fact] + public async Task Connect_without_token_fails_upgrade() + { + var client = _factory.Server.CreateWebSocketClient(); + var uri = new Uri(_factory.Server.BaseAddress, "device"); + var act = async () => await client.ConnectAsync(uri, CancellationToken.None); + await act.Should().ThrowAsync(); + } + + [Fact] + public async Task Hello_ack_includes_config() + { + var token = await PairAndGetTokenAsync("ws1@example.com"); + var client = _factory.Server.CreateWebSocketClient(); + client.ConfigureRequest = req => req.Headers["Authorization"] = $"Bearer {token}"; + + var uri = new Uri(_factory.Server.BaseAddress, "device"); + var ws = await client.ConnectAsync(uri, CancellationToken.None); + + await SendJson(ws, new { type = "hello", device_id = Guid.NewGuid(), client_version = "0.1" }); + var ack = await ReceiveJson(ws); + ack.GetProperty("type").GetString().Should().Be("hello_ack"); + ack.GetProperty("config").GetProperty("voice").GetString().Should().NotBeNullOrEmpty(); + ack.GetProperty("config").GetProperty("model").GetString().Should().NotBeNullOrEmpty(); + ack.GetProperty("config").GetProperty("enabled_tools").GetArrayLength().Should().BeGreaterThan(0); + } + + [Fact] + public async Task Ping_replies_with_pong() + { + var token = await PairAndGetTokenAsync("ws2@example.com"); + var client = _factory.Server.CreateWebSocketClient(); + client.ConfigureRequest = req => req.Headers["Authorization"] = $"Bearer {token}"; + var uri = new Uri(_factory.Server.BaseAddress, "device"); + var ws = await client.ConnectAsync(uri, CancellationToken.None); + + await SendJson(ws, new { type = "hello", device_id = Guid.NewGuid(), client_version = "0.1" }); + await ReceiveJson(ws); + await SendJson(ws, new { type = "ping" }); + var pong = await ReceiveJson(ws); + pong.GetProperty("type").GetString().Should().Be("pong"); + } + + private static async Task SendJson(WebSocket ws, object payload) + { + var bytes = Encoding.UTF8.GetBytes(JsonSerializer.Serialize(payload)); + await ws.SendAsync(bytes, WebSocketMessageType.Text, true, CancellationToken.None); + } + + private static async Task ReceiveJson(WebSocket ws) + { + var buf = new byte[8192]; + using var stream = new MemoryStream(); + while (true) + { + var r = await ws.ReceiveAsync(buf, CancellationToken.None); + stream.Write(buf, 0, r.Count); + if (r.EndOfMessage) break; + } + return JsonDocument.Parse(stream.ToArray()).RootElement.Clone(); + } +} diff --git a/backend/DeviceHub/DeviceAuth.cs b/backend/DeviceHub/DeviceAuth.cs new file mode 100644 index 0000000..fc9dc0d --- /dev/null +++ b/backend/DeviceHub/DeviceAuth.cs @@ -0,0 +1,21 @@ +using backend.Data; +using backend.Devices; +using Microsoft.EntityFrameworkCore; + +namespace backend.DeviceHub; + +public class DeviceAuth(AppDbContext db, DeviceTokenService tokens) +{ + public async Task AuthenticateAsync(string? authHeader, CancellationToken ct) + { + if (string.IsNullOrEmpty(authHeader)) return null; + const string prefix = "Bearer "; + if (!authHeader.StartsWith(prefix, StringComparison.Ordinal)) return null; + var token = authHeader[prefix.Length..].Trim(); + if (string.IsNullOrEmpty(token)) return null; + var hash = tokens.Hash(token); + var device = await db.Devices.Include(d => d.Config) + .FirstOrDefaultAsync(d => d.TokenHash == hash && !d.IsRevoked, ct); + return device; + } +} diff --git a/backend/DeviceHub/DeviceHubEndpoint.cs b/backend/DeviceHub/DeviceHubEndpoint.cs new file mode 100644 index 0000000..98f4c9d --- /dev/null +++ b/backend/DeviceHub/DeviceHubEndpoint.cs @@ -0,0 +1,116 @@ +using System.Net.WebSockets; +using System.Text; +using System.Text.Json; +using System.Text.Json.Nodes; +using backend.Data; +using backend.Devices; +using Microsoft.EntityFrameworkCore; + +namespace backend.DeviceHub; + +public static class DeviceHubEndpoint +{ + public static IEndpointRouteBuilder MapDeviceHub(this IEndpointRouteBuilder app) + { + app.Map("/device", async (HttpContext ctx, DeviceAuth auth, DeviceRegistry reg, AppDbContext db) => + { + if (!ctx.WebSockets.IsWebSocketRequest) + { + ctx.Response.StatusCode = StatusCodes.Status400BadRequest; + return; + } + + var device = await auth.AuthenticateAsync( + ctx.Request.Headers.Authorization.ToString(), ctx.RequestAborted); + if (device is null) + { + ctx.Response.StatusCode = StatusCodes.Status401Unauthorized; + return; + } + + using var ws = await ctx.WebSockets.AcceptWebSocketAsync(); + await using var active = new ActiveDevice(device.Id, ws); + if (!reg.TryRegister(active)) + { + await ws.CloseAsync(WebSocketCloseStatus.PolicyViolation, + "device already connected", ctx.RequestAborted); + return; + } + + try + { + await HandleAsync(active, device, db, ctx.RequestAborted); + } + finally + { + reg.Unregister(device.Id); + } + }); + return app; + } + + private static async Task HandleAsync(ActiveDevice active, Device device, AppDbContext db, CancellationToken ct) + { + var buffer = new byte[8192]; + var ws = active.Ws; + + while (ws.State == WebSocketState.Open && !ct.IsCancellationRequested) + { + using var msg = new MemoryStream(); + WebSocketReceiveResult res; + while (true) + { + res = await ws.ReceiveAsync(buffer, ct); + if (res.MessageType == WebSocketMessageType.Close) return; + msg.Write(buffer, 0, res.Count); + if (res.EndOfMessage) break; + } + + if (res.MessageType == WebSocketMessageType.Binary) + { + if (msg.Length != 3840) continue; + // Binary routing into a session lives in Task 18. + continue; + } + + var text = Encoding.UTF8.GetString(msg.ToArray()); + using var doc = JsonDocument.Parse(text); + var type = doc.RootElement.GetProperty("type").GetString(); + switch (type) + { + case "hello": + await SendHelloAckAsync(active, device, ct); + break; + + case "ping": + device.LastSeenAt = DateTimeOffset.UtcNow; + await db.SaveChangesAsync(ct); + await active.SendEnvelopeAsync(new PongEnvelope("pong"), ct); + break; + } + } + } + + private static async Task SendHelloAckAsync(ActiveDevice active, Device device, CancellationToken ct) + { + var enabled = device.Config?.EnabledToolsJson ?? "[]"; + var config = JsonNode.Parse($$""" + { + "voice": "{{Escape(device.Config?.Voice)}}", + "model": "{{Escape(device.Config?.Model)}}", + "system_prompt": "{{Escape(device.Config?.SystemPrompt)}}", + "idle_timeout_seconds": {{device.Config?.IdleTimeoutSeconds ?? 30}}, + "enabled_tools": {{enabled}} + } + """)!.AsObject(); + + await active.WriteEnvelopeAsync(new JsonObject + { + ["type"] = "hello_ack", + ["config"] = config, + }, ct); + } + + private static string Escape(string? s) => + (s ?? "").Replace("\\", "\\\\").Replace("\"", "\\\""); +} diff --git a/backend/Program.cs b/backend/Program.cs index 02fa7ca..08461a8 100644 --- a/backend/Program.cs +++ b/backend/Program.cs @@ -1,5 +1,6 @@ using backend.Auth; using backend.Data; +using backend.DeviceHub; using backend.Devices; using backend.Install; using backend.Pairing; @@ -41,6 +42,7 @@ builder.Services.AddSingleton( builder.Services.AddSingleton(); builder.Services.AddScoped(); builder.Services.AddSingleton(); +builder.Services.AddScoped(); var app = builder.Build(); @@ -69,6 +71,9 @@ app.MapInstallEndpoints(); app.MapGet("/health", () => Results.Ok(new { ok = true })); +app.UseWebSockets(); +app.MapDeviceHub(); + app.Run(); public partial class Program { }