WebDrop/server/tests/WebDrop.Api.Tests/StreamingTests.cs

using System.Net;
using System.Net.WebSockets;
using System.Text;
using System.Text.Json;
using WebDrop.Core;
using static WebDrop.Api.Tests.ServerFixture;

namespace WebDrop.Api.Tests;

public class StreamingTests(ServerFixture s) : IClassFixture<ServerFixture>
{
    private static readonly TimeSpan Timeout = TimeSpan.FromSeconds(15);

    private static byte[] Small(string item, byte fill = 0x42) => Frame(item, 0, FinalEnc, EncPayload, fill);

    private async Task<ClientWebSocket> ConnectWs(Chan ch, Who who, long from = 0)
    {
        var ws = new ClientWebSocket();
        ws.Options.AddSubProtocol("wd.v1");
        ws.Options.AddSubProtocol($"wd.role.{who.Role}");
        if (who.Token is not null) ws.Options.AddSubProtocol($"wd.auth.{B64(who.Token)}");
        ws.Options.AddSubProtocol($"wd.ep.{B64(who.EndpointSecret)}");

        var uri = new UriBuilder(s.BaseUri) { Scheme = "ws", Path = $"/api/v1/channels/{ch.Id}/ws", Query = $"from={from}" }.Uri;
        using var cts = new CancellationTokenSource(Timeout);
        await ws.ConnectAsync(uri, cts.Token);
        return ws;
    }

    private static async Task<(WebSocketMessageType Type, byte[] Data)> Receive(ClientWebSocket ws)
    {
        using var cts = new CancellationTokenSource(Timeout);
        var buffer = new byte[2 * 1024 * 1024];
        int total = 0;
        while (true)
        {
            var r = await ws.ReceiveAsync(buffer.AsMemory(total), cts.Token);
            if (r.MessageType == WebSocketMessageType.Close) return (WebSocketMessageType.Close, []);
            total += r.Count;
            if (r.EndOfMessage) return (r.MessageType, buffer[..total]);
        }
    }

    [Fact]
    public async Task WebSocketReplaysHistoryThenPushesNewFrames()
    {
        var ch = await s.CreateAsync();
        string a = NewId(), b = NewId();
        await s.AppendAsync(ch, ch.Owner, Small(a, 1));

        using var ws = await ConnectWs(ch, ch.Owner);
        Assert.Equal("wd.v1", ws.SubProtocol);

        var (type, data) = await Receive(ws);
        Assert.Equal(WebSocketMessageType.Binary, type);
        var first = ParseOne(data);
        Assert.Equal(1, first.Seq);
        Assert.Equal(Small(a, 1), first.Body);

        // 连接建立之后追加的帧会被实时推送
        await s.AppendAsync(ch, ch.Owner, Small(b, 2));
        var (_, pushed) = await Receive(ws);
        Assert.Equal(2, ParseOne(pushed).Seq);
        Assert.Equal(Small(b, 2), ParseOne(pushed).Body);
    }

    [Fact]
    public async Task WebSocketResumesFromTheRequestedSeq()
    {
        var ch = await s.CreateAsync();
        for (int i = 0; i < 4; i++) await s.AppendAsync(ch, ch.Owner, Small(NewId(), (byte)i));

        using var ws = await ConnectWs(ch, ch.Owner, from: 3);
        Assert.Equal(3, ParseOne((await Receive(ws)).Data).Seq);
        Assert.Equal(4, ParseOne((await Receive(ws)).Data).Seq);
    }

    [Fact]
    public async Task WebSocketDeliversUploadsFromAnotherMember()
    {
        var ch = await s.CreateAsync();
        await s.JoinAsync(ch, ch.Member);

        using var owner = await ConnectWs(ch, ch.Owner);
        using var member = await ConnectWs(ch, ch.Member);

        string item = NewId();
        using var put = await s.AppendAsync(ch, ch.Member, Small(item, 9));
        Assert.Equal(HttpStatusCode.Created, put.StatusCode);

        foreach (var ws in new[] { owner, member })
        {
            var (_, data) = await Receive(ws);
            Assert.Equal(Small(item, 9), ParseOne(data).Body);
        }
    }

    [Fact]
    public async Task WebSocketUpgradeIsRefusedWithoutValidCredentials()
    {
        var ch = await s.CreateAsync();

        // 令牌错误:在升级之前就以 HTTP 401 拒绝,而不是先建立连接再断开
        var bad = new Who("owner", new byte[32], ch.OwnerSecret);
        var ex = await Assert.ThrowsAsync<WebSocketException>(() => ConnectWs(ch, bad));
        Assert.Contains("401", ex.Message);

        // 未入席
        var stranger = new Who("write", ch.WriteToken, new byte[32]);
        var ex2 = await Assert.ThrowsAsync<WebSocketException>(() => ConnectWs(ch, stranger));
        Assert.Contains("403", ex2.Message);
    }

