DepotDownloader: Changes to support connection reuse when downloading Steam2 content.

DepotDownloader/ContentDownloader.cs

@@ -24,6 +24,53 @@ namespace DepotDownloader
         private static Steam3Session steam3;
         private static Steam3Session.Credentials steam3Credentials;
 
+        private static IPEndPoint lastSteam2ContentServer;
+
+        private abstract class IDepotDownloadInfo
+        {
+            public int id { get; protected set; }
+            public string installDir { get; protected set; }
+            public string contentName { get; protected set; }
+
+            public abstract DownloadSource GetDownloadType();
+        }
+
+        private sealed class DepotDownloadInfo2 : IDepotDownloadInfo
+        {
+            public int version { get; private set; }
+
+            public IPEndPoint[] contentServers = null;
+            public Steam2Manifest manifest = null;
+            public List<int> NodesToDownload = new List<int>();
+            public byte[] cryptKey = null;
+
+            public override DownloadSource GetDownloadType() { return DownloadSource.Steam2; }
+
+            public DepotDownloadInfo2(int id, int version, string installDir, string contentName)
+            {
+                this.id = id;
+                this.version = version;
+                this.installDir = installDir;
+                this.contentName = contentName;
+            }
+        }
+
+        private sealed class DepotDownloadInfo3 : IDepotDownloadInfo
+        {
+            public ulong manifestId { get; private set; }
+            public byte[] depotKey;
+
+            public override DownloadSource GetDownloadType() { return DownloadSource.Steam3; }
+
+            public DepotDownloadInfo3(int depotid, ulong manifestId, string installDir, string contentName)
+            {
+                this.id = depotid;
+                this.manifestId = manifestId;
+                this.installDir = installDir;
+                this.contentName = contentName;
+            }
+        }
+
         static bool CreateDirectories( int depotId, uint depotVersion, out string installDir )
         {
             installDir = null;
@@ -404,6 +451,9 @@ namespace DepotDownloader
                 return;
             }
 
+            var infos2 = new List<DepotDownloadInfo2>();
+            var infos3 = new List<DepotDownloadInfo3>();
+
             foreach (var depot in depotIDs)
             {
                 int depotVersion = 0;
@@ -419,11 +469,74 @@ namespace DepotDownloader
                     }
                 }
 
-                DownloadDepot(depot, depotVersion, appId);
+                IDepotDownloadInfo info = GetDepotInfo(depot, depotVersion, appId);
+                if (info != null)
+                {
+                    if (info.GetDownloadType() == DownloadSource.Steam2)
+                        infos2.Add((DepotDownloadInfo2)info);
+                    else if (info.GetDownloadType() == DownloadSource.Steam3)
+                        infos3.Add((DepotDownloadInfo3)info);
+                }
             }
+
+            if( infos2.Count() > 0 )
+                DownloadSteam2( infos2 );
+
+            if( infos3.Count() > 0 )
+                DownloadSteam3( infos3 );
         }
 
