From 57fc4d2bfd54770571337bcb6bcb982ba086ae81 Mon Sep 17 00:00:00 2001
From: Lucas Pedroza <lucas.pedroza@fauna.com>
Date: Tue, 22 Oct 2024 15:38:14 -0700
Subject: [PATCH] fix: Handle unmaterialize set endings in UTF8FaunaReader

---
 .../Serialization/Utf8FaunaReader.Tests.cs    | 26 +++++++++++++++++++
 Fauna/Serialization/Utf8FaunaReader.cs        | 20 +++++++++++++-
 2 files changed, 45 insertions(+), 1 deletion(-)

diff --git a/Fauna.Test/Serialization/Utf8FaunaReader.Tests.cs b/Fauna.Test/Serialization/Utf8FaunaReader.Tests.cs
index 0113be36..3011a15c 100644
--- a/Fauna.Test/Serialization/Utf8FaunaReader.Tests.cs
+++ b/Fauna.Test/Serialization/Utf8FaunaReader.Tests.cs
@@ -593,4 +593,30 @@ public void SkipNestedArrays()
         Assert.AreEqual(TokenType.FieldName, reader.CurrentTokenType);
         Assert.AreEqual("k2", reader.GetString());
     }
+
+    [Test]
+    public void HandlesUnmaterializedSets()
+    {
+        const string test =
+            @"{""products"":{""@set"":""hdWCxmdQcm9kdWN0gcqEamJ5Q2F0ZWdvcnmBxmJ2MPT2gc1KBbkdRAugBgAEAvbBghpnGCOgGjW0K0AQ""},""name"":""foo""}";
+        var reader = new Utf8FaunaReader(test);
+        reader.Read(); // {
+        Assert.AreEqual(TokenType.StartObject, reader.CurrentTokenType);
+        reader.Read(); // "products"
+        Assert.AreEqual(TokenType.FieldName, reader.CurrentTokenType);
+        Assert.AreEqual("products", reader.GetString());
+        reader.Read(); // { "@set":
+        Assert.AreEqual(TokenType.StartPage, reader.CurrentTokenType);
+        reader.Read(); // "hdWCxmdQcm9kdWN0gcqEamJ5Q2F0ZWdvcnmBxmJ2MPT2gc1KBbkdRAugBgAEAvbBghpnGCOgGjW0K0AQ"
+        Assert.AreEqual("hdWCxmdQcm9kdWN0gcqEamJ5Q2F0ZWdvcnmBxmJ2MPT2gc1KBbkdRAugBgAEAvbBghpnGCOgGjW0K0AQ", reader.GetString());
+        reader.Read(); // "}"
+        Assert.AreEqual(TokenType.EndPage, reader.CurrentTokenType);
+        reader.Read(); // "name"
+        Assert.AreEqual(TokenType.FieldName, reader.CurrentTokenType);
+        Assert.AreEqual("name", reader.GetString());
+        reader.Read(); // "foo"
+        Assert.AreEqual("foo", reader.GetString());
+        reader.Read(); // "}"
+        Assert.AreEqual(TokenType.EndObject, reader.CurrentTokenType);
+    }
 }
diff --git a/Fauna/Serialization/Utf8FaunaReader.cs b/Fauna/Serialization/Utf8FaunaReader.cs
index 39f998b2..fdcde5ea 100644
--- a/Fauna/Serialization/Utf8FaunaReader.cs
+++ b/Fauna/Serialization/Utf8FaunaReader.cs
@@ -16,6 +16,7 @@ public ref struct Utf8FaunaReader
     private Utf8JsonReader _json;
     private readonly Stack<object> _tokenStack = new();
     private TokenType? _bufferedTokenType = null;
+    private object? _bufferedTokenValue = null;
 
     private readonly HashSet<TokenType> _closers = new()
     {
@@ -37,6 +38,7 @@ private enum TokenTypeInternal
     {
         /// <summary>The token type is the start of an escaped Fauna object.</summary>
         StartEscapedObject,
+        StartPageUnmaterialized,
     }
 
     /// <summary>
@@ -106,6 +108,9 @@ public bool Read()
             return true;
         }
 
+        // At this point we're passed a buffered token type so this read should unset this no matter what
+        _bufferedTokenValue = null;
+
         if (!Advance())
         {
             return false;
@@ -188,6 +193,7 @@ public bool Read()
 
         try
         {
+            if (_bufferedTokenValue != null) return (string)_bufferedTokenValue;
             return _json.GetString();
         }
         catch (Exception e)
@@ -598,7 +604,16 @@ private void HandleStartObject()
                     case "@set":
                         AdvanceTrue();
                         CurrentTokenType = TokenType.StartPage;
-                        _tokenStack.Push(TokenType.StartPage);
+                        if (_json.TokenType == JsonTokenType.String)
+                        {
+                            _bufferedTokenValue = _json.GetString();
+                            _bufferedTokenType = TokenType.String;
+                            _tokenStack.Push(TokenTypeInternal.StartPageUnmaterialized);
+                        }
+                        else
+                        {
+                            _tokenStack.Push(TokenType.StartPage);
+                        }
                         break;
                     case "@time":
                         HandleTaggedString(TokenType.Time);
@@ -636,6 +651,9 @@ private void HandleEndObject()
                 CurrentTokenType = TokenType.EndPage;
                 AdvanceTrue();
                 break;
+            case TokenTypeInternal.StartPageUnmaterialized:
+                CurrentTokenType = TokenType.EndPage;
+                break;
             case TokenType.StartRef:
                 CurrentTokenType = TokenType.EndRef;
                 AdvanceTrue();