    [Fact]
    public async Task WebSocketNotifiesSeatChangesAndClosesWhenRevoked()
    {
        var ch = await s.CreateAsync();
        var joined = await s.JoinAsync(ch, ch.Member);
        string memberHash = (await ReadJsonAsync(joined)).GetProperty("endpoint_hash").GetString()!;

        using var member = await ConnectWs(ch, ch.Member);

        // 所有者释放席位:订阅者先收到 seats 控制消息,随后被以 not_joined 关闭
        using var release = await s.Http.SendAsync(s.Req(HttpMethod.Delete, $"/api/v1/channels/{ch.Id}/seats/{memberHash}", ch.Owner));
        Assert.Equal(HttpStatusCode.NoContent, release.StatusCode);

        bool closed = false;
        for (int i = 0; i < 5 && !closed; i++)
        {
            var (type, data) = await Receive(member);
            if (type == WebSocketMessageType.Close)
            {
                closed = true;
                Assert.Equal(WebSocketCloseStatus.PolicyViolation, member.CloseStatus);
                Assert.Equal("not_joined", member.CloseStatusDescription);
            }
            else if (type == WebSocketMessageType.Text)
            {
                Assert.Equal("seats", JsonDocument.Parse(data).RootElement.GetProperty("t").GetString());
            }
        }
        Assert.True(closed);
    }

    [Fact]
    public async Task WebSocketEndsWhenTheChannelIsDeleted()
    {
        var ch = await s.CreateAsync();
        using var ws = await ConnectWs(ch, ch.Owner);

        using var del = await s.Http.SendAsync(s.Req(HttpMethod.Delete, $"/api/v1/channels/{ch.Id}", ch.Owner));
        Assert.Equal(HttpStatusCode.NoContent, del.StatusCode);

        var (type, _) = await Receive(ws);
        Assert.Equal(WebSocketMessageType.Close, type);
        Assert.Equal("not_found", ws.CloseStatusDescription);
    }

    [Fact]
    public async Task WebSocketClientCanCloseCleanly()
    {
        var ch = await s.CreateAsync();
        using var ws = await ConnectWs(ch, ch.Owner);
        using var cts = new CancellationTokenSource(Timeout);
        await ws.CloseAsync(WebSocketCloseStatus.NormalClosure, null, cts.Token);
        Assert.Equal(WebSocketState.Closed, ws.State);
    }

    // ───── SSE ─────

    private sealed record SseEvent(string Event, string? Id, string Data);

    private async IAsyncEnumerable<SseEvent> ReadSse(string path, Who who, Dictionary<string, string>? extraHeaders = null)
    {
        using var req = s.Req(HttpMethod.Get, path, who);
        foreach (var (k, v) in extraHeaders ?? []) req.Headers.Add(k, v);
        using var cts = new CancellationTokenSource(Timeout);
        using var resp = await s.Http.SendAsync(req, HttpCompletionOption.ResponseHeadersRead, cts.Token);
        Assert.Equal(HttpStatusCode.OK, resp.StatusCode);
        Assert.Equal("text/event-stream", resp.Content.Headers.ContentType!.MediaType);

        await using var stream = await resp.Content.ReadAsStreamAsync(cts.Token);
        using var reader = new StreamReader(stream, Encoding.UTF8);

        string ev = "message", data = "";
        string? id = null;
        while (await reader.ReadLineAsync(cts.Token) is { } line)
        {
            if (line.Length == 0)
            {
                if (data.Length > 0) yield return new SseEvent(ev, id, data);
                ev = "message"; data = ""; id = null;
            }
            else if (line.StartsWith("event: ")) ev = line[7..];
            else if (line.StartsWith("id: ")) id = line[4..];
            else if (line.StartsWith("data: ")) data = line[6..];
        }
    }

    private static byte[] DecodeB64(string text) => System.Buffers.Text.Base64Url.DecodeFromChars(text);

    [Fact]
    public async Task SseStreamsHistoryThenNewFrames()
    {
        var ch = await s.CreateAsync();
        string a = NewId(), b = NewId();
        await s.AppendAsync(ch, ch.Owner, Small(a, 1));

        await using var events = ReadSse($"/api/v1/channels/{ch.Id}/events", ch.Owner).GetAsyncEnumerator();

        Assert.True(await events.MoveNextAsync());
        Assert.Equal("log", events.Current.Event);
        Assert.Equal("1", events.Current.Id);
        Assert.Equal(Small(a, 1), ParseOne(DecodeB64(events.Current.Data)).Body);

        await s.AppendAsync(ch, ch.Owner, Small(b, 2));
        Assert.True(await events.MoveNextAsync());
        Assert.Equal("2", events.Current.Id);
        Assert.Equal(Small(b, 2), ParseOne(DecodeB64(events.Current.Data)).Body);
    }

    [Fact]
    public async Task SseResumesFromLastEventId()
    {
        var ch = await s.CreateAsync();
        for (int i = 0; i < 3; i++) await s.AppendAsync(ch, ch.Owner, Small(NewId(), (byte)i));

        await using var events = ReadSse($"/api/v1/channels/{ch.Id}/events", ch.Owner,
            new() { ["Last-Event-ID"] = "2" }).GetAsyncEnumerator();

        Assert.True(await events.MoveNextAsync());
        Assert.Equal("3", events.Current.Id); // 断线前处理到 2,续传从 3 开始
    }

    [Fact]
    public async Task SseRefusesBadCredentialsWithAJsonError()
    {
        var ch = await s.CreateAsync();
        var bad = new Who("owner", new byte[32], ch.OwnerSecret);
        await AssertErrorAsync(await s.Http.SendAsync(s.Req(HttpMethod.Get, $"/api/v1/channels/{ch.Id}/events", bad)), 401, "unauthorized");
    }
}