fix(TomlTable): better key name checks (#33)

SamboyCoding · Jul 16, 2023 · 8dc1129 · 8dc1129
1 parent 4729d95
commit 8dc1129
Show file tree

Hide file tree

Showing 4 changed files with 93 additions and 42 deletions.
diff --git a/Tomlet.Tests/DictionaryTests.cs b/Tomlet.Tests/DictionaryTests.cs
@@ -27,4 +27,30 @@ public void DictionariesAsFieldsWork()
         //Just make sure this doesn't throw
         var serialized = TomletMain.TomlStringFrom(obj); 
     }
+
+    [Fact]
+    public void DictionaryKeysShouldBeProperlyEscaped()
+    {
+        var dictionary = new Dictionary<string, string>
+        {
+            {"normal-key", "normal-key"},
+            {"normal_key", "normal_key"},
+            {"normalkey", "normalkey"},
+            {"key with space", "\"key with spaces\""},
+            {"key!with{}(*%&)random[other+symbols", "\"key!with{}(*%&)random[other+symbols\""},
+            {"key/with/slashes", "\"key/with/slashes\""},
+            {"Nam\\e", "\"Nam\\\\e\""}
+        };
+
+        var obj = new ClassWithDictionary
+        {
+            GenericDictionary = dictionary
+        };
+
+        var serialized = TomletMain.TomlStringFrom(obj);
+        foreach(var (_, expectedValue) in dictionary)
+        {
+            Assert.Contains(expectedValue, serialized);
+        }
+    }
 }
diff --git a/Tomlet.Tests/TableTests.cs b/Tomlet.Tests/TableTests.cs
@@ -46,6 +46,26 @@ public void TablesCanHaveQuotedKeyNames()
             var tomlString = document.SerializedValue.Trim();
             Assert.Equal(inputString, tomlString);
         }
+
+         [Theory]
+         [InlineData("normal-key", "normal-key")]
+         [InlineData("normal_key", "normal_key")]
+         [InlineData("normalkey", "normalkey")]
+         [InlineData("\"key with spaces\"", "\"key with spaces\"")]
+         [InlineData("key!with{}(*%&)random[other+symbols", "\"key!with{}(*%&)random[other+symbols\"")]
+         [InlineData("key/with/slashes", "\"key/with/slashes\"")]
+         [InlineData("Nam\\e", "\"Nam\\\\e\"")]
+         public void KeysShouldBeSerializedCorrectly(string inputKey, string expectedKey)
+         {
+             var inputString = $"""
+             ["Test Table"]
+             {inputKey} = "value"
+             """;
+
+             var document = GetDocument(inputString);
+             var serializedDocument = document.SerializedValue.Trim();
+             Assert.Contains(expectedKey, serializedDocument);
+         }
 
         [Fact]
         public void TablesCanHaveQuotedDottedNames()
@@ -58,23 +78,23 @@ public void TablesCanHaveQuotedDottedNames()
         }
 
         [Fact]
-        public void TableIteratorEqualsEntries()
-        {
-            var document = GetDocument(TestResources.TableIteratorInput);
-
-            Assert.Single(document.Entries);
-
-            Assert.NotNull(document.GetSubTable("table-1"));
-
-            var table1 = document.GetSubTable("table-1");
-
-            Assert.Equal(4, table1.Entries.Count);
-
-            foreach (var entry in table1)
-            {
-                Assert.True(table1.Entries.ContainsKey(entry.Key));
-                Assert.Equal(table1.Entries[entry.Key], entry.Value);
-            }
+        public void TableIteratorEqualsEntries()
+        {
+            var document = GetDocument(TestResources.TableIteratorInput);
+
+            Assert.Single(document.Entries);
+
+            Assert.NotNull(document.GetSubTable("table-1"));
+
+            var table1 = document.GetSubTable("table-1");
+
+            Assert.Equal(4, table1.Entries.Count);
+
+            foreach (var entry in table1)
+            {
+                Assert.True(table1.Entries.ContainsKey(entry.Key));
+                Assert.Equal(table1.Entries[entry.Key], entry.Value);
+            }
         }
     }
 }
diff --git a/Tomlet.Tests/TestModelClasses/ClassWithDictionary.cs b/Tomlet.Tests/TestModelClasses/ClassWithDictionary.cs
@@ -5,4 +5,5 @@ namespace Tomlet.Tests.TestModelClasses;
 public class ClassWithDictionary
 {
     public Dictionary<string, Subname> name;
+    public Dictionary<string, string> GenericDictionary;
 }
diff --git a/Tomlet/Models/TomlTable.cs b/Tomlet/Models/TomlTable.cs
@@ -1,5 +1,5 @@
 using System;
-using System.Collections;
+using System.Collections;
 using System.Collections.Generic;
 using System.Diagnostics.CodeAnalysis;
 using System.Linq;
@@ -37,7 +37,7 @@ public override string SerializedValue
 
                 var builder = new StringBuilder("{ ");
 
-                builder.Append(string.Join(", ", Entries.Select(o => o.Key + " = " + o.Value.SerializedValue).ToArray()));
+                builder.Append(string.Join(", ", Entries.Select(o => EscapeKeyIfNeeded(o.Key) + " = " + o.Value.SerializedValue).ToArray()));
 
                 builder.Append(" }");
 
@@ -137,10 +137,8 @@ private void WriteValueToStringBuilder(string? keyName, string subKey, StringBui
             builder.Append('\n');
         }
 
-        private string EscapeKeyIfNeeded(string key)
+        private static string EscapeKeyIfNeeded(string key)
         {
-            var didQuote = false;
-
             if (key.StartsWith("\"") && key.EndsWith("\"") && key.Count(c => c == '"') == 2)
                 //Already double quoted
                 return key;
@@ -149,18 +147,24 @@ private string EscapeKeyIfNeeded(string key)
                 //Already single quoted
                 return key;
 
-            if (key.Contains("\"") || key.Contains("'"))
+            if (IsValidKey(key))
+                return key;
+
+            key = TomlUtils.EscapeStringValue(key);
+            return TomlUtils.AddCorrectQuotes(key);
+        }
+
+        private static bool IsValidKey(string key)
+        {
+            foreach (var c in key)
             {
-                key = TomlUtils.AddCorrectQuotes(key);
-                didQuote = true;
+                if (!char.IsLetterOrDigit(c) && c != '_' && c != '-')
+                {
+                    return false;
+                }
             }
 
-            var escaped = TomlUtils.EscapeStringValue(key);
-
-            if (escaped.Contains(" ") || escaped.Contains("\\") && !didQuote)
-                escaped = TomlUtils.AddCorrectQuotes(escaped);
-
-            return escaped;
+            return true;
         }
 
         internal void ParserPutValue(string key, TomlValue value, int lineNumber)
@@ -455,16 +459,16 @@ public TomlTable GetSubTable(string key)
                 throw new TomlTypeMismatchException(typeof(TomlTable), value.GetType(), typeof(TomlTable));
 
             return tbl;
-        }
-
-        IEnumerator IEnumerable.GetEnumerator()
-        {
-            return GetEnumerator();
-        }
-
-        public IEnumerator<KeyValuePair<string, TomlValue>> GetEnumerator()
-        {
-            return Entries.GetEnumerator();
-        }
+        }
+
+        IEnumerator IEnumerable.GetEnumerator()
+        {
+            return GetEnumerator();
+        }
+
+        public IEnumerator<KeyValuePair<string, TomlValue>> GetEnumerator()
+        {
+            return Entries.GetEnumerator();
+        }
     }
 }