1 files changed, 138 insertions, 10 deletions
diff --git a/Jellyfin.Server/Filters/AdditionalModelFilter.cs b/Jellyfin.Server/Filters/AdditionalModelFilter.cs
index 645696e31..bf38f741c 100644
--- a/Jellyfin.Server/Filters/AdditionalModelFilter.cs
+++ b/Jellyfin.Server/Filters/AdditionalModelFilter.cs
@@ -1,12 +1,16 @@
 using System;
+using System.Collections.Generic;
+using System.ComponentModel;
 using System.Linq;
+using System.Reflection;
 using Jellyfin.Extensions;
 using Jellyfin.Server.Migrations;
 using MediaBrowser.Common.Plugins;
 using MediaBrowser.Controller.Configuration;
-using MediaBrowser.Controller.LiveTv;
+using MediaBrowser.Controller.Net;
+using MediaBrowser.Controller.Net.WebSocketMessages;
+using MediaBrowser.Controller.Net.WebSocketMessages.Outbound;
 using MediaBrowser.Model.ApiClient;
-using MediaBrowser.Model.Entities;
 using MediaBrowser.Model.Session;
 using MediaBrowser.Model.SyncPlay;
 using Microsoft.OpenApi.Any;
@@ -36,17 +40,141 @@ namespace Jellyfin.Server.Filters
         /// <inheritdoc />
         public void Apply(OpenApiDocument swaggerDoc, DocumentFilterContext context)
         {
-            context.SchemaGenerator.GenerateSchema(typeof(LibraryUpdateInfo), context.SchemaRepository);
             context.SchemaGenerator.GenerateSchema(typeof(IPlugin), context.SchemaRepository);
-            context.SchemaGenerator.GenerateSchema(typeof(PlayRequest), context.SchemaRepository);
-            context.SchemaGenerator.GenerateSchema(typeof(PlaystateRequest), context.SchemaRepository);
-            context.SchemaGenerator.GenerateSchema(typeof(TimerEventInfo), context.SchemaRepository);
-            context.SchemaGenerator.GenerateSchema(typeof(SendCommand), context.SchemaRepository);
-            context.SchemaGenerator.GenerateSchema(typeof(GeneralCommandType), context.SchemaRepository);
 
-            context.SchemaGenerator.GenerateSchema(typeof(GroupUpdate<object>), context.SchemaRepository);
+            var webSocketTypes = typeof(WebSocketMessage).Assembly.GetTypes()
+                .Where(t => t.IsSubclassOf(typeof(WebSocketMessage))
+                            && !t.IsGenericType
+                            && t != typeof(WebSocketMessageInfo))
+                .ToList();
+
+            var inboundWebSocketSchemas = new List<OpenApiSchema>();
+            var inboundWebSocketDiscriminators = new Dictionary<string, string>();
+            foreach (var type in webSocketTypes.Where(t => typeof(IInboundWebSocketMessage).IsAssignableFrom(t)))
+            {
+                var messageType = (SessionMessageType?)type.GetProperty(nameof(WebSocketMessage.MessageType))?.GetCustomAttribute<DefaultValueAttribute>()?.Value;
+                if (messageType is null)
+                {
+                    continue;
+                }
+
+                var schema = context.SchemaGenerator.GenerateSchema(type, context.SchemaRepository);
+                inboundWebSocketSchemas.Add(schema);
+                inboundWebSocketDiscriminators[messageType.ToString()!] = schema.Reference.ReferenceV3;
+            }
+
+            var inboundWebSocketMessageSchema = new OpenApiSchema
+            {
+                Type = "object",
+                Description = "Represents the list of possible inbound websocket types",
+                Reference = new OpenApiReference
+                {
+                    Id = nameof(InboundWebSocketMessage),
+                    Type = ReferenceType.Schema
+                },
+                OneOf = inboundWebSocketSchemas,
+                Discriminator = new OpenApiDiscriminator
+                {
+                    PropertyName = nameof(WebSocketMessage.MessageType),
+                    Mapping = inboundWebSocketDiscriminators
+                }
+            };
+
+            context.SchemaRepository.AddDefinition(nameof(InboundWebSocketMessage), inboundWebSocketMessageSchema);
+
+            var outboundWebSocketSchemas = new List<OpenApiSchema>();
+            var outboundWebSocketDiscriminators = new Dictionary<string, string>();
+            foreach (var type in webSocketTypes.Where(t => typeof(IOutboundWebSocketMessage).IsAssignableFrom(t)))
+            {
+                var messageType = (SessionMessageType?)type.GetProperty(nameof(WebSocketMessage.MessageType))?.GetCustomAttribute<DefaultValueAttribute>()?.Value;
+                if (messageType is null)
+                {
+                    continue;
+                }
+
+                // Additional discriminator needed for GroupUpdate models...
+                if (messageType == SessionMessageType.SyncPlayGroupUpdate && type != typeof(SyncPlayGroupUpdateCommandMessage))
+                {
+                    continue;
+                }
+
+                var schema = context.SchemaGenerator.GenerateSchema(type, context.SchemaRepository);
+                outboundWebSocketSchemas.Add(schema);
+                outboundWebSocketDiscriminators.Add(messageType.ToString()!, schema.Reference.ReferenceV3);
+            }
+
+            var outboundWebSocketMessageSchema = new OpenApiSchema
+            {
+                Type = "object",
+                Description = "Represents the list of possible outbound websocket types",
+                Reference = new OpenApiReference
+                {
+                    Id = nameof(OutboundWebSocketMessage),
+                    Type = ReferenceType.Schema
+                },
+                OneOf = outboundWebSocketSchemas,
+                Discriminator = new OpenApiDiscriminator
+                {
+                    PropertyName = nameof(WebSocketMessage.MessageType),
+                    Mapping = outboundWebSocketDiscriminators
+                }
+            };
+
+            context.SchemaRepository.AddDefinition(nameof(OutboundWebSocketMessage), outboundWebSocketMessageSchema);
+            context.SchemaRepository.AddDefinition(
+                nameof(WebSocketMessage),
+                new OpenApiSchema
+                {
+                    Type = "object",
+                    Description = "Represents the possible websocket types",
+                    Reference = new OpenApiReference
+                    {
+                        Id = nameof(WebSocketMessage),
+                        Type = ReferenceType.Schema
+                    },
+                    OneOf = new[]
+                    {
+                        inboundWebSocketMessageSchema,
+                        outboundWebSocketMessageSchema
+                    }
+                });
+
+            // Manually generate sync play GroupUpdate messages.
+            if (!context.SchemaRepository.Schemas.TryGetValue(nameof(GroupUpdate), out var groupUpdateSchema))
+            {
+                groupUpdateSchema = context.SchemaGenerator.GenerateSchema(typeof(GroupUpdate), context.SchemaRepository);
+            }
+
+            var groupUpdateOfGroupInfoSchema = context.SchemaGenerator.GenerateSchema(typeof(GroupUpdate<GroupInfoDto>), context.SchemaRepository);
+            var groupUpdateOfGroupStateSchema = context.SchemaGenerator.GenerateSchema(typeof(GroupUpdate<GroupStateUpdate>), context.SchemaRepository);
+            var groupUpdateOfStringSchema = context.SchemaGenerator.GenerateSchema(typeof(GroupUpdate<string>), context.SchemaRepository);
+            var groupUpdateOfPlayQueueSchema = context.SchemaGenerator.GenerateSchema(typeof(GroupUpdate<PlayQueueUpdate>), context.SchemaRepository);
+
+            groupUpdateSchema.OneOf = new List<OpenApiSchema>
+            {
+                groupUpdateOfGroupInfoSchema,
+                groupUpdateOfGroupStateSchema,
+                groupUpdateOfStringSchema,
+                groupUpdateOfPlayQueueSchema
+            };
+
+            groupUpdateSchema.Discriminator = new OpenApiDiscriminator
+            {
+                PropertyName = nameof(GroupUpdate.Type),
+                Mapping = new Dictionary<string, string>
+                {
+                    { GroupUpdateType.UserJoined.ToString(), groupUpdateOfStringSchema.Reference.ReferenceV3 },
+                    { GroupUpdateType.UserLeft.ToString(), groupUpdateOfStringSchema.Reference.ReferenceV3 },
+                    { GroupUpdateType.GroupJoined.ToString(), groupUpdateOfGroupInfoSchema.Reference.ReferenceV3 },
+                    { GroupUpdateType.GroupLeft.ToString(), groupUpdateOfStringSchema.Reference.ReferenceV3 },
+                    { GroupUpdateType.StateUpdate.ToString(), groupUpdateOfGroupStateSchema.Reference.ReferenceV3 },
+                    { GroupUpdateType.PlayQueue.ToString(), groupUpdateOfPlayQueueSchema.Reference.ReferenceV3 },
+                    { GroupUpdateType.NotInGroup.ToString(), groupUpdateOfStringSchema.Reference.ReferenceV3 },
+                    { GroupUpdateType.GroupDoesNotExist.ToString(), groupUpdateOfStringSchema.Reference.ReferenceV3 },
+                    { GroupUpdateType.LibraryAccessDenied.ToString(), groupUpdateOfStringSchema.Reference.ReferenceV3 }
+                }
+            };
 
-            context.SchemaGenerator.GenerateSchema(typeof(SessionMessageType), context.SchemaRepository);
             context.SchemaGenerator.GenerateSchema(typeof(ServerDiscoveryInfo), context.SchemaRepository);
 
             foreach (var configuration in _serverConfigurationManager.GetConfigurationStores())