From 6b4c112a2b0a81096834e5e46b3352932389f459 Mon Sep 17 00:00:00 2001
From: JKorf <jankorf91@gmail.com>
Date: Tue, 23 Apr 2024 09:10:35 +0200
Subject: [PATCH] Fix to prevent endless looping when reconnecting an
 authenticated websocket

---
 .../SpotApi/MexcSocketClientSpotApi.cs        | 31 +++++++++++++++++--
 Mexc.Net/Mexc.Net.xml                         |  6 ++++
 2 files changed, 34 insertions(+), 3 deletions(-)

diff --git a/Mexc.Net/Clients/SpotApi/MexcSocketClientSpotApi.cs b/Mexc.Net/Clients/SpotApi/MexcSocketClientSpotApi.cs
index 95a9bda..c84548a 100644
--- a/Mexc.Net/Clients/SpotApi/MexcSocketClientSpotApi.cs
+++ b/Mexc.Net/Clients/SpotApi/MexcSocketClientSpotApi.cs
@@ -131,24 +131,49 @@ public async Task<CallResult<UpdateSubscription>> SubscribeToMiniTickerUpdatesAs
         /// <inheritdoc />
         public async Task<CallResult<UpdateSubscription>> SubscribeToAccountUpdatesAsync(string listenKey, Action<DataEvent<MexcAccountUpdate>> handler, CancellationToken ct = default)
         {
-            var subscription = new MexcSubscription<MexcAccountUpdate>(_logger, new[] { "spot@private.account.v3.api" }, handler, false);
+            var subscription = new MexcSubscription<MexcAccountUpdate>(_logger, new[] { "spot@private.account.v3.api" }, handler, true);
             return await SubscribeAsync(BaseAddress + "?listenKey=" + listenKey, subscription, ct).ConfigureAwait(false);
         }
 
         /// <inheritdoc />
         public async Task<CallResult<UpdateSubscription>> SubscribeToOrderUpdatesAsync(string listenKey, Action<DataEvent<MexcUserOrderUpdate>> handler, CancellationToken ct = default)
         {
-            var subscription = new MexcSubscription<MexcUserOrderUpdate>(_logger, new[] { "spot@private.orders.v3.api" }, handler, false);
+            var subscription = new MexcSubscription<MexcUserOrderUpdate>(_logger, new[] { "spot@private.orders.v3.api" }, handler, true);
             return await SubscribeAsync(BaseAddress + "?listenKey=" + listenKey, subscription, ct).ConfigureAwait(false);
         }
 
         /// <inheritdoc />
         public async Task<CallResult<UpdateSubscription>> SubscribeToUserTradeUpdatesAsync(string listenKey, Action<DataEvent<MexcUserTradeUpdate>> handler, CancellationToken ct = default)
         {
-            var subscription = new MexcSubscription<MexcUserTradeUpdate>(_logger, new[] { "spot@private.deals.v3.api" }, handler, false);
+            var subscription = new MexcSubscription<MexcUserTradeUpdate>(_logger, new[] { "spot@private.deals.v3.api" }, handler, true);
             return await SubscribeAsync(BaseAddress + "?listenKey=" + listenKey, subscription, ct).ConfigureAwait(false);
         }
 
+        /// <inheritdoc />
+        protected override async Task<Uri?> GetReconnectUriAsync(SocketConnection connection)
+        {
+            if (connection.Subscriptions.Any(s => s.Authenticated))
+            {
+                // If any of the subs on the connection is authenticated we request a new listenkey
+                // to prevent endlessly looping if the listenkey happens to be expired
+                var creds = ApiOptions.ApiCredentials ?? ClientOptions.ApiCredentials;
+                var client = new MexcRestClient(opts =>
+                {
+                    if (creds != null)
+                        opts.ApiCredentials = creds;
+                });
+
+                var listenKeyResult = await client.SpotApi.Account.StartUserStreamAsync().ConfigureAwait(false);
+                if (listenKeyResult)
+                    return new Uri(BaseAddress + "?listenKey=" + listenKeyResult.Data);
+            }
+
+            return await base.GetReconnectUriAsync(connection).ConfigureAwait(false);
+        }
+
+        /// <inheritdoc />
+        protected override Query? GetAuthenticationRequest() => null;
+
         private static string GetIntervalString(KlineInterval interval)
             => interval switch
             {
diff --git a/Mexc.Net/Mexc.Net.xml b/Mexc.Net/Mexc.Net.xml
index aaab199..62779ef 100644
--- a/Mexc.Net/Mexc.Net.xml
+++ b/Mexc.Net/Mexc.Net.xml
@@ -302,6 +302,12 @@
         <member name="M:Mexc.Net.Clients.SpotApi.MexcSocketClientSpotApi.SubscribeToUserTradeUpdatesAsync(System.String,System.Action{CryptoExchange.Net.Objects.Sockets.DataEvent{Mexc.Net.Objects.Models.Spot.MexcUserTradeUpdate}},System.Threading.CancellationToken)">
             <inheritdoc />
         </member>
+        <member name="M:Mexc.Net.Clients.SpotApi.MexcSocketClientSpotApi.GetReconnectUriAsync(CryptoExchange.Net.Sockets.SocketConnection)">
+            <inheritdoc />
+        </member>
+        <member name="M:Mexc.Net.Clients.SpotApi.MexcSocketClientSpotApi.GetAuthenticationRequest">
+            <inheritdoc />
+        </member>
         <member name="T:Mexc.Net.Enums.AccountType">
             <summary>
             Account type