DigestMiddleware: Move code to use IncrementalHash

Refresh.Common/Extensions/IncrementalHashExtensions.cs

@@ -0,0 +1,9 @@
+using System.Security.Cryptography;
+using System.Text;
+
+namespace Refresh.Common.Extensions;
+
+public static class IncrementalHashExtensions
+{
+    public static void WriteString(this IncrementalHash hash, string str) => hash.AppendData(Encoding.UTF8.GetBytes(str));
+}

Refresh.GameServer/Middlewares/DigestMiddleware.cs

@@ -25,24 +25,24 @@ public record PspVersionInfo(short ExeVersion, short DataVersion) {}
     public static string CalculateDigest(
         string digest, 
         string route,
-        Stream body,
+        ReadOnlySpan<byte> body,
         string auth, 
         PspVersionInfo? pspVersionInfo, 
         bool isUpload, 
         bool hmacDigest)
     {
-        // Init a MemoryStream with the known final capacity capacity
-        using MemoryStream ms = new((int)(auth.Length + route.Length + digest.Length + (isUpload ? 0 : body.Length)) + (pspVersionInfo == null ? 0 : 4));
+        IncrementalHash hash = hmacDigest
+            ? IncrementalHash.CreateHMAC(HashAlgorithmName.SHA1, Encoding.UTF8.GetBytes(digest))
+            : IncrementalHash.CreateHash(HashAlgorithmName.SHA1);
 
         // If this is not an upload endpoint, then we need to copy the body of the request into the digest calculation
         if (!isUpload)
         {
-            body.CopyTo(ms);
-            body.Seek(0, SeekOrigin.Begin);
+            hash.AppendData(body);
         }
-        
-        ms.WriteString(auth);
-        ms.WriteString(route);
+
+        hash.WriteString(auth);
+        hash.WriteString(route);
         if (pspVersionInfo != null)
         {
             Span<byte> bytes = stackalloc byte[2];
@@ -51,27 +51,19 @@ public static string CalculateDigest(
             // If we are on a big endian system, we need to flip the bytes
             if(!BitConverter.IsLittleEndian)
                 bytes.Reverse();
-            ms.Write(bytes);
+            hash.AppendData(bytes);
 
             BitConverter.TryWriteBytes(bytes, pspVersionInfo.DataVersion);
             // If we are on a big endian system, we need to flip the bytes
             if(!BitConverter.IsLittleEndian)
                 bytes.Reverse();
-            ms.Write(bytes);
+            hash.AppendData(bytes);
         }
+
         if(!hmacDigest)
-            ms.WriteString(digest);
+            hash.WriteString(digest);
 
-        ms.Position = 0;
-
-        if (hmacDigest)
-        {
-            using HMACSHA1 hmac = new(Encoding.UTF8.GetBytes(digest));
-            return Convert.ToHexString(hmac.ComputeHash(ms)).ToLower();
-        }
-
-        using SHA1 sha = SHA1.Create();
-        return Convert.ToHexString(sha.ComputeHash(ms)).ToLower();
+        return Convert.ToHexString(hash.GetCurrentHash()).ToLower();
     }
 
     public void HandleRequest(ListenerContext context, Lazy<IDatabaseContext> database, Action next)
@@ -85,6 +77,9 @@ public void HandleRequest(ListenerContext context, Lazy<IDatabaseContext> databa
             return;
         }
 
+        // TODO: once we have PS4 support, check if the token is a PS4 token
+        bool isPs4 = context.RequestHeaders["User-Agent"]?.Contains("MM CHTTPClient LBP3 01.26") ?? false;
+
         PspVersionInfo? pspVersionInfo = null;
         // Try to acquire the exe and data version, this is only accounted for in the client digests, not the server digests
         if (short.TryParse(context.RequestHeaders["X-Exe-V"], out short exeVer) &&
@@ -107,50 +102,70 @@ public void HandleRequest(ListenerContext context, Lazy<IDatabaseContext> databa
 
         Token? token = gameDatabase.GetTokenFromTokenData(auth, TokenType.Game);
 
+        byte[] responseBody = context.ResponseStream.ToArray();
+        byte[] requestBody = context.InputStream.ToArray();
+
         // Make sure the digest calculation reads the whole response stream
         context.ResponseStream.Seek(0, SeekOrigin.Begin);
 
         // If the digest is already saved on the token, use the token's digest
         if (token is { Digest: not null })
         {
-            SetDigestResponse(context, CalculateDigest(token.Digest, route, context.ResponseStream, auth, null, isUpload, token.IsHmacDigest));
+            SetDigestResponse(context, CalculateDigest(token.Digest, route, responseBody, auth, null, isUpload, token.IsHmacDigest));
             return;
         }
 
-        (string digest, bool hmac)? foundDigest = this.FindBestKey(clientDigest, route, context.InputStream, auth, pspVersionInfo, isUpload, false) ??
-                                                  this.FindBestKey(clientDigest, route, context.InputStream, auth, pspVersionInfo, isUpload, true);
+        // If the client asks for a particular digest index, use that digest
+        if (int.TryParse(context.RequestHeaders["Refresh-Ps3-Digest-Index"], out int ps3DigestIndex) && 
+            int.TryParse(context.RequestHeaders["Refresh-Ps4-Digest-Index"], out int ps4DigestIndex))
+        {
+            string digest = isPs4
+                ? this._config.HmacDigestKeys[ps4DigestIndex]
+                : this._config.Sha1DigestKeys[ps3DigestIndex];
+
+            SetDigestResponse(context, CalculateDigest(digest, route, responseBody, auth, null, isUpload, isPs4));
+
+            if(token != null)
+                gameDatabase.SetTokenDigestInfo(token, digest, isPs4);
+
+            return;
+        }
+
+        (string digest, bool hmac)? foundDigest = this.FindBestKey(clientDigest, route, requestBody, auth, pspVersionInfo, isUpload, false) ??
+                                                  this.FindBestKey(clientDigest, route, requestBody, auth, pspVersionInfo, isUpload, true);
 
         if (foundDigest != null)
         {
             string digest = foundDigest.Value.digest;
             bool hmac = foundDigest.Value.hmac;
 
-            SetDigestResponse(context, CalculateDigest(digest, route, context.ResponseStream, auth, null, isUpload, hmac));
+            SetDigestResponse(context, CalculateDigest(digest, route, responseBody, auth, null, isUpload, hmac));
 
             if(token != null)
                 gameDatabase.SetTokenDigestInfo(token, digest, hmac); 
         }
         else
         {
             // If we were unable to find any digests, just use the first one specified as a backup
-            // TODO: once we have PS4 support, check if the token is a PS4 token
-            bool isPs4 = context.RequestHeaders["User-Agent"] == "MM CHTTPClient LBP3 01.26";
             string firstDigest = isPs4 ? this._config.HmacDigestKeys[0] : this._config.Sha1DigestKeys[0];
 
-            SetDigestResponse(context, CalculateDigest(firstDigest, route, context.ResponseStream, auth, null, isUpload, isPs4));
+            SetDigestResponse(context, CalculateDigest(firstDigest, route, responseBody, auth, null, isUpload, isPs4));
 
-            if(token != null)
+            // If theres no token, or the client didnt provide any client digest, lock the token into the found digest
+            // The second condition is to make sure that we dont lock in a digest for endpoints which send no client digest,
+            // but expect a server digest.
+            if(token != null && !string.IsNullOrEmpty(clientDigest))
                 gameDatabase.SetTokenDigestInfo(token, firstDigest, isPs4);
         }
     }
 
-    private (string digest, bool hmac)? FindBestKey(string clientDigest, string route, MemoryStream inputStream, string auth, PspVersionInfo? pspVersionInfo, bool isUpload, bool hmac)
+    private (string digest, bool hmac)? FindBestKey(string clientDigest, string route, ReadOnlySpan<byte> requestData, string auth, PspVersionInfo? pspVersionInfo, bool isUpload, bool hmac)
     {
         string[] keys = hmac ? this._config.HmacDigestKeys : this._config.Sha1DigestKeys;
 
         foreach (string digest in keys)
         {
-            string calculatedClientDigest = CalculateDigest(digest, route, inputStream, auth, pspVersionInfo, isUpload, hmac);
+            string calculatedClientDigest = CalculateDigest(digest, route, requestData, auth, pspVersionInfo, isUpload, hmac);
 
             // If the calculated client digest is invalid, then this isn't the digest the game is using, so check the next one
             if (calculatedClientDigest != clientDigest) 

Refresh.HttpsProxy/Config/ProxyConfig.cs

@@ -2,12 +2,12 @@ namespace Refresh.HttpsProxy.Config;
 
 public class ProxyConfig : Bunkum.Core.Configuration.Config
 {
-    public override int CurrentConfigVersion => 1;
+    public override int CurrentConfigVersion => 2;
     public override int Version { get; set; }
 
     protected override void Migrate(int oldVer, dynamic oldConfig) { }
 
     public string TargetServerUrl { get; set; } = "https://lbp.littlebigrefresh.com";
-    public string Ps3Digest { get; set; } = "CustomServerDigest";
-    public string Ps4Digest { get; set; } = "CustomServerDigest";
+    public int Ps3DigestIndex { get; set; } = 0;
+    public int Ps4DigestIndex { get; set; } = 0;
 }

Refresh.HttpsProxy/Middlewares/ProxyMiddleware.cs

@@ -63,6 +63,9 @@ public void HandleRequest(ListenerContext context, Lazy<IDatabaseContext> databa
             requestMessage.Headers.TryAddWithoutValidation(key, value);
         }
 
+        requestMessage.Headers.Add("Refresh-Ps3-Digest-Index", config.Ps3DigestIndex.ToString());
+        requestMessage.Headers.Add("Refresh-Ps4-Digest-Index", config.Ps4DigestIndex.ToString());
+
         // Send our HTTP request
         HttpResponseMessage response = client.Send(requestMessage);
 

Refresh.HttpsProxy/Program.cs

@@ -1,6 +1,4 @@
-// See https://aka.ms/new-console-template for more information
-
-using Bunkum.Core;
+using Bunkum.Core;
 using Bunkum.Core.Configuration;
 using Bunkum.Protocols.Http;
 using Bunkum.Protocols.Https;
@@ -30,13 +28,11 @@
 httpsServer.Initialize = s =>
 {
     s.AddMiddleware(new ProxyMiddleware(config));
-    s.AddMiddleware(new DigestMiddleware(config));
 };
 
 httpServer.Initialize = s =>
 {
     s.AddMiddleware(new ProxyMiddleware(config));
-    s.AddMiddleware(new DigestMiddleware(config));
 };
 
 // Start the server in multi-threaded mode, and let Bunkum manage the rest.