natsdotnet/tests/NATS.Server.Transport.Tests/WebSocket/WsCompressionNegotiationTests.cs

// Tests for WebSocket permessage-deflate parameter negotiation (E10).
// Verifies RFC 7692 extension parameter parsing and negotiation during
// WebSocket upgrade handshake.
// Reference: golang/nats-server/server/websocket.go — wsPMCExtensionSupport (line 885).

using System.Text;
using NATS.Server.WebSocket;

namespace NATS.Server.Transport.Tests.WebSocket;

public class WsCompressionNegotiationTests
{
    // ─── WsDeflateNegotiator.Negotiate tests ──────────────────────────────

    [Fact]
    public void Negotiate_NullHeader_ReturnsNull()
    {
        // Go parity: wsPMCExtensionSupport — no extension header means no compression
        var result = WsDeflateNegotiator.Negotiate(null);
        result.ShouldBeNull();
    }

    [Fact]
    public void Negotiate_EmptyHeader_ReturnsNull()
    {
        var result = WsDeflateNegotiator.Negotiate("");
        result.ShouldBeNull();
    }

    [Fact]
    public void Negotiate_NoPermessageDeflate_ReturnsNull()
    {
        var result = WsDeflateNegotiator.Negotiate("x-webkit-deflate-frame");
        result.ShouldBeNull();
    }

    [Fact]
    public void Negotiate_BarePermessageDeflate_ReturnsDefaults()
    {
        // Go parity: wsPMCExtensionSupport — basic extension without parameters
        var result = WsDeflateNegotiator.Negotiate("permessage-deflate");

        result.ShouldNotBeNull();
        // NATS always enforces no_context_takeover
        result.Value.ServerNoContextTakeover.ShouldBeTrue();
        result.Value.ClientNoContextTakeover.ShouldBeTrue();
        result.Value.ServerMaxWindowBits.ShouldBe(15);
        result.Value.ClientMaxWindowBits.ShouldBe(15);
    }

    [Fact]
    public void Negotiate_WithServerNoContextTakeover()
    {
        // Go parity: wsPMCExtensionSupport — server_no_context_takeover parameter
        var result = WsDeflateNegotiator.Negotiate("permessage-deflate; server_no_context_takeover");

        result.ShouldNotBeNull();
        result.Value.ServerNoContextTakeover.ShouldBeTrue();
    }

    [Fact]
    public void Negotiate_WithClientNoContextTakeover()
    {
        // Go parity: wsPMCExtensionSupport — client_no_context_takeover parameter
        var result = WsDeflateNegotiator.Negotiate("permessage-deflate; client_no_context_takeover");

        result.ShouldNotBeNull();
        result.Value.ClientNoContextTakeover.ShouldBeTrue();
    }

    [Fact]
    public void Negotiate_WithBothNoContextTakeover()
    {
        // Go parity: wsPMCExtensionSupport — both no_context_takeover parameters
        var result = WsDeflateNegotiator.Negotiate(
            "permessage-deflate; server_no_context_takeover; client_no_context_takeover");

        result.ShouldNotBeNull();
        result.Value.ServerNoContextTakeover.ShouldBeTrue();
        result.Value.ClientNoContextTakeover.ShouldBeTrue();
    }

    [Fact]
    public void Negotiate_WithServerMaxWindowBits()
    {
        // RFC 7692 Section 7.1.2.1: server_max_window_bits parameter
        var result = WsDeflateNegotiator.Negotiate("permessage-deflate; server_max_window_bits=10");

        result.ShouldNotBeNull();
        result.Value.ServerMaxWindowBits.ShouldBe(10);
    }

    [Fact]
    public void Negotiate_WithClientMaxWindowBits_Value()
    {
        // RFC 7692 Section 7.1.2.2: client_max_window_bits with explicit value
        var result = WsDeflateNegotiator.Negotiate("permessage-deflate; client_max_window_bits=12");

        result.ShouldNotBeNull();
        result.Value.ClientMaxWindowBits.ShouldBe(12);
    }

    [Fact]
    public void Negotiate_WithClientMaxWindowBits_NoValue()
    {
        // RFC 7692 Section 7.1.2.2: client_max_window_bits with no value means
        // client supports any value 8-15; defaults to 15
        var result = WsDeflateNegotiator.Negotiate("permessage-deflate; client_max_window_bits");

        result.ShouldNotBeNull();
        result.Value.ClientMaxWindowBits.ShouldBe(15);
    }

