From e2194fc8733d74444f3f807901e3127770d8d180 Mon Sep 17 00:00:00 2001
From: Pavel Djundik <xPaw@users.noreply.github.com>
Date: Mon, 26 Aug 2024 17:24:31 +0300
Subject: [PATCH] Attempt to get cdn auth when getting 403

---
 DepotDownloader/ContentDownloader.cs | 42 ++++++++++++++++++++++++++--
 DepotDownloader/Steam3Session.cs     | 27 ++++++++++++++++++
 2 files changed, 67 insertions(+), 2 deletions(-)

diff --git a/DepotDownloader/ContentDownloader.cs b/DepotDownloader/ContentDownloader.cs
index 6d49bd27..fa97b98b 100644
--- a/DepotDownloader/ContentDownloader.cs
+++ b/DepotDownloader/ContentDownloader.cs
@@ -773,6 +773,13 @@ namespace DepotDownloader
                         {
                             connection = cdnPool.GetConnection(cts.Token);
 
+                            string cdnToken = null;
+                            if (steam3.CDNAuthTokens.TryGetValue((depot.DepotId, connection.Host), out var authTokenCallbackPromise))
+                            {
+                                var result = await authTokenCallbackPromise.Task;
+                                cdnToken = result.Token;
+                            }
+
                             var now = DateTime.Now;
 
                             // In order to download this manifest, we need the current manifest request code
@@ -806,7 +813,8 @@ namespace DepotDownloader
                                 manifestRequestCode,
                                 connection,
                                 depot.DepotKey,
-                                cdnPool.ProxyServer).ConfigureAwait(false);
+                                cdnPool.ProxyServer,
+                                cdnToken).ConfigureAwait(false);
 
                             cdnPool.ReturnConnection(connection);
                         }
@@ -816,6 +824,16 @@ namespace DepotDownloader
                         }
                         catch (SteamKitWebRequestException e)
                         {
+                            // If the CDN returned 403, attempt to get a cdn auth if we didn't yet
+                            if (e.StatusCode == HttpStatusCode.Forbidden && !steam3.CDNAuthTokens.ContainsKey((depot.DepotId, connection.Host)))
+                            {
+                                await steam3.RequestCDNAuthToken(depot.AppId, depot.DepotId, connection);
+
+                                cdnPool.ReturnConnection(connection);
+
+                                continue;
+                            }
+
                             cdnPool.ReturnBrokenConnection(connection);
 
                             if (e.StatusCode == HttpStatusCode.Unauthorized || e.StatusCode == HttpStatusCode.Forbidden)
@@ -1215,6 +1233,13 @@ namespace DepotDownloader
                     {
                         connection = cdnPool.GetConnection(cts.Token);
 
+                        string cdnToken = null;
+                        if (steam3.CDNAuthTokens.TryGetValue((depot.DepotId, connection.Host), out var authTokenCallbackPromise))
+                        {
+                            var result = await authTokenCallbackPromise.Task;
+                            cdnToken = result.Token;
+                        }
+
                         DebugLog.WriteLine("ContentDownloader", "Downloading chunk {0} from {1} with {2}", chunkID, connection, cdnPool.ProxyServer != null ? cdnPool.ProxyServer : "no proxy");
                         written = await cdnPool.CDNClient.DownloadDepotChunkAsync(
                             depot.DepotId,
@@ -1222,7 +1247,8 @@ namespace DepotDownloader
                             connection,
                             chunkBuffer,
                             depot.DepotKey,
-                            cdnPool.ProxyServer).ConfigureAwait(false);
+                            cdnPool.ProxyServer,
+                            cdnToken).ConfigureAwait(false);
 
                         cdnPool.ReturnConnection(connection);
 
@@ -1234,6 +1260,18 @@ namespace DepotDownloader
                     }
                     catch (SteamKitWebRequestException e)
                     {
+                        // If the CDN returned 403, attempt to get a cdn auth if we didn't yet,
+                        // if auth task already exists, make sure it didn't complete yet, so that it gets awaited above
+                        if (e.StatusCode == HttpStatusCode.Forbidden &&
+                            (!steam3.CDNAuthTokens.TryGetValue((depot.DepotId, connection.Host), out var authTokenCallbackPromise) || !authTokenCallbackPromise.Task.IsCompleted))
+                        {
+                            await steam3.RequestCDNAuthToken(depot.AppId, depot.DepotId, connection);
+
+                            cdnPool.ReturnConnection(connection);
+
+                            continue;
+                        }
+
                         cdnPool.ReturnBrokenConnection(connection);
 
                         if (e.StatusCode == HttpStatusCode.Unauthorized || e.StatusCode == HttpStatusCode.Forbidden)
diff --git a/DepotDownloader/Steam3Session.cs b/DepotDownloader/Steam3Session.cs
index 0246272a..92620636 100644
--- a/DepotDownloader/Steam3Session.cs
+++ b/DepotDownloader/Steam3Session.cs
@@ -1,4 +1,5 @@
 using System;
+using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.Collections.ObjectModel;
 using System.Linq;
@@ -7,6 +8,7 @@ using System.Threading.Tasks;
 using QRCoder;
 using SteamKit2;
 using SteamKit2.Authentication;
+using SteamKit2.CDN;
 using SteamKit2.Internal;
 
 namespace DepotDownloader
@@ -24,6 +26,7 @@ namespace DepotDownloader
         public Dictionary<uint, ulong> AppTokens { get; } = [];
         public Dictionary<uint, ulong> PackageTokens { get; } = [];
         public Dictionary<uint, byte[]> DepotKeys { get; } = [];
+        public ConcurrentDictionary<(uint, string), TaskCompletionSource<SteamApps.CDNAuthTokenCallback>> CDNAuthTokens { get; } = [];
         public Dictionary<uint, SteamApps.PICSProductInfoCallback.PICSProductInfo> AppInfo { get; } = [];
         public Dictionary<uint, SteamApps.PICSProductInfoCallback.PICSProductInfo> PackageInfo { get; } = [];
         public Dictionary<string, byte[]> AppBetaPasswords { get; } = [];
@@ -282,6 +285,30 @@ namespace DepotDownloader
             return requestCode;
         }
 
+        public async Task RequestCDNAuthToken(uint appid, uint depotid, Server server)
+        {
+            var cdnKey = (depotid, server.Host);
+            var completion = new TaskCompletionSource<SteamApps.CDNAuthTokenCallback>();
+
+            if (bAborted || !CDNAuthTokens.TryAdd(cdnKey, completion))
+            {
+                return;
+            }
+
+            DebugLog.WriteLine(nameof(Steam3Session), $"Requesting CDN auth token for {server.Host}");
+
+            var cdnAuth = await steamApps.GetCDNAuthToken(appid, depotid, server.Host);
+
+            Console.WriteLine($"Got CDN auth token for {server.Host} result: {cdnAuth.Result} (expires {cdnAuth.Expiration})");
+
+            if (cdnAuth.Result != EResult.OK)
+            {
+                return;
+            }
+
+            completion.TrySetResult(cdnAuth);
+        }
+
         public void CheckAppBetaPassword(uint appid, string password)
         {
             var completed = false;