-        public static void DownloadDepot(int depotId, int depotVersionRequested, int appId = 0 )
+        public static void DownloadDepotsForGame(string game)
+        {
+            var infos2 = new List<DepotDownloadInfo2>();
+            var infos3 = new List<DepotDownloadInfo3>();
+            List<int> depotIDs = CDRManager.GetDepotIDsForGameserver(game, ContentDownloader.Config.DownloadAllPlatforms);
+            foreach (var depot in depotIDs)
+            {
+                int depotVersion = CDRManager.GetLatestDepotVersion(depot, ContentDownloader.Config.PreferBetaVersions);
+                if (depotVersion == -1)
+                {
+                    Console.WriteLine("Error: Unable to find DepotID {0} in the CDR!", depot);
+                    ContentDownloader.ShutdownSteam3();
+                    continue;
+                }
+
+                IDepotDownloadInfo info = GetDepotInfo(depot, depotVersion, 0);
+                if (info.GetDownloadType() == DownloadSource.Steam2)
+                {
+                    infos2.Add((DepotDownloadInfo2)info);
+                }
+                else if (info.GetDownloadType() == DownloadSource.Steam3)
+                {
+                    infos3.Add((DepotDownloadInfo3)info);
+                }
+            }
+
+            if (infos2.Count() > 0)
+                DownloadSteam2(infos2);
+
+            if (infos3.Count() > 0)
+                DownloadSteam3(infos3);
+        }
+
+        public static void DownloadDepot(int depotId, int depotVersion, int appId = 0)
+        {
+            IDepotDownloadInfo info = GetDepotInfo(depotId, depotVersion, appId);
+            if (info.GetDownloadType() == DownloadSource.Steam2)
+            {
+                var infos = new List<DepotDownloadInfo2>();
+                infos.Add((DepotDownloadInfo2)info);
+                DownloadSteam2(infos);
+            }
+            else if (info.GetDownloadType() == DownloadSource.Steam3)
+            {
+                var infos = new List<DepotDownloadInfo3>();
+                infos.Add((DepotDownloadInfo3)info);
+                DownloadSteam3(infos);
+            }
+        }
+
+        static IDepotDownloadInfo GetDepotInfo(int depotId, int depotVersion, int appId)
         {
             if(steam3 != null && appId > 0)
                 steam3.RequestAppInfo((uint)appId);
@@ -434,25 +547,26 @@ namespace DepotDownloader
             {    
                 Console.WriteLine("Depot {0} ({1}) is not available from this account.", depotId, contentName);
 
-                return;
+                return null;
             }
 
             DownloadSource source = GetAppDownloadSource(appId);
-            uint depotVersion = (uint)depotVersionRequested;
+
+            uint uVersion = (uint)depotVersion;
 
             if (source == DownloadSource.Steam3)
             {
-                depotVersion = GetSteam3AppChangeNumber(appId);
+                uVersion = GetSteam3AppChangeNumber(appId);
             }
 
             string installDir;
-            if (!CreateDirectories(depotId, depotVersion, out installDir))
+            if (!CreateDirectories(depotId, uVersion, out installDir))
             {
                 Console.WriteLine("Error: Unable to create install directories!");
-                return;
+                return null;
             }
 
-            Console.WriteLine("Downloading \"{0}\" version {1} ...", contentName, depotVersion);
+            Console.WriteLine("Downloading \"{0}\" version {1} ...", contentName, uVersion);
 
             if(steam3 != null)
                 steam3.RequestAppTicket((uint)depotId);
@@ -463,29 +577,40 @@ namespace DepotDownloader
                 if (manifestID == 0)
                 {
                     Console.WriteLine("Depot {0} ({1}) missing public subsection or manifest section.", depotId, contentName);
-                    return;
+                    return null;
                 }
 
                 steam3.RequestDepotKey( ( uint )depotId, ( uint )appId );
                 byte[] depotKey = steam3.DepotKeys[(uint)depotId];
 
-                DownloadSteam3(depotId, manifestID, depotKey, installDir);
+                var info = new DepotDownloadInfo3( depotId, manifestID, installDir, contentName );
+                info.depotKey = depotKey;
+                return info;
             }
             else
             {
                 // steam2 path
-                DownloadSteam2(depotId, depotVersionRequested, installDir);
+                var info = new DepotDownloadInfo2(depotId, depotVersion, installDir, contentName);
+                return info;
             }
         }
 
-        private static void DownloadSteam3( int depotId, ulong depot_manifest, byte[] depotKey, string installDir )
+        private static void DownloadSteam3( List<DepotDownloadInfo3> depots )
         {
+            foreach (var depot in depots)
+            {
+                int depotId = depot.id;
+                ulong depot_manifest = depot.manifestId;
+                byte[] depotKey = depot.depotKey;
+                string installDir = depot.installDir;
+
+                Console.WriteLine("Downloading depot {0} - {1}", depot.id, depot.contentName);
                 Console.Write("Finding content servers...");
 
                 List<IPEndPoint> serverList = steam3.steamClient.GetServersOfType(EServerType.CS);
 
                 List<CDNClient.ClientEndPoint> cdnServers = null;
-            int tries = 0, counterDeferred = 0;
+                int counterDeferred = 0;
 
                 for (int i = 0; ; i++)
                 {
@@ -499,31 +624,17 @@ namespace DepotDownloader
                         break;
 
                     if (((i + 1) % serverList.Count) == 0)
-                {
-                    if (++tries > MAX_CONNECT_RETRIES)
-                    {
-                        Console.WriteLine("\nGiving up finding Steam3 content server.");
-                        return;
-                    }
-
-                    Console.Write("\nSearching for content servers... (deferred: {0})", counterDeferred);
-                    counterDeferred = 0;
-                    Thread.Sleep(1000);
-                }
-            }
-
-            if (cdnServers == null || cdnServers.Count == 0)
                     {
                         Console.WriteLine("Unable to find any Steam3 content servers");
                         return;
                     }
+                }
 
                 Console.WriteLine(" Done!");
                 Console.Write("Downloading depot manifest...");
 
                 List<CDNClient.ClientEndPoint> cdnEndpoints = cdnServers.Where((ep) => { return ep.Type == "CDN"; }).ToList();
                 List<CDNClient.ClientEndPoint> csEndpoints = cdnServers.Where((ep) => { return ep.Type == "CS"; }).ToList();
-
                 List<CDNClient> cdnClients = new List<CDNClient>();
 
                 foreach (var server in csEndpoints)
@@ -538,7 +649,6 @@ namespace DepotDownloader
                             break;
                     }
                 }
