Some internal changes to command handling

compujuckel · Jan 27, 2024 · 21f9ffe · 21f9ffe
1 parent 56e21ac
commit 21f9ffe
Show file tree

Hide file tree

Showing 8 changed files with 39 additions and 47 deletions.
diff --git a/AssettoServer/Commands/Attributes/RequireAdminAttribute.cs b/AssettoServer/Commands/Attributes/RequireAdminAttribute.cs
@@ -13,8 +13,7 @@ public override async ValueTask<CheckResult> CheckAsync(CommandContext context)
     {
         switch (context)
         {
-            case RconCommandContext:
-            case ChatCommandContext { Client.IsAdministrator: true }:
+            case BaseCommandContext { IsAdministrator: true }:
                 return CheckResult.Successful;
             case ChatCommandContext chatContext:
             {
@@ -33,11 +32,10 @@ public override async ValueTask<CheckResult> CheckAsync(CommandContext context)
                         }
                     }
                 }
-
-                return CheckResult.Failed("You are not an administrator.");
+                goto default;
             }
             default:
-                return CheckResult.Failed("Invalid command context.");
+                return CheckResult.Failed("You are not an administrator.");
         }
     }
 }
diff --git a/AssettoServer/Commands/ChatService.cs b/AssettoServer/Commands/ChatService.cs
@@ -17,25 +17,24 @@ public class ChatService
 {
     private readonly EntryCarManager _entryCarManager;
     private readonly Func<ACTcpClient, ChatCommandContext> _chatContextFactory;
-    private readonly Func<RconClient, int, RconCommandContext> _rconContextFactory;
-    private readonly CommandService _commandService = new(new CommandServiceConfiguration
-    {
-        DefaultRunMode = RunMode.Parallel
-    });
+    private readonly CommandService _commandService;
 
     public event EventHandler<ACTcpClient, ChatEventArgs>? MessageReceived;
 
-    public ChatService(ACPluginLoader loader, Func<ACTcpClient, ChatCommandContext> chatContextFactory, ACClientTypeParser acClientTypeParser, EntryCarManager entryCarManager, Func<RconClient, int, RconCommandContext> rconContextFactory)
+    public ChatService(ACPluginLoader loader,
+        Func<ACTcpClient, ChatCommandContext> chatContextFactory,
+        ACClientTypeParser acClientTypeParser,
+        EntryCarManager entryCarManager,
+        CommandService commandService)
     {
         _chatContextFactory = chatContextFactory;
         _entryCarManager = entryCarManager;
-        _rconContextFactory = rconContextFactory;
+        _commandService = commandService;
         _entryCarManager.ClientConnected += OnClientConnected;
 
         _commandService.AddModules(Assembly.GetEntryAssembly());
         _commandService.AddTypeParser(acClientTypeParser);
         _commandService.CommandExecutionFailed += OnCommandExecutionFailed;
-        _commandService.CommandExecuted += OnCommandExecuted;
 
         foreach (var plugin in loader.LoadedPlugins)
         { 
@@ -48,42 +47,17 @@ private void OnClientConnected(ACTcpClient sender, EventArgs args)
         sender.ChatMessageReceived += OnChatMessageReceived;
     }
 
-    private static ValueTask OnCommandExecuted(object? sender, CommandExecutedEventArgs args)
-    {
-        if (args.Context is RconCommandContext context)
-        {
-            context.SendRconResponse();
-        }
-
-        return ValueTask.CompletedTask;
-    }
-
     private async Task ProcessCommandAsync(ACTcpClient client, ChatMessage message)
-    {
-        var context = _chatContextFactory(client);
-        var result = await _commandService.ExecuteAsync(message.Message, context);
+        => await ProcessCommandAsync(_chatContextFactory(client), message.Message);
 
-        if (result is ChecksFailedResult checksFailedResult)
-            context.Reply(checksFailedResult.FailedChecks[0].Result.FailureReason);
-        else if (result is FailedResult failedResult)
-            context.Reply(failedResult.FailureReason);
-    }
-
-    public async Task ProcessCommandAsync(RconClient client, int requestId, string command)
+    public async Task ProcessCommandAsync(BaseCommandContext context, string command)
     {
-        var context = _rconContextFactory(client, requestId);
         var result = await _commandService.ExecuteAsync(command, context);
 
         if (result is ChecksFailedResult checksFailedResult)
-        {
             context.Reply(checksFailedResult.FailedChecks[0].Result.FailureReason);
-            context.SendRconResponse();
-        }
         else if (result is FailedResult failedResult)
-        {
             context.Reply(failedResult.FailureReason);
-            context.SendRconResponse();
-        }
     }
 
     private ValueTask OnCommandExecutionFailed(object? sender, CommandExecutionFailedEventArgs e)

diff --git a/AssettoServer/Commands/Contexts/BaseCommandContext.cs b/AssettoServer/Commands/Contexts/BaseCommandContext.cs
@@ -11,6 +11,8 @@ public abstract class BaseCommandContext(
         IServiceProvider? serviceProvider = null)
     : CommandContext(serviceProvider)
 {
+    public virtual bool IsAdministrator => false;
+
     public abstract void Reply(string message);
 
     public virtual void Broadcast(string message)

diff --git a/AssettoServer/Commands/Contexts/ChatCommandContext.cs b/AssettoServer/Commands/Contexts/ChatCommandContext.cs
@@ -13,6 +13,8 @@ public class ChatCommandContext(
 {
     public ACTcpClient Client { get; } = client;
 
+    public override bool IsAdministrator => Client.IsAdministrator;
+
     public override void Reply(string message)
     {
         Client.SendPacket(new ChatMessage { SessionId = 255, Message = message });

diff --git a/AssettoServer/Commands/Contexts/RconCommandContext.cs b/AssettoServer/Commands/Contexts/RconCommandContext.cs
@@ -16,6 +16,8 @@ public class RconCommandContext(
     public int RconRequestId { get; } = rconRequestId;
     public StringBuilder RconResponseBuilder { get; } = new();
 
+    public override bool IsAdministrator => true;
+
     public override void Reply(string message)
     {
         RconResponseBuilder.AppendLine(message);

diff --git a/AssettoServer/Commands/Modules/AdminModule.cs b/AssettoServer/Commands/Modules/AdminModule.cs
@@ -221,4 +221,12 @@ public void Say([Remainder] string message)
     {
         Broadcast("CONSOLE: " + message);
     }
+
+#if DEBUG
+    [Command("exception")]
+    public void ThrowException()
+    {
+        throw new Exception("test");
+    }
+#endif
 }
diff --git a/AssettoServer/Network/Http/Startup.cs b/AssettoServer/Network/Http/Startup.cs
@@ -28,6 +28,7 @@
 using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Hosting;
 using Prometheus;
+using Qmmands;
 
 namespace AssettoServer.Network.Http;
 
@@ -88,6 +89,7 @@ public void ConfigureContainer(ContainerBuilder builder)
         builder.RegisterType<ACClientAuthentication>().AsSelf().SingleInstance().AutoActivate();
         builder.RegisterType<HttpInfoCache>().AsSelf().As<IAssettoServerAutostart>().SingleInstance();
         builder.RegisterType<DefaultCMContentProvider>().As<ICMContentProvider>().SingleInstance();
+        builder.RegisterType<CommandService>().AsSelf().SingleInstance();
 
         if (_configuration.Extra.EnableLegacyPluginInterface)
         {

diff --git a/AssettoServer/Network/Rcon/RconClient.cs b/AssettoServer/Network/Rcon/RconClient.cs
@@ -5,6 +5,7 @@
 using System.Threading.Channels;
 using System.Threading.Tasks;
 using AssettoServer.Commands;
+using AssettoServer.Commands.Contexts;
 using AssettoServer.Server.Configuration;
 using AssettoServer.Shared.Network.Packets;
 using AssettoServer.Shared.Network.Packets.Outgoing;
@@ -22,16 +23,18 @@ public class RconClient
     private readonly Channel<IOutgoingNetworkPacket> _outgoingPacketChannel = Channel.CreateBounded<IOutgoingNetworkPacket>(256);
     private readonly Memory<byte> _tcpSendBuffer = new byte[4096];
     private readonly CancellationTokenSource _disconnectTokenSource = new ();
+    private readonly Func<RconClient, int, RconCommandContext> _rconContextFactory;
 
     private bool _isAuthenticated = false;
     private bool _isDisconnectRequested = false;
 
     private Task SendLoopTask { get; set; } = null!;
 
-    public RconClient(ACServerConfiguration configuration, TcpClient client, ChatService chatService)
+    public RconClient(ACServerConfiguration configuration, TcpClient client, ChatService chatService, Func<RconClient, int, RconCommandContext> rconContextFactory)
     {
         _client = client;
         _chatService = chatService;
+        _rconContextFactory = rconContextFactory;
         _configuration = configuration;
         _stream = client.GetStream();
     }
@@ -122,13 +125,14 @@ private async Task ReceiveLoopAsync()
                 }
                 else if (_isAuthenticated)
                 {
-                    switch (type)
+                    if (type == RconProtocolIn.ExecCommand)
                     {
-                        case RconProtocolIn.ExecCommand:
-                            var packet = reader.ReadPacket<ExecCommandPacket>();
-                            Log.Debug("RCON ({IpEndpoint}): {Command}", _client.Client.RemoteEndPoint?.ToString(), packet.Command);
-                            await _chatService.ProcessCommandAsync(this, requestId, packet.Command);
-                            break;
+                        var packet = reader.ReadPacket<ExecCommandPacket>();
+                        Log.Debug("RCON ({IpEndpoint}): {Command}", _client.Client.RemoteEndPoint?.ToString(),
+                            packet.Command);
+                        var context = _rconContextFactory(this, requestId);
+                        await _chatService.ProcessCommandAsync(context, packet.Command);
+                        context.SendRconResponse();
                     }
                 }
             }