natsdotnet/tests/NATS.Server.Tests/WebSocket/WsJwtAuthTests.cs

// Tests for WebSocket JWT authentication during upgrade (E11).
// Verifies JWT extraction from Authorization header, cookie, and query parameter.
// Reference: golang/nats-server/server/websocket.go — cookie JWT extraction (line 856),
//            websocket_test.go — TestWSReloadTLSConfig (line 4066).

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

namespace NATS.Server.Tests.WebSocket;

public class WsJwtAuthTests
{
    // ─── Authorization header JWT extraction ─────────────────────────────

    [Fact]
    public async Task Upgrade_AuthorizationBearerHeader_ExtractsJwt()
    {
        // JWT from Authorization: Bearer <token> header (standard HTTP auth)
        var jwt = "eyJhbGciOiJFZDI1NTE5IiwidHlwIjoiSldUIn0.test-payload.test-sig";
        var request = BuildValidRequest(extraHeaders:
            $"Authorization: Bearer {jwt}\r\n");
        var (inputStream, outputStream) = CreateStreamPair(request);

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

        result.Success.ShouldBeTrue();
        result.Jwt.ShouldBe(jwt);
    }

    [Fact]
    public async Task Upgrade_AuthorizationBearerCaseInsensitive()
    {
        // RFC 7235: "bearer" scheme is case-insensitive
        var jwt = "my-jwt-token-123";
        var request = BuildValidRequest(extraHeaders:
            $"Authorization: bearer {jwt}\r\n");
        var (inputStream, outputStream) = CreateStreamPair(request);

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

        result.Success.ShouldBeTrue();
        result.Jwt.ShouldBe(jwt);
    }

    [Fact]
    public async Task Upgrade_AuthorizationBareToken_ExtractsJwt()
    {
        // Some clients send the token directly without "Bearer" prefix
        var jwt = "raw-jwt-token-456";
        var request = BuildValidRequest(extraHeaders:
            $"Authorization: {jwt}\r\n");
        var (inputStream, outputStream) = CreateStreamPair(request);

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

        result.Success.ShouldBeTrue();
        result.Jwt.ShouldBe(jwt);
    }

    // ─── Cookie JWT extraction ──────────────────────────────────────────

    [Fact]
    public async Task Upgrade_JwtCookie_ExtractsJwt()
    {
        // Go parity: websocket.go line 856 — JWT from configured cookie name
        var jwt = "cookie-jwt-token-789";
        var request = BuildValidRequest(extraHeaders:
            $"Cookie: jwt={jwt}; other=value\r\n");
        var (inputStream, outputStream) = CreateStreamPair(request);

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

        result.Success.ShouldBeTrue();
        result.CookieJwt.ShouldBe(jwt);
        // Cookie JWT is used as fallback when no Authorization header is present
        result.Jwt.ShouldBe(jwt);
    }

    [Fact]
    public async Task Upgrade_AuthorizationHeader_TakesPriorityOverCookie()
    {
        // Authorization header has higher priority than cookie
        var headerJwt = "auth-header-jwt";
        var cookieJwt = "cookie-jwt";
        var request = BuildValidRequest(extraHeaders:
            $"Authorization: Bearer {headerJwt}\r\n" +
            $"Cookie: jwt={cookieJwt}\r\n");
        var (inputStream, outputStream) = CreateStreamPair(request);

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

        result.Success.ShouldBeTrue();
        result.Jwt.ShouldBe(headerJwt);
        result.CookieJwt.ShouldBe(cookieJwt); // Cookie value is still preserved
    }

    // ─── Query parameter JWT extraction ─────────────────────────────────

    [Fact]
    public async Task Upgrade_QueryParamJwt_ExtractsJwt()
    {
        // JWT from ?jwt= query parameter (useful for browser clients)
        var jwt = "query-jwt-token-abc";
        var request = BuildValidRequest(path: $"/?jwt={jwt}");
        var (inputStream, outputStream) = CreateStreamPair(request);

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

        result.Success.ShouldBeTrue();
        result.Jwt.ShouldBe(jwt);
    }

    [Fact]
    public async Task Upgrade_QueryParamJwt_UrlEncoded()
    {
        // JWT value may be URL-encoded
        var jwt = "eyJ0eXAiOiJKV1QifQ.payload.sig";
        var encoded = Uri.EscapeDataString(jwt);
        var request = BuildValidRequest(path: $"/?jwt={encoded}");
        var (inputStream, outputStream) = CreateStreamPair(request);

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

        result.Success.ShouldBeTrue();
        result.Jwt.ShouldBe(jwt);
    }