-
                 if (cdnClients.Count == 0)
                 {
                     Console.WriteLine("\nCould not initialize connection with CDN.");
@@ -666,10 +776,25 @@ namespace DepotDownloader
                     Console.WriteLine();
                 }
             }
+        }
 
         private static ContentServerClient.StorageSession GetSteam2StorageSession(IPEndPoint [] contentServers, ContentServerClient csClient, int depotId, int depotVersion)
         {
             ContentServerClient.StorageSession session = null;
+            if (csClient.IsConnected && contentServers.Contains(lastSteam2ContentServer))
+            {
+                try
+                {
+                    session = csClient.OpenStorage( (uint)depotId, (uint)depotVersion, (uint)Config.CellID, null, false );
+
+                    return session;
+                }
+                catch ( Steam2Exception )
+                {
+                    csClient.Disconnect();
+                }
+            }
+
             int tries = 0;
             int counterSocket = 0, counterSteam2 = 0;
             for (int i = 0; ; i++)
@@ -679,7 +804,8 @@ namespace DepotDownloader
                 try
                 {
                     csClient.Connect( endpoint );
-                    session = csClient.OpenStorage( (uint)depotId, (uint)depotVersion, (uint)Config.CellID, GetSteam2Credentials( (uint)depotId ) );
+                    session = csClient.OpenStorage( (uint)depotId, (uint)depotVersion, (uint)Config.CellID, GetSteam2Credentials( (uint)depotId ), true );
+                    lastSteam2ContentServer = endpoint;
                     break;
                 }
                 catch ( SocketException )
@@ -708,63 +834,61 @@ namespace DepotDownloader
             }
             return session;
         }
