From b85e45b2b2791840d5c97102823aab106cd48606 Mon Sep 17 00:00:00 2001
From: Pavel Djundik <xPaw@users.noreply.github.com>
Date: Thu, 27 Jun 2024 22:27:51 +0300
Subject: [PATCH] Implement StringPool for binary kv1 used by Steam

---
 .../Abstraction/IParsingVisitationListener.cs     |  2 ++
 .../Deserialization/KVObjectBuilder.cs            |  2 ++
 .../Deserialization/KeyValues1/KV1BinaryReader.cs | 15 ++++++++++++---
 ValveKeyValue/ValveKeyValue/KVSerializer.cs       |  5 ++++-
 .../ValveKeyValue/KVSerializerOptions.cs          |  5 +++++
 5 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/ValveKeyValue/ValveKeyValue/Abstraction/IParsingVisitationListener.cs b/ValveKeyValue/ValveKeyValue/Abstraction/IParsingVisitationListener.cs
index ceb0c2fe..ad711efc 100644
--- a/ValveKeyValue/ValveKeyValue/Abstraction/IParsingVisitationListener.cs
+++ b/ValveKeyValue/ValveKeyValue/Abstraction/IParsingVisitationListener.cs
@@ -2,6 +2,8 @@ namespace ValveKeyValue.Abstraction
 {
     interface IParsingVisitationListener : IVisitationListener
     {
+        public string[] StringPool { get; set; }
+
         void DiscardCurrentObject();
 
         IParsingVisitationListener GetMergeListener();
diff --git a/ValveKeyValue/ValveKeyValue/Deserialization/KVObjectBuilder.cs b/ValveKeyValue/ValveKeyValue/Deserialization/KVObjectBuilder.cs
index 36c82a2a..21203220 100644
--- a/ValveKeyValue/ValveKeyValue/Deserialization/KVObjectBuilder.cs
+++ b/ValveKeyValue/ValveKeyValue/Deserialization/KVObjectBuilder.cs
@@ -6,6 +6,8 @@ class KVObjectBuilder : IParsingVisitationListener
     {
         readonly IList<KVObjectBuilder> associatedBuilders = new List<KVObjectBuilder>();
 
+        public string[] StringPool { get; set; }
+
         public KVObject GetObject()
         {
             if (stateStack.Count != 1)
diff --git a/ValveKeyValue/ValveKeyValue/Deserialization/KeyValues1/KV1BinaryReader.cs b/ValveKeyValue/ValveKeyValue/Deserialization/KeyValues1/KV1BinaryReader.cs
index 7c4be940..474cff83 100644
--- a/ValveKeyValue/ValveKeyValue/Deserialization/KeyValues1/KV1BinaryReader.cs
+++ b/ValveKeyValue/ValveKeyValue/Deserialization/KeyValues1/KV1BinaryReader.cs
@@ -8,7 +8,7 @@ class KV1BinaryReader : IVisitingReader
     {
         public const int BinaryMagicHeader = 0x564B4256; // VBKV
 
-        public KV1BinaryReader(Stream stream, IVisitationListener listener)
+        public KV1BinaryReader(Stream stream, IParsingVisitationListener listener)
         {
             Require.NotNull(stream, nameof(stream));
             Require.NotNull(listener, nameof(listener));
@@ -25,7 +25,7 @@ public KV1BinaryReader(Stream stream, IVisitationListener listener)
 
         readonly Stream stream;
         readonly BinaryReader reader;
-        readonly IVisitationListener listener;
+        readonly IParsingVisitationListener listener;
         bool disposed;
         KV1BinaryNodeType endMarker = KV1BinaryNodeType.End;
 
@@ -76,9 +76,18 @@ void ReadObjectCore()
 
         void ReadValue(KV1BinaryNodeType type)
         {
-            var name = Encoding.UTF8.GetString(ReadNullTerminatedBytes());
+            string name;
             KVValue value;
 
+            if (listener.StringPool != null)
+            {
+                name = listener.StringPool[reader.ReadInt32()];
+            }
+            else
+            {
+                name = Encoding.UTF8.GetString(ReadNullTerminatedBytes());
+            }
+
             switch (type)
             {
                 case KV1BinaryNodeType.ChildObject:
diff --git a/ValveKeyValue/ValveKeyValue/KVSerializer.cs b/ValveKeyValue/ValveKeyValue/KVSerializer.cs
index 965765d8..01337067 100644
--- a/ValveKeyValue/ValveKeyValue/KVSerializer.cs
+++ b/ValveKeyValue/ValveKeyValue/KVSerializer.cs
@@ -34,7 +34,10 @@ public static KVSerializer Create(KVSerializationFormat format)
         public KVDocument Deserialize(Stream stream, KVSerializerOptions options = null)
         {
             Require.NotNull(stream, nameof(stream));
-            var builder = new KVObjectBuilder();
+            var builder = new KVObjectBuilder
+            {
+                StringPool = options?.StringPool
+            };
 
             using (var reader = MakeReader(stream, builder, options ?? KVSerializerOptions.DefaultOptions))
             {
diff --git a/ValveKeyValue/ValveKeyValue/KVSerializerOptions.cs b/ValveKeyValue/ValveKeyValue/KVSerializerOptions.cs
index d3b3c8d6..d2f1364a 100644
--- a/ValveKeyValue/ValveKeyValue/KVSerializerOptions.cs
+++ b/ValveKeyValue/ValveKeyValue/KVSerializerOptions.cs
@@ -32,6 +32,11 @@ public sealed class KVSerializerOptions
         /// </summary>
         public static KVSerializerOptions DefaultOptions => new();
 
+        /// <summary>
+        /// If KV1 is serialized using a string pool for the keys, provide the string pool here.
+        /// </summary>
+        public string[] StringPool;
+
         static IEnumerable<string> GetDefaultConditions()
         {
             // TODO: In the future we will want to skip this for consoles and mobile devices?