    [Fact]
    public async Task Upgrade_AuthorizationHeader_TakesPriorityOverQueryParam()
    {
        // Authorization header > query parameter
        var headerJwt = "auth-header-jwt";
        var queryJwt = "query-jwt";
        var request = BuildValidRequest(
            path: $"/?jwt={queryJwt}",
            extraHeaders: $"Authorization: Bearer {headerJwt}\r\n");
        var (inputStream, outputStream) = CreateStreamPair(request);

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

        result.Success.ShouldBeTrue();
        result.Jwt.ShouldBe(headerJwt);
    }

    [Fact]
    public async Task Upgrade_Cookie_TakesPriorityOverQueryParam()
    {
        // Cookie > query parameter
        var cookieJwt = "cookie-jwt";
        var queryJwt = "query-jwt";
        var request = BuildValidRequest(
            path: $"/?jwt={queryJwt}",
            extraHeaders: $"Cookie: jwt_token={cookieJwt}\r\n");
        var (inputStream, outputStream) = CreateStreamPair(request);

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

        result.Success.ShouldBeTrue();
        result.Jwt.ShouldBe(cookieJwt);
    }

    // ─── No JWT scenarios ───────────────────────────────────────────────

    [Fact]
    public async Task Upgrade_NoJwtAnywhere_JwtIsNull()
    {
        // No JWT in any source
        var request = BuildValidRequest();
        var (inputStream, outputStream) = CreateStreamPair(request);

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

        result.Success.ShouldBeTrue();
        result.Jwt.ShouldBeNull();
    }

    [Fact]
    public async Task Upgrade_EmptyAuthorizationHeader_JwtIsEmpty()
    {
        // Empty authorization header should produce empty string (non-null)
        var request = BuildValidRequest(extraHeaders: "Authorization: \r\n");
        var (inputStream, outputStream) = CreateStreamPair(request);

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

        result.Success.ShouldBeTrue();
        // Empty auth header is treated as null/no JWT
        result.Jwt.ShouldBeNull();
    }

    // ─── ExtractBearerToken unit tests ──────────────────────────────────

    [Fact]
    public void ExtractBearerToken_BearerPrefix()
    {
        WsUpgrade.ExtractBearerToken("Bearer my-token").ShouldBe("my-token");
    }

    [Fact]
    public void ExtractBearerToken_BearerPrefixLowerCase()
    {
        WsUpgrade.ExtractBearerToken("bearer my-token").ShouldBe("my-token");
    }

    [Fact]
    public void ExtractBearerToken_BareToken()
    {
        WsUpgrade.ExtractBearerToken("raw-token").ShouldBe("raw-token");
    }

    [Fact]
    public void ExtractBearerToken_Null()
    {
        WsUpgrade.ExtractBearerToken(null).ShouldBeNull();
    }

    [Fact]
    public void ExtractBearerToken_Empty()
    {
        WsUpgrade.ExtractBearerToken("").ShouldBeNull();
    }

    [Fact]
    public void ExtractBearerToken_Whitespace()
    {
        WsUpgrade.ExtractBearerToken("   ").ShouldBeNull();
    }

    // ─── ParseQueryString unit tests ────────────────────────────────────

    [Fact]
    public void ParseQueryString_SingleParam()
    {
        var result = WsUpgrade.ParseQueryString("?jwt=token123");
        result["jwt"].ShouldBe("token123");
    }

    [Fact]
    public void ParseQueryString_MultipleParams()
    {
        var result = WsUpgrade.ParseQueryString("?jwt=token&user=admin");
        result["jwt"].ShouldBe("token");
        result["user"].ShouldBe("admin");
    }

    [Fact]
    public void ParseQueryString_UrlEncoded()
    {
        var result = WsUpgrade.ParseQueryString("?jwt=a%20b%3Dc");
        result["jwt"].ShouldBe("a b=c");
    }

    [Fact]
    public void ParseQueryString_NoQuestionMark()
    {
        var result = WsUpgrade.ParseQueryString("jwt=token");
        result["jwt"].ShouldBe("token");
    }

    // ─── FailUnauthorizedAsync ──────────────────────────────────────────

    [Fact]
    public async Task FailUnauthorizedAsync_Returns401()
    {
        var output = new MemoryStream();
        var result = await WsUpgrade.FailUnauthorizedAsync(output, "invalid JWT");

        result.Success.ShouldBeFalse();
        output.Position = 0;
        var response = Encoding.ASCII.GetString(output.ToArray());
        response.ShouldContain("401");
        response.ShouldContain("invalid JWT");
    }

    // ─── Query param path routing still works with query strings ────────

    [Fact]
    public async Task Upgrade_PathWithQueryParam_StillRoutesCorrectly()
    {
        // /leafnode?jwt=token should still detect as leaf kind
        var request = BuildValidRequest(path: "/leafnode?jwt=my-token");
        var (inputStream, outputStream) = CreateStreamPair(request);

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

        result.Success.ShouldBeTrue();
        result.Kind.ShouldBe(WsClientKind.Leaf);
        result.Jwt.ShouldBe("my-token");
    }

    // ─── 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());
    }
}