-        private static void DownloadSteam2( int depotId, int depotVersion, string installDir )
+        private static void DownloadSteam2( List<DepotDownloadInfo2> depots )
         {
-            Console.Write("Finding content servers...");
+            Console.WriteLine("Found depots:");
-            IPEndPoint[] contentServers = GetStorageServer(depotId, depotVersion, Config.CellID);
+            foreach (var depot in depots)
+            {
+                Console.WriteLine(" - {0}\t{1}", depot.id, depot.contentName);
+            }
 
-            if (contentServers == null || contentServers.Length == 0)
+            Console.Write("Finding content servers...");
+            foreach( var depot in depots )
+            {
+                depot.contentServers = GetStorageServer(depot.id, depot.version, Config.CellID);
+                if (depot.contentServers == null || depot.contentServers.Length == 0)
                 {
-                Console.WriteLine("\nError: Unable to find any Steam2 content servers for depot {0}, version {1}", depotId, depotVersion);
+                    Console.WriteLine("\nError: Unable to find any Steam2 content servers for depot {0}, version {1}", depot.id, depot.version);
                     return;
                 }
+            }
 
             Console.WriteLine(" Done!");
-            Console.Write("Downloading depot manifest...");
 
-            string txtManifest = Path.Combine(installDir, "manifest.txt");
+            if (!Config.DownloadManifestOnly)
+                Console.WriteLine("Building list of files to download and checking existing files...");
 
             ContentServerClient csClient = new ContentServerClient();
             csClient.ConnectionTimeout = TimeSpan.FromSeconds(STEAM2_CONNECT_TIMEOUT_SECONDS);
 
-            ContentServerClient.StorageSession session = GetSteam2StorageSession(contentServers, csClient, depotId, depotVersion);
+            ContentServerClient.StorageSession session;
+
+            foreach (var depot in depots)
+            {
+                session = GetSteam2StorageSession(depot.contentServers, csClient, depot.id, depot.version);
                 if (session == null)
-                return;
+                    continue;
 
-            Steam2Manifest manifest = null;
+                Console.Write(String.Format("Downloading manifest for depot {0}...", depot.id));
+
+                string txtManifest = Path.Combine(depot.installDir, "manifest.txt");
                 Steam2ChecksumData checksums = null;
-            List<int> NodesToDownload = new List<int>();
                 StringBuilder manifestBuilder = new StringBuilder();
-            byte[] cryptKey = CDRManager.GetDepotEncryptionKey( depotId, depotVersion );
                 string[] excludeList = null;
-            using ( session )
-            {
-                manifest = session.DownloadManifest();
-
-                Console.WriteLine( " Done!" );
-
-                if(!Config.DownloadManifestOnly)
-                {
-                    // Skip downloading checksums if we're only interested in manifests.   
-                    Console.Write("Downloading depot checksums...");
 
-                    checksums = session.DownloadChecksums();
+                depot.cryptKey = CDRManager.GetDepotEncryptionKey(depot.id, depot.version);
+                depot.manifest = session.DownloadManifest();
 
                 Console.WriteLine(" Done!");
-                }
 
                 if (Config.UsingExclusionList)
-                    excludeList = GetExcludeList( session, manifest );
+                    excludeList = GetExcludeList(session, depot.manifest);
-            }
-            csClient.Disconnect();
 
-            if(!Config.DownloadManifestOnly)
-                Console.WriteLine("Building list of files to download and checking existing files...");
                 // Build a list of files that need downloading.
-            for ( int x = 0 ; x < manifest.Nodes.Count ; ++x )
+                for (int x = 0; x < depot.manifest.Nodes.Count; ++x)
                 {
-                var dirEntry = manifest.Nodes[ x ];
+                    var dirEntry = depot.manifest.Nodes[x];
-                string downloadPath = Path.Combine( installDir, dirEntry.FullName.ToLower() );
+                    string downloadPath = Path.Combine(depot.installDir, dirEntry.FullName.ToLower());
                     if (Config.DownloadManifestOnly)
                     {
                         if (dirEntry.FileID == -1)
@@ -795,10 +919,20 @@ namespace DepotDownloader
                         continue;
                     }
 
+                    if (checksums == null)
+                    {
+                        // Skip downloading checksums if we're only interested in manifests.   
+                        Console.Write(String.Format("Downloading checksums for depot {0}...", depot.id));
+
+                        checksums = session.DownloadChecksums();
+
+                        Console.WriteLine(" Done!");
+                    }
+
                     FileInfo fi = new FileInfo(downloadPath);
                     if (fi.Exists)
                     {
-                    float perc = ( ( float )x / ( float )manifest.Nodes.Count ) * 100.0f;
+                        float perc = ((float)x / (float)depot.manifest.Nodes.Count) * 100.0f;
                         Console.WriteLine("{0,6:#00.00}%\t{1}", perc, downloadPath);
                         // Similar file, let's check checksums
                         if (fi.Length == dirEntry.SizeOrCount &&
@@ -811,7 +945,7 @@ namespace DepotDownloader
                         // This will keep symbolic/hard link targets from being overwritten.
                         fi.Delete();
                     }
-                NodesToDownload.Add(x);
+                    depot.NodesToDownload.Add(x);
                 }
 
                 if (Config.DownloadManifestOnly)
@@ -819,24 +953,27 @@ namespace DepotDownloader
                     File.WriteAllText(txtManifest, manifestBuilder.ToString());
                     return;
                 }
+            }
 
+            Console.WriteLine("Downloading selected files.");
+            foreach( var depot in depots )
+            {
+                if (depot.NodesToDownload.Count == 0)
+                    continue;
 
-            session = GetSteam2StorageSession(contentServers, csClient, depotId, depotVersion);
+                session = GetSteam2StorageSession(depot.contentServers, csClient, depot.id, depot.version);
                 if(session == null)
-                return;
+                    continue;
 
-            using ( session )
+                for ( int x = 0 ; x < depot.NodesToDownload.Count ; ++x )
-            {
-                Console.WriteLine("Downloading selected files.");
-                for ( int x = 0 ; x < NodesToDownload.Count ; ++x )
                 {
-                    var dirEntry = manifest.Nodes[ NodesToDownload[ x ] ];
+                    var dirEntry = depot.manifest.Nodes[ depot.NodesToDownload[ x ] ];
-                    string downloadPath = Path.Combine( installDir, dirEntry.FullName.ToLower() );
+                    string downloadPath = Path.Combine( depot.installDir, dirEntry.FullName.ToLower() );
 
-                    float perc = ( ( float )x / ( float )NodesToDownload.Count ) * 100.0f;
+                    float perc = ( ( float )x / ( float )depot.NodesToDownload.Count ) * 100.0f;
                     Console.WriteLine("{0,6:#00.00}%\t{1}", perc, downloadPath);
                     
-                    var file = session.DownloadFile( dirEntry, ContentServerClient.StorageSession.DownloadPriority.High, cryptKey );
+                    var file = session.DownloadFile( dirEntry, ContentServerClient.StorageSession.DownloadPriority.High, depot.cryptKey );
                     File.WriteAllBytes( downloadPath, file );
                 }
             }