    [Fact]
    public void Negotiate_WindowBits_ClampedToValidRange()
    {
        // RFC 7692: valid range is 8-15
        var result = WsDeflateNegotiator.Negotiate(
            "permessage-deflate; server_max_window_bits=5; client_max_window_bits=20");

        result.ShouldNotBeNull();
        result.Value.ServerMaxWindowBits.ShouldBe(8);  // Clamped up from 5
        result.Value.ClientMaxWindowBits.ShouldBe(15); // Clamped down from 20
    }

    [Fact]
    public void Negotiate_FullParameters()
    {
        // All parameters specified
        var result = WsDeflateNegotiator.Negotiate(
            "permessage-deflate; server_no_context_takeover; client_no_context_takeover; server_max_window_bits=9; client_max_window_bits=11");

        result.ShouldNotBeNull();
        result.Value.ServerNoContextTakeover.ShouldBeTrue();
        result.Value.ClientNoContextTakeover.ShouldBeTrue();
        result.Value.ServerMaxWindowBits.ShouldBe(9);
        result.Value.ClientMaxWindowBits.ShouldBe(11);
    }

    [Fact]
    public void Negotiate_CaseInsensitive()
    {
        // RFC 7692 extension names are case-insensitive
        var result = WsDeflateNegotiator.Negotiate("Permessage-Deflate; Server_No_Context_Takeover");

        result.ShouldNotBeNull();
        result.Value.ServerNoContextTakeover.ShouldBeTrue();
    }

    [Fact]
    public void Negotiate_MultipleExtensions_PicksDeflate()
    {
        // Header may contain multiple comma-separated extensions
        var result = WsDeflateNegotiator.Negotiate(
            "x-custom-ext, permessage-deflate; server_no_context_takeover, other-ext");

        result.ShouldNotBeNull();
        result.Value.ServerNoContextTakeover.ShouldBeTrue();
    }

    [Fact]
    public void Negotiate_WhitespaceHandling()
    {
        // Extra whitespace around parameters
        var result = WsDeflateNegotiator.Negotiate(
            "  permessage-deflate ;  server_no_context_takeover ;  client_max_window_bits = 10  ");

        result.ShouldNotBeNull();
        result.Value.ServerNoContextTakeover.ShouldBeTrue();
        result.Value.ClientMaxWindowBits.ShouldBe(10);
    }

    // ─── NatsAlwaysEnforcesNoContextTakeover ─────────────────────────────

    [Fact]
    public void Negotiate_AlwaysEnforcesNoContextTakeover()
    {
        // NATS Go server always returns server_no_context_takeover and
        // client_no_context_takeover regardless of what the client requests
        var result = WsDeflateNegotiator.Negotiate("permessage-deflate");

        result.ShouldNotBeNull();
        result.Value.ServerNoContextTakeover.ShouldBeTrue();
        result.Value.ClientNoContextTakeover.ShouldBeTrue();
    }

    // ─── WsDeflateParams.ToResponseHeaderValue tests ────────────────────

    [Fact]
    public void DefaultParams_ResponseHeader_ContainsNoContextTakeover()
    {
        var header = WsDeflateParams.Default.ToResponseHeaderValue();

        header.ShouldContain("permessage-deflate");
        header.ShouldContain("server_no_context_takeover");
        header.ShouldContain("client_no_context_takeover");
        header.ShouldNotContain("server_max_window_bits");
        header.ShouldNotContain("client_max_window_bits");
    }

    [Fact]
    public void CustomWindowBits_ResponseHeader_IncludesValues()
    {
        var params_ = new WsDeflateParams(
            ServerNoContextTakeover: true,
            ClientNoContextTakeover: true,
            ServerMaxWindowBits: 10,
            ClientMaxWindowBits: 12);

        var header = params_.ToResponseHeaderValue();

        header.ShouldContain("server_max_window_bits=10");
        header.ShouldContain("client_max_window_bits=12");
    }

    [Fact]
    public void DefaultWindowBits_ResponseHeader_OmitsValues()
    {
        // RFC 7692: window bits of 15 is the default and should not be sent
        var params_ = new WsDeflateParams(
            ServerNoContextTakeover: true,
            ClientNoContextTakeover: true,
            ServerMaxWindowBits: 15,
            ClientMaxWindowBits: 15);

        var header = params_.ToResponseHeaderValue();

        header.ShouldNotContain("server_max_window_bits");
        header.ShouldNotContain("client_max_window_bits");
    }

