pgsql-jellyfin/Jellyfin.Api/Middleware/WebSocketAuthenticationMiddleware.cs

// <copyright file="WebSocketAuthenticationMiddleware.cs" company="PlaceholderCompany">
// Copyright (c) PlaceholderCompany. All rights reserved.
// </copyright>

namespace Jellyfin.Api.Middleware;

using System;
using System.Globalization;
using System.Security.Claims;
using System.Threading.Tasks;
using Jellyfin.Api.Constants;
using Jellyfin.Data;
using Jellyfin.Database.Implementations.Enums;
using MediaBrowser.Controller.Net;
using Microsoft.AspNetCore.Http;
using Microsoft.Extensions.Logging;

/// <summary>
/// WebSocket Authentication Middleware.
/// Extracts and validates API tokens from WebSocket query strings.
/// </summary>
public class WebSocketAuthenticationMiddleware
{
    private readonly RequestDelegate _next;
    private readonly ILogger<WebSocketAuthenticationMiddleware> _logger;

    /// <summary>
    /// Initializes a new instance of the <see cref="WebSocketAuthenticationMiddleware"/> class.
    /// </summary>
    /// <param name="next">The next delegate in the pipeline.</param>
    /// <param name="logger">The logger.</param>
    public WebSocketAuthenticationMiddleware(RequestDelegate next, ILogger<WebSocketAuthenticationMiddleware> logger)
    {
        _next = next;
        _logger = logger;
    }

    /// <summary>
    /// Executes the middleware action.
    /// </summary>
    /// <param name="httpContext">The current HTTP context.</param>
    /// <param name="authService">The authentication service.</param>
    /// <returns>The async task.</returns>
    public async Task Invoke(HttpContext httpContext, IAuthService authService)
    {
        // Only process WebSocket upgrade requests on the /socket endpoint
        if (httpContext.WebSockets.IsWebSocketRequestUpgrade
            && httpContext.Request.Path.StartsWithSegments("/socket", StringComparison.OrdinalIgnoreCase))
        {
            await AuthenticateWebSocketRequest(httpContext, authService).ConfigureAwait(false);
        }

        await _next(httpContext).ConfigureAwait(false);
    }

    /// <summary>
    /// Authenticates WebSocket requests by extracting api_key from query string.
    /// </summary>
    /// <param name="httpContext">The HTTP context.</param>
    /// <param name="authService">The authentication service.</param>
    /// <returns>The async task.</returns>
    private async Task AuthenticateWebSocketRequest(HttpContext httpContext, IAuthService authService)
    {
        try
        {
            // Extract api_key from query string
            if (!httpContext.Request.Query.TryGetValue("api_key", out var apiKeyValues) || apiKeyValues.Count == 0)
            {
                _logger.LogDebug("WebSocket connection attempted without api_key query parameter. IP: {IP}", 
                    httpContext.Connection.RemoteIpAddress);
                return;
            }

            var apiKey = apiKeyValues[0];
            if (string.IsNullOrWhiteSpace(apiKey))
            {
                _logger.LogDebug("WebSocket connection attempted with empty api_key. IP: {IP}", 
                    httpContext.Connection.RemoteIpAddress);
                return;
            }

            // Authenticate using the provided api_key
            var authorizationInfo = await authService.Authenticate(httpContext.Request).ConfigureAwait(false);

            if (!authorizationInfo.HasToken)
            {
                _logger.LogDebug("WebSocket authentication failed: Invalid or expired token. IP: {IP}", 
                    httpContext.Connection.RemoteIpAddress);
                return;
            }

            // Set the authenticated user in the context for downstream middleware
            var role = UserRoles.User;
            if (authorizationInfo.IsApiKey
                || (authorizationInfo.User?.HasPermission(PermissionKind.IsAdministrator) ?? false))
            {
                role = UserRoles.Administrator;
            }

            var claims = new[]
            {
                new Claim(ClaimTypes.Name, authorizationInfo.User?.Username ?? string.Empty),
                new Claim(ClaimTypes.Role, role),
                new Claim(InternalClaimTypes.UserId, authorizationInfo.UserId.ToString("N", CultureInfo.InvariantCulture)),
                new Claim(InternalClaimTypes.DeviceId, authorizationInfo.DeviceId ?? string.Empty),
                new Claim(InternalClaimTypes.Device, authorizationInfo.Device ?? string.Empty),
                new Claim(InternalClaimTypes.Client, authorizationInfo.Client ?? string.Empty),
                new Claim(InternalClaimTypes.Version, authorizationInfo.Version ?? string.Empty),
                new Claim(InternalClaimTypes.Token, authorizationInfo.Token),
                new Claim(InternalClaimTypes.IsApiKey, authorizationInfo.IsApiKey.ToString(CultureInfo.InvariantCulture))
            };

            var identity = new ClaimsIdentity(claims, "WebSocket");
            var principal = new ClaimsPrincipal(identity);
            httpContext.User = principal;

            _logger.LogDebug("WebSocket authentication successful for user {Username}. IP: {IP}",
                authorizationInfo.User?.Username ?? "Unknown",
                httpContext.Connection.RemoteIpAddress);
        }
        catch (Exception ex)
        {
            _logger.LogDebug(ex, "Error during WebSocket authentication. IP: {IP}",
                httpContext.Connection.RemoteIpAddress);
        }
    }
}
        /// <summary>
        /// Adds WebSocket authentication to the application pipeline.
        /// </summary>
        /// <param name="appBuilder">The application builder.</param>
        /// <returns>The updated application builder.</returns>
        public static IApplicationBuilder UseWebSocketAuthentication(this IApplicationBuilder appBuilder)
        {
            return appBuilder.UseMiddleware<WebSocketAuthenticationMiddleware>();
        }