From b5f399623659bbb0342ff10169d0229359e07e99 Mon Sep 17 00:00:00 2001
From: EliteMasterEric <ericmyllyoja@gmail.com>
Date: Fri, 27 Sep 2024 12:21:17 -0400
Subject: [PATCH] Properly sort difficulties/variations. Validate variation IDs
 (alphanumeric only)

---
 source/funkin/play/song/Song.hx | 53 +++++++++++++++++++++++----------
 1 file changed, 38 insertions(+), 15 deletions(-)

diff --git a/source/funkin/play/song/Song.hx b/source/funkin/play/song/Song.hx
index 9023b872e..a6449bc10 100644
--- a/source/funkin/play/song/Song.hx
+++ b/source/funkin/play/song/Song.hx
@@ -156,6 +156,11 @@ class Song implements IPlayStateScriptedClass implements IRegistryEntry<SongMeta
     {
       for (vari in _data.playData.songVariations)
       {
+        if (!validateVariationId(vari)) {
+          trace('  [WARN] Variation id "$vari" is invalid, skipping...');
+          continue;
+        }
+
         var variMeta:Null<SongMetadata> = fetchVariationMetadata(id, vari);
         if (variMeta != null)
         {
@@ -407,7 +412,6 @@ class Song implements IPlayStateScriptedClass implements IRegistryEntry<SongMeta
     if (possibleVariations == null)
     {
       possibleVariations = getVariationsByCharacter(currentCharacter);
-      possibleVariations.sort(SortUtil.defaultsThenAlphabetically.bind(Constants.DEFAULT_VARIATION_LIST));
     }
     if (diffId == null) diffId = listDifficulties(null, possibleVariations)[0];
 
@@ -428,7 +432,12 @@ class Song implements IPlayStateScriptedClass implements IRegistryEntry<SongMeta
    */
   public function getVariationsByCharacter(?char:PlayableCharacter):Array<String>
   {
-    if (char == null) return variations;
+    if (char == null)
+    {
+      var result = variations;
+      result.sort(SortUtil.defaultsThenAlphabetically.bind(Constants.DEFAULT_VARIATION_LIST));
+      return result;
+    }
 
     var result = [];
     trace('Evaluating variations for ${this.id} ${char.id}: ${this.variations}');
@@ -445,6 +454,8 @@ class Song implements IPlayStateScriptedClass implements IRegistryEntry<SongMeta
       }
     }
 
+    result.sort(SortUtil.defaultsThenAlphabetically.bind(Constants.DEFAULT_VARIATION_LIST));
+
     return result;
   }
 
@@ -465,18 +476,11 @@ class Song implements IPlayStateScriptedClass implements IRegistryEntry<SongMeta
 
     if (variationIds.length == 0) return [];
 
-    // The difficulties array contains entries like 'normal', 'nightmare-erect', and 'normal-pico',
-    // so we have to map it to the actual difficulty names.
-    // We also filter out difficulties that don't match the variation or that don't exist.
-
-    var diffFiltered:Array<String> = difficulties.keys()
-      .array()
-      .map(function(diffId:String):Null<String> {
-        var difficulty:Null<SongDifficulty> = difficulties.get(diffId);
-        if (difficulty == null) return null;
-        if (variationIds.length > 0 && !variationIds.contains(difficulty.variation)) return null;
-        return difficulty.difficulty;
-      })
+    var diffFiltered:Array<String> = variationIds.map(function(variationId:String):Array<String> {
+      var metadata = _metadata.get(variationId);
+      return metadata?.playData?.difficulties ?? [];
+    })
+      .flatten()
       .filterNull()
       .distinct();
 
@@ -489,11 +493,15 @@ class Song implements IPlayStateScriptedClass implements IRegistryEntry<SongMeta
       return false;
     });
 
-    diffFiltered.sort(SortUtil.defaultsThenAlphabetically.bind(Constants.DEFAULT_DIFFICULTY_LIST));
+    diffFiltered.sort(SortUtil.defaultsThenAlphabetically.bind(Constants.DEFAULT_DIFFICULTY_LIST_FULL));
 
     return diffFiltered;
   }
 
+  /**
+   * TODO: This line of code makes me sad, but you can't really fix it without a breaking migration.
+   * @return `easy`, `erect`, `normal-pico`, etc.
+   */
   public function listSuffixedDifficulties(variationIds:Array<String>, ?showLocked:Bool, ?showHidden:Bool):Array<String>
   {
     var result = [];
@@ -509,6 +517,8 @@ class Song implements IPlayStateScriptedClass implements IRegistryEntry<SongMeta
       }
     }
 
+    result.sort(SortUtil.defaultsThenAlphabetically.bind(Constants.DEFAULT_DIFFICULTY_LIST_FULL));
+
     return result;
   }
 
@@ -629,6 +639,19 @@ class Song implements IPlayStateScriptedClass implements IRegistryEntry<SongMeta
     var meta:Null<SongMetadata> = SongRegistry.instance.parseEntryMetadataWithMigration(id, vari, version);
     return meta;
   }
+
+  static final VARIATION_REGEX = ~/^[a-z][a-z0-9]+$/;
+
+  /**
+   * Validate that the variation ID is valid.
+   * Auto-accept if it's one of the base game default variations.
+   * Reject if the ID starts with a number, or contains invalid characters.
+   */
+  static function validateVariationId(variation:String):Bool {
+    if (Constants.DEFAULT_VARIATION_LIST.contains(variation)) return true;
+
+    return VARIATION_REGEX.match(variation);
+  }
 }
 
 class SongDifficulty