    // ─── Integration with WsUpgrade ─────────────────────────────────────

    [Fact]
    public async Task Upgrade_WithDeflateParams_NegotiatesCompression()
    {
        // Go parity: WebSocket upgrade with permessage-deflate parameters
        var request = BuildValidRequest(extraHeaders:
            "Sec-WebSocket-Extensions: permessage-deflate; server_no_context_takeover; client_no_context_takeover; server_max_window_bits=10\r\n");
        var (inputStream, outputStream) = CreateStreamPair(request);

        var opts = new WebSocketOptions { NoTls = true, Compression = true };
        var result = await WsUpgrade.TryUpgradeAsync(inputStream, outputStream, opts);

        result.Success.ShouldBeTrue();
        result.Compress.ShouldBeTrue();
        result.DeflateParams.ShouldNotBeNull();
        result.DeflateParams.Value.ServerNoContextTakeover.ShouldBeTrue();
        result.DeflateParams.Value.ClientNoContextTakeover.ShouldBeTrue();
        result.DeflateParams.Value.ServerMaxWindowBits.ShouldBe(10);
    }

    [Fact]
    public async Task Upgrade_WithDeflateParams_ResponseIncludesNegotiatedParams()
    {
        var request = BuildValidRequest(extraHeaders:
            "Sec-WebSocket-Extensions: permessage-deflate; server_no_context_takeover; client_max_window_bits=10\r\n");
        var (inputStream, outputStream) = CreateStreamPair(request);

        var opts = new WebSocketOptions { NoTls = true, Compression = true };
        await WsUpgrade.TryUpgradeAsync(inputStream, outputStream, opts);

        var response = ReadResponse(outputStream);
        response.ShouldContain("permessage-deflate");
        response.ShouldContain("server_no_context_takeover");
        response.ShouldContain("client_no_context_takeover");
        response.ShouldContain("client_max_window_bits=10");
    }

    [Fact]
    public async Task Upgrade_CompressionDisabled_NoDeflateParams()
    {
        var request = BuildValidRequest(extraHeaders:
            "Sec-WebSocket-Extensions: permessage-deflate; server_no_context_takeover\r\n");
        var (inputStream, outputStream) = CreateStreamPair(request);

        var opts = new WebSocketOptions { NoTls = true, Compression = false };
        var result = await WsUpgrade.TryUpgradeAsync(inputStream, outputStream, opts);

        result.Success.ShouldBeTrue();
        result.Compress.ShouldBeFalse();
        result.DeflateParams.ShouldBeNull();
    }

    [Fact]
    public async Task Upgrade_NoExtensionHeader_NoCompression()
    {
        var request = BuildValidRequest();
        var (inputStream, outputStream) = CreateStreamPair(request);

        var opts = new WebSocketOptions { NoTls = true, Compression = true };
        var result = await WsUpgrade.TryUpgradeAsync(inputStream, outputStream, opts);

        result.Success.ShouldBeTrue();
        result.Compress.ShouldBeFalse();
        result.DeflateParams.ShouldBeNull();
    }

    // ─── Helpers ─────────────────────────────────────────────────────────

    private static string BuildValidRequest(string path = "/", string? extraHeaders = null)
    {
        var sb = new StringBuilder();
        sb.Append($"GET {path} HTTP/1.1\r\n");
        sb.Append("Host: localhost:4222\r\n");
        sb.Append("Upgrade: websocket\r\n");
        sb.Append("Connection: Upgrade\r\n");
        sb.Append("Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==\r\n");
        sb.Append("Sec-WebSocket-Version: 13\r\n");
        if (extraHeaders != null)
            sb.Append(extraHeaders);
        sb.Append("\r\n");
        return sb.ToString();
    }

    private static (Stream input, MemoryStream output) CreateStreamPair(string httpRequest)
    {
        var inputBytes = Encoding.ASCII.GetBytes(httpRequest);
        return (new MemoryStream(inputBytes), new MemoryStream());
    }

    private static string ReadResponse(MemoryStream output)
    {
        output.Position = 0;
        return Encoding.ASCII.GetString(output.ToArray());
    }
}