Plerx2493 · Dec 19, 2024 · Jul 10, 2024 · Jul 10, 2024 · Jul 10, 2024 · Jul 11, 2024
diff --git a/DSharpPlus/Clients/DiscordClientBuilder.cs b/DSharpPlus/Clients/DiscordClientBuilder.cs
@@ -6,6 +6,9 @@
 using DSharpPlus.Logging;
 using DSharpPlus.Net;
 using DSharpPlus.Net.Gateway;
+using DSharpPlus.Net.Gateway.Compression;
+using DSharpPlus.Net.Gateway.Compression.Zlib;
+using DSharpPlus.Net.Gateway.Compression.Zstd;
 
 using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Logging;
@@ -95,6 +98,36 @@ public static DiscordClientBuilder CreateSharded
         return builder;
     }
 
+    /// <summary>
+    /// Sets the gateway compression used to zstd. This requires zstd natives to be available to the application.
+    /// </summary>
+    /// <returns>The current instance for chaining.</returns>
+    public DiscordClientBuilder UseZstdCompression()
+    {
+        this.serviceCollection.Replace<IPayloadDecompressor, ZstdDecompressor>();
+        return this;
+    }
+
+    /// <summary>
+    /// Sets the gateway compression used to zlib. This is the default compression mode.
+    /// </summary>
+    /// <returns>The current instance for chaining.</returns>
+    public DiscordClientBuilder UseZlibCompression()
+    {
+        this.serviceCollection.Replace<IPayloadDecompressor, ZlibStreamDecompressor>();
+        return this;
+    }
+
+    /// <summary>
+    /// Disables gateway compression entirely.
+    /// </summary>
+    /// <returns>The current instance for chaining.</returns>
+    public DiscordClientBuilder DisableGatewayCompression()
+    {
+        this.serviceCollection.Replace<IPayloadDecompressor, NullDecompressor>();
+        return this;
+    }
+
     /// <summary>
     /// Disables the DSharpPlus default logger for this DiscordClientBuilder.
     /// </summary>

diff --git a/DSharpPlus/Clients/MultiShardOrchestrator.cs b/DSharpPlus/Clients/MultiShardOrchestrator.cs
@@ -6,7 +6,7 @@
 using DSharpPlus.Net;
 using DSharpPlus.Net.Abstractions;
 using DSharpPlus.Net.Gateway;
-using DSharpPlus.Net.WebSocket;
+using DSharpPlus.Net.Gateway.Compression;
 
 using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Options;
@@ -22,7 +22,7 @@ public sealed class MultiShardOrchestrator : IShardOrchestrator
     private readonly DiscordApiClient apiClient;
     private readonly ShardingOptions options;
     private readonly IServiceProvider serviceProvider;
-    private readonly PayloadDecompressor decompressor;
+    private readonly IPayloadDecompressor decompressor;
 
     private uint shardCount;
     private uint stride;
@@ -54,7 +54,7 @@ public MultiShardOrchestrator
         IServiceProvider serviceProvider, 
         IOptions<ShardingOptions> options,
         DiscordApiClient apiClient,
