From 104376b43fba7ce12aae9bce9e217de62453614b Mon Sep 17 00:00:00 2001
From: miniduikboot <mini@duikbo.at>
Date: Sun, 17 Nov 2024 16:18:49 +0100
Subject: [PATCH] Record and restore player colors

This prevents situations where a player that requested red has a
different color when rejoining because red is the default color of a
PlayerInfo object. This is most noticeable if the host is late with
rejoining, and a lot of players are alreading in WaitingForHost limbo.

When the host finally rejoins, they will not assign Red immediately to
players that requested Red as it is "in use", and instead assign the
next color.

By recording and restoring this is avoided, and players tend to get
their own colors back in most cases.

This is not something that officials do at the moment, but I don't think
this has an impact on mods, as a normal checkcolor/setcolor cycle will
take place as usual.
---
 src/Impostor.Server/Net/ClientBase.cs                  |  3 +++
 .../Net/Inner/Objects/InnerPlayerControl.cs            | 10 ++++++++++
 src/Impostor.Server/Net/State/Game.Data.cs             |  8 ++++++++
 3 files changed, 21 insertions(+)

diff --git a/src/Impostor.Server/Net/ClientBase.cs b/src/Impostor.Server/Net/ClientBase.cs
index a8945fa..1aee622 100644
--- a/src/Impostor.Server/Net/ClientBase.cs
+++ b/src/Impostor.Server/Net/ClientBase.cs
@@ -3,6 +3,7 @@ using System.Collections.Generic;
 using System.Threading.Tasks;
 using Impostor.Api;
 using Impostor.Api.Innersloth;
+using Impostor.Api.Innersloth.Customization;
 using Impostor.Api.Net;
 using Impostor.Server.Net.State;
 
@@ -39,6 +40,8 @@ namespace Impostor.Server.Net
 
         public ClientPlayer? Player { get; set; }
 
+        public ColorType? PreviousColor { get; set; } = null;
+
         IClientPlayer? IClient.Player => Player;
 
         public virtual ValueTask<bool> ReportCheatAsync(CheatContext context, CheatCategory category, string message)
diff --git a/src/Impostor.Server/Net/Inner/Objects/InnerPlayerControl.cs b/src/Impostor.Server/Net/Inner/Objects/InnerPlayerControl.cs
index 6503ed4..e9cec15 100644
--- a/src/Impostor.Server/Net/Inner/Objects/InnerPlayerControl.cs
+++ b/src/Impostor.Server/Net/Inner/Objects/InnerPlayerControl.cs
@@ -845,6 +845,16 @@ namespace Impostor.Server.Net.Inner.Objects
 
             PlayerInfo.CurrentOutfit.Color = color;
 
+            // Record the color so it can be restored on the next game
+            if (!Game.IsHostAuthoritive && Game.TryGetPlayer(OwnerId, out var clientPlayer))
+            {
+                clientPlayer.Client.PreviousColor = color;
+            }
+            else
+            {
+                _logger.LogWarning("Tried to record color, but couldn't get player with id {PlayerId}", OwnerId);
+            }
+
             return true;
         }
 
diff --git a/src/Impostor.Server/Net/State/Game.Data.cs b/src/Impostor.Server/Net/State/Game.Data.cs
index 9d0fb55..dcd57e0 100644
--- a/src/Impostor.Server/Net/State/Game.Data.cs
+++ b/src/Impostor.Server/Net/State/Game.Data.cs
@@ -499,6 +499,14 @@ namespace Impostor.Server.Net.State
             playerInfo.ClientId = sender.Client.Id;
             playerInfo.PlayerId = GameNet.GameData.GetNextAvailablePlayerId();
 
+            // If player played a previous game, restore their color
+            var prevColor = sender.Client.PreviousColor;
+            if (prevColor.HasValue)
+            {
+                _logger.LogTrace("Color restored to {Color}", prevColor.Value);
+                playerInfo.CurrentOutfit.Color = prevColor.Value;
+            }
+
             if (!AddNetObject(playerInfo))
             {
                 _logger.LogError("Couldn't spawn PlayerInfo for {Name} ({ClientId})", sender.Client.Name, sender.Client.Id);
-- 
2.39.5