Support Content-Disposition header for download filenames (#419)

Add parameter to allow downloads to prefer to use the name provided by the server for downloaded files

Author: Top-Cat

ModAssistant/Classes/External Interfaces/Playlists.cs

@@ -2,7 +2,6 @@
 using System.IO;
 using System.Linq;
 using System.Threading.Tasks;
-using System.Web;
 using System.Windows;
 using static ModAssistant.Http;
 
@@ -34,14 +33,11 @@ public static async Task DownloadAll(Uri uri)
 
         public static async Task<string> Get(Uri url)
         {
-            string filename = HttpUtility.UrlDecode(url.Segments.Last());
-            string absolutePath = Path.Combine(BeatSaberPath, PlaylistsFolder, filename);
             try
             {
                 CreatePlaylistsFolder();
-                await Utils.DownloadAsset(url.ToString(), PlaylistsFolder, filename);
-
-                return absolutePath;
+                string filename = await Utils.DownloadAsset(url.ToString(), PlaylistsFolder, preferContentDisposition: true);
+                return Path.Combine(BeatSaberPath, PlaylistsFolder, filename);
             }
             catch
             {

ModAssistant/Classes/External Interfaces/Utils.cs

@@ -29,50 +29,49 @@ public static void SetMessage(string message)
             }
         }
 
-        public static async Task DownloadAsset(string link, string folder, bool showNotifcation, string fileName = null)
-        {
-            await DownloadAsset(link, folder, fileName, null, showNotifcation);
-        }
-
-        public static async Task DownloadAsset(string link, string folder, string fileName = null, string displayName = null)
-        {
-            await DownloadAsset(link, folder, fileName, displayName, true);
-        }
-
-        public static async Task DownloadAsset(string link, string folder, string fileName, string displayName, bool showNotification, bool beatsaver = false)
+        public static async Task<string> DownloadAsset(string link, string folder, string fileName = null, string displayName = null, bool showNotification = true, bool beatsaver = false, bool preferContentDisposition = false)
         {
             if (string.IsNullOrEmpty(BeatSaberPath))
             {
                 ModAssistant.Utils.SendNotify((string)Application.Current.FindResource("OneClick:InstallDirNotFound"));
             }
             try
             {
-                Directory.CreateDirectory(Path.Combine(BeatSaberPath, folder));
+                var parentDir = Path.Combine(BeatSaberPath, folder);
+                Directory.CreateDirectory(parentDir);
+
                 if (string.IsNullOrEmpty(fileName))
                 {
-                    fileName = WebUtility.UrlDecode(Path.Combine(BeatSaberPath, folder, new Uri(link).Segments.Last()));
+                    fileName = new Uri(link).Segments.Last();
+                }
+
+                if (beatsaver)
+                {
+                    fileName = WebUtility.UrlDecode(Path.Combine(parentDir, fileName));
+                    await BeatSaver.Download(link, fileName);
                 }
                 else
                 {
-                    fileName = WebUtility.UrlDecode(Path.Combine(BeatSaberPath, folder, fileName));
+                    fileName = await ModAssistant.Utils.Download(link, parentDir, fileName, preferContentDisposition);
                 }
+
                 if (string.IsNullOrEmpty(displayName))
                 {
                     displayName = Path.GetFileNameWithoutExtension(fileName);
                 }
 
-                if (beatsaver) await BeatSaver.Download(link, fileName);
-                else await ModAssistant.Utils.Download(link, fileName);
-
                 if (showNotification)
                 {
                     SetMessage(string.Format((string)Application.Current.FindResource("OneClick:InstalledAsset"), displayName));
                 }
+
+                return fileName;
             }
             catch
             {
                 SetMessage((string)Application.Current.FindResource("OneClick:AssetInstallFailed"));
                 App.CloseWindowOnFinish = false;
+                return null;
             }
         }
     }

ModAssistant/Classes/Updater.cs

@@ -77,7 +77,7 @@ public static async Task StartUpdate()
 
                 File.Move(Utils.ExePath, OldExe);
 
-                await Utils.Download(DownloadLink, NewExe);
+                await Utils.Download(DownloadLink, "", NewExe);
                 RunNew();
             }
         }

ModAssistant/Classes/Utils.cs

@@ -5,6 +5,7 @@
 using System.IO;
 using System.Linq;
 using System.Management;
+using System.Net;
 using System.Security.Cryptography;
 using System.Security.Principal;
 using System.Text;
@@ -438,14 +439,28 @@ public static void Log(string message, string severity = "LOG")
             File.AppendAllText(logFile, $"[{DateTime.UtcNow:yyyy-mm-dd HH:mm:ss.ffffff}][{severity.ToUpper()}] {message}\n");
         }
 
-        public static async Task Download(string link, string output)
+        public static async Task<string> Download(string link, string folder, string output, bool preferContentDisposition = false)
         {
             var resp = await HttpClient.GetAsync(link);
+            var cdFilename = resp.Content.Headers.ContentDisposition.FileName.Trim('"');
+            // Prevent path traversal
+            if (cdFilename.IndexOfAny(Path.GetInvalidFileNameChars()) >= 0)
+            {
+                cdFilename = null;
+            }
+
+            var filename = WebUtility.UrlDecode(Path.Combine(
+                folder,
+                (preferContentDisposition ? cdFilename : null) ?? output
+            ));
+
             using (var stream = await resp.Content.ReadAsStreamAsync())
-            using (var fs = new FileStream(output, FileMode.OpenOrCreate, FileAccess.Write))
+            using (var fs = new FileStream(filename, FileMode.OpenOrCreate, FileAccess.Write))
             {
                 await stream.CopyToAsync(fs);
             }
+
+            return filename;
         }
 
         private delegate void ShowMessageBoxDelegate(string Message, string Caption);