-        PayloadDecompressor decompressor
+        IPayloadDecompressor decompressor
     )
     {
         this.apiClient = apiClient;
@@ -99,9 +99,9 @@ public async ValueTask StartAsync(DiscordActivity? activity, DiscordUserStatus?
         gwuri.AddParameter("v", "10")
              .AddParameter("encoding", "json");
 
-        if (this.decompressor.CompressionLevel == GatewayCompressionLevel.Stream)
+        if (this.decompressor.IsTransportCompression)
         {
-            gwuri.AddParameter("compress", "zlib-stream");
+            gwuri.AddParameter("compress", this.decompressor.Name);
         }
 
         this.shards = new IGatewayClient[startShards];

diff --git a/DSharpPlus/Clients/SingleShardOrchestrator.cs b/DSharpPlus/Clients/SingleShardOrchestrator.cs
@@ -4,7 +4,7 @@
 using DSharpPlus.Entities;
 using DSharpPlus.Net;
 using DSharpPlus.Net.Gateway;
-using DSharpPlus.Net.WebSocket;
+using DSharpPlus.Net.Gateway.Compression;
 
 namespace DSharpPlus.Clients;
 
@@ -15,7 +15,7 @@ public sealed class SingleShardOrchestrator : IShardOrchestrator
 {
     private readonly IGatewayClient gatewayClient;
     private readonly DiscordApiClient apiClient;
-    private readonly PayloadDecompressor decompressor;
+    private readonly IPayloadDecompressor decompressor;
 
     /// <summary>
     /// Creates a new instance of this type.
@@ -24,7 +24,7 @@ public SingleShardOrchestrator
     (
         IGatewayClient gatewayClient,
         DiscordApiClient apiClient,
-        PayloadDecompressor decompressor
+        IPayloadDecompressor decompressor
     )
     {
         this.gatewayClient = gatewayClient;
@@ -77,9 +77,9 @@ public async ValueTask StartAsync(DiscordActivity? activity, DiscordUserStatus?
         gwuri.AddParameter("v", "10")
              .AddParameter("encoding", "json");
 
-        if (this.decompressor.CompressionLevel == GatewayCompressionLevel.Stream)
+        if (this.decompressor.IsTransportCompression)
         {
-            gwuri.AddParameter("compress", "zlib-stream");
+            gwuri.AddParameter("compress", this.decompressor.Name);
         }
 
         await this.gatewayClient.ConnectAsync(gwuri.Build(), activity, status, idleSince);

diff --git a/DSharpPlus/DSharpPlus.csproj b/DSharpPlus/DSharpPlus.csproj
@@ -3,6 +3,7 @@
     <PackageId>DSharpPlus</PackageId>
     <Description>A C# API for Discord based off DiscordSharp, but rewritten to fit the API standards.</Description>
     <PackageTags>$(PackageTags), webhooks</PackageTags>
+    <AllowUnsafeBlocks>true</AllowUnsafeBlocks>
     <IsPackable>true</IsPackable>
     <AllowUnsafeBlocks>true</AllowUnsafeBlocks>
     <LangVersion>Preview</LangVersion>

diff --git a/DSharpPlus/DiscordConfiguration.cs b/DSharpPlus/DiscordConfiguration.cs
@@ -10,14 +10,6 @@ namespace DSharpPlus;
 /// </summary>
 public sealed class DiscordConfiguration
 {
-    /// <summary>
-    /// <para>Sets the level of compression for WebSocket traffic.</para>
-    /// <para>Disabling this option will increase the amount of traffic sent via WebSocket. Setting <see cref="GatewayCompressionLevel.Payload"/> will enable compression for READY and GUILD_CREATE payloads. Setting <see cref="Stream"/> will enable compression for the entire WebSocket stream, drastically reducing amount of traffic.</para>
-    /// <para>Defaults to <see cref="GatewayCompressionLevel.None"/>.</para>
-    /// </summary>
-    /// <remarks>This property's default has been changed from <see cref="GatewayCompressionLevel.Stream"/> due to a bug in the client wherein rapid reconnections causes payloads to be decompressed into a mangled string.</remarks>
-    // Here be dragons, ye who use compression.
-    public GatewayCompressionLevel GatewayCompressionLevel { internal get; set; } = GatewayCompressionLevel.None;
 
     /// <summary>
     /// Sets whether the client should attempt to cache members if exclusively using unprivileged intents.
@@ -74,7 +66,6 @@ public DiscordConfiguration()
     /// <param name="other">Client configuration to clone.</param>
     public DiscordConfiguration(DiscordConfiguration other)
     {
-        this.GatewayCompressionLevel = other.GatewayCompressionLevel;
         this.UdpClientFactory = other.UdpClientFactory;
         this.LogUnknownEvents = other.LogUnknownEvents;
         this.LogUnknownAuditlogs = other.LogUnknownAuditlogs;

diff --git a/DSharpPlus/Extensions/ServiceCollectionExtensions.InternalSetup.cs b/DSharpPlus/Extensions/ServiceCollectionExtensions.InternalSetup.cs
@@ -1,4 +1,5 @@
 using System.Net.Http;
+using System.Runtime.InteropServices;
 using System.Threading.Channels;
 
 using DSharpPlus.Clients;
@@ -7,7 +8,9 @@
 using DSharpPlus.Net.Gateway;
 using DSharpPlus.Net.InboundWebhooks;
 using DSharpPlus.Net.InboundWebhooks.Transport;
-using DSharpPlus.Net.WebSocket;
+using DSharpPlus.Net.Gateway.Compression;
+using DSharpPlus.Net.Gateway.Compression.Zlib;
+using DSharpPlus.Net.Gateway.Compression.Zstd;
 
 using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Logging;
@@ -56,7 +59,7 @@ DiscordIntents intents
             ))
             .AddTransient<ITransportService, TransportService>()
             .AddTransient<IGatewayClient, GatewayClient>()
-            .AddTransient<PayloadDecompressor>()
+            .RegisterBestDecompressor()
             .AddSingleton<IShardOrchestrator, SingleShardOrchestrator>()
             .AddSingleton<IEventDispatcher, DefaultEventDispatcher>()
             .AddSingleton<DiscordClient>();
@@ -115,7 +118,7 @@ DiscordIntents intents
             .AddKeyedSingleton("DSharpPlus.Gateway.EventChannel", Channel.CreateUnbounded<GatewayPayload>(new UnboundedChannelOptions { SingleReader = true }))
             .AddTransient<ITransportService, TransportService>()
             .AddTransient<IGatewayClient, GatewayClient>()
-            .AddTransient<PayloadDecompressor>()
+            .RegisterBestDecompressor()
             .AddSingleton<IShardOrchestrator, MultiShardOrchestrator>()
             .AddSingleton<IEventDispatcher, DefaultEventDispatcher>()
             .AddSingleton<DiscordClient>();
@@ -140,4 +143,18 @@ DiscordIntents intents
 
         return serviceCollection;
     }
+
+    private static IServiceCollection RegisterBestDecompressor(this IServiceCollection services)
+    {
+        if (NativeLibrary.TryLoad("libzstd", out _))
+        {
+            services.AddTransient<IPayloadDecompressor, ZstdDecompressor>();
+        }
+        else
+        {
+            services.AddTransient<IPayloadDecompressor, ZlibStreamDecompressor>();
+        }
+
+        return services;
+    }
 }
diff --git a/DSharpPlus/Extensions/ServiceCollectionExtensions.cs b/DSharpPlus/Extensions/ServiceCollectionExtensions.cs
@@ -1,6 +1,10 @@
 using System;
 using System.Linq;
 
+using DSharpPlus.Net.Gateway.Compression;
+using DSharpPlus.Net.Gateway.Compression.Zlib;
+using DSharpPlus.Net.Gateway.Compression.Zstd;
+
 using Microsoft.Extensions.DependencyInjection;
 
 namespace DSharpPlus.Extensions;
@@ -52,6 +56,30 @@ DiscordIntents intents
         return services;
     }
 
+    /// <summary>
+    /// Forces DSharpPlus to use zlib compression for the gateway.
+    /// </summary>
+    /// <param name="services">The service collection to configure this for.</param>
+    /// <returns>The current instance for chaining.</returns>
+    public static IServiceCollection UseZlibCompression(this IServiceCollection services)
+        => services.Replace<IPayloadDecompressor, ZlibStreamDecompressor>();
+
+    /// <summary>
+    /// Forces DSharpPlus to use zstd compression for the gateway.
+    /// </summary>
+    /// <param name="services">The service collection to configure this for.</param>
+    /// <returns>The current instance for chaining.</returns>
+    public static IServiceCollection UseZstdCompression(this IServiceCollection services)
+        => services.Replace<IPayloadDecompressor, ZstdDecompressor>();
+
+    /// <summary>
+    /// Forces DSharpPlus to disable gateway compression entirely.
+    /// </summary>
+    /// <param name="services">The service collection to configure this for.</param>
+    /// <returns>The current instance for chaining.</returns>
+    public static IServiceCollection DisableGatewayCompression(this IServiceCollection services)
+        => services.Replace<IPayloadDecompressor, NullDecompressor>();
+
     /// <summary>
 	/// Decorates a given <typeparamref name="TInterface"/> with a decorator of type <typeparamref name="TDecorator"/>.
 	/// </summary>

diff --git a/DSharpPlus/GatewayCompressionLevel.cs b/DSharpPlus/GatewayCompressionLevel.cs
diff --git a/DSharpPlus/Net/Abstractions/Gateway/GatewayIdentifyResume.cs b/DSharpPlus/Net/Abstractions/Gateway/GatewayIdentifyResume.cs
@@ -22,7 +22,7 @@ internal sealed class GatewayIdentify
     /// <summary>
     /// Gets or sets whether to encrypt websocket traffic.
     /// </summary>
-    [JsonProperty("compress")]
+    [JsonProperty("compress", DefaultValueHandling = DefaultValueHandling.Ignore)]
     public bool Compress { get; set; }
 
     /// <summary>

diff --git a/DSharpPlus/Net/Gateway/Compression/IPayloadDecompressor.cs b/DSharpPlus/Net/Gateway/Compression/IPayloadDecompressor.cs
@@ -0,0 +1,39 @@
+using System;
+
+using CommunityToolkit.HighPerformance.Buffers;
+
+namespace DSharpPlus.Net.Gateway.Compression;
+
+/// <summary>
+/// Contains functionality for decompressing inbound gateway payloads.
+/// </summary>
+public interface IPayloadDecompressor : IDisposable
+{
+    /// <summary>
+    /// Gets the name of the decompressor.
+    /// </summary>
+    public string? Name { get; }
+
+    /// <summary>
+    /// Indicates whether the present compression format is connection-wide.
+    /// </summary>
+    public bool IsTransportCompression { get; }
+
+    /// <summary>
+    /// Attempts to decompress the provided payload.
+    /// </summary>
+    /// <param name="compressed">The raw, compressed data.</param>
+    /// <param name="decompressed">A buffer writer for writing decompressed data.</param>
+    /// <returns>A value indicating whether the operation was successful.</returns>
+    public bool TryDecompress(ReadOnlySpan<byte> compressed, ArrayPoolBufferWriter<byte> decompressed);
+
+    /// <summary>
+    /// Initializes the decompressor for a new connection.
+    /// </summary>
+    public void Initialize();
+
+    /// <summary>
+    /// Frees and destroys all internal state when a connection has terminated.
+    /// </summary>
+    public void Reset();
+}
diff --git a/DSharpPlus/Net/Gateway/Compression/NullDecompressor.cs b/DSharpPlus/Net/Gateway/Compression/NullDecompressor.cs
@@ -0,0 +1,44 @@
+using System;
+
+using CommunityToolkit.HighPerformance;
+using CommunityToolkit.HighPerformance.Buffers;
+
+namespace DSharpPlus.Net.Gateway.Compression;
+
+/// <summary>
+/// Represents a decompressor that doesn't decompress at all.
+/// </summary>
+public sealed class NullDecompressor : IPayloadDecompressor
+{
+    /// <inheritdoc/>
+    public string? Name => null;
+
+    // this decompressor *technically* applies transport-wide, and this simplifies composing the IDENTIFY payload.
+    /// <inheritdoc/>
+    public bool IsTransportCompression => true;
+
+    /// <inheritdoc/>
+    public bool TryDecompress(ReadOnlySpan<byte> compressed, ArrayPoolBufferWriter<byte> decompressed)
+    {
+        decompressed.Write(compressed);
+        return true;
+    }
+
+    /// <inheritdoc/>
+    public void Dispose()
+    {
+
+    }
+
+    /// <inheritdoc/>
+    public void Reset()
+    {
+
+    }
+
+    /// <inheritdoc/>
+    public void Initialize()
+    {
+
+    }
+}