From ffa37d73a8cb52e330742689ca3c4072af21f4f4 Mon Sep 17 00:00:00 2001
From: SiuTung08 <siutung08@gmail.com>
Date: Tue, 12 Nov 2024 20:20:06 +0000
Subject: [PATCH] making HeaderWriter internal again

---
 .../{ => Internal}/HeaderWriter.cs               | 16 ++++++----------
 src/NATS.Client.Core/NatsHeaders.cs              | 14 ++++++++++----
 .../NATS.Client.CoreUnit.Tests/NatsHeaderTest.cs |  4 ++--
 3 files changed, 18 insertions(+), 16 deletions(-)
 rename src/NATS.Client.Core/{ => Internal}/HeaderWriter.cs (88%)

diff --git a/src/NATS.Client.Core/HeaderWriter.cs b/src/NATS.Client.Core/Internal/HeaderWriter.cs
similarity index 88%
rename from src/NATS.Client.Core/HeaderWriter.cs
rename to src/NATS.Client.Core/Internal/HeaderWriter.cs
index 7cb1b597..286fd999 100644
--- a/src/NATS.Client.Core/HeaderWriter.cs
+++ b/src/NATS.Client.Core/Internal/HeaderWriter.cs
@@ -1,11 +1,10 @@
 using System.Buffers;
 using System.Text;
 using NATS.Client.Core.Commands;
-using NATS.Client.Core.Internal;
 
-namespace NATS.Client.Core;
+namespace NATS.Client.Core.Internal;
 
-public class HeaderWriter
+internal class HeaderWriter
 {
     private const byte ByteCr = (byte)'\r';
     private const byte ByteLf = (byte)'\n';
@@ -21,7 +20,7 @@ public class HeaderWriter
 
     private static ReadOnlySpan<byte> ColonSpace => new[] { ByteColon, ByteSpace };
 
-    public long GetBytesLength(NatsHeaders headers)
+    internal static long GetBytesLength(NatsHeaders headers, Encoding encoding)
     {
         var len = CommandConstants.NatsHeaders10NewLine.Length;
         foreach (var kv in headers)
@@ -31,14 +30,14 @@ public long GetBytesLength(NatsHeaders headers)
                 if (value != null)
                 {
                     // key length
-                    var keyLength = _encoding.GetByteCount(kv.Key);
+                    var keyLength = encoding.GetByteCount(kv.Key);
                     len += keyLength;
 
                     // colon space length
                     len += ColonSpace.Length;
 
                     // value length
-                    var valueLength = _encoding.GetByteCount(value);
+                    var valueLength = encoding.GetByteCount(value);
                     len += valueLength;
 
                     // CrLf length
@@ -68,10 +67,7 @@ internal long Write(IBufferWriter<byte> bufferWriter, NatsHeaders headers)
                     var keySpan = bufferWriter.GetSpan(keyLength);
                     _encoding.GetBytes(kv.Key, keySpan);
                     if (!ValidateKey(keySpan.Slice(0, keyLength)))
-                    {
-                        throw new NatsException(
-                            $"Invalid header key '{kv.Key}': contains colon, space, or other non-printable ASCII characters");
-                    }
+                        throw new NatsException($"Invalid header key '{kv.Key}': contains colon, space, or other non-printable ASCII characters");
 
                     bufferWriter.Advance(keyLength);
                     len += keyLength;
diff --git a/src/NATS.Client.Core/NatsHeaders.cs b/src/NATS.Client.Core/NatsHeaders.cs
index 3dcf429b..d8d447ed 100644
--- a/src/NATS.Client.Core/NatsHeaders.cs
+++ b/src/NATS.Client.Core/NatsHeaders.cs
@@ -1,7 +1,9 @@
 using System.Collections;
 using System.Diagnostics;
 using System.Diagnostics.CodeAnalysis;
+using System.Text;
 using Microsoft.Extensions.Primitives;
+using NATS.Client.Core.Internal;
 
 namespace NATS.Client.Core;
 
@@ -297,11 +299,15 @@ public void CopyTo(KeyValuePair<string, StringValues>[] array, int arrayIndex)
     /// <summary>
     /// Returns the bytes length of the header
     /// </summary>
-    /// <param name="headerWriter">an instance of headerWriter</param>
-    /// <returns>bytes length of th header</returns>
-    public long GetBytesLength(HeaderWriter headerWriter)
+    /// <param name="encoding">Encoding used.  Default to utf8 if not provided</param>
+    /// <returns>Bytes length of the header</returns>
+    public long GetBytesLength(Encoding? encoding = null)
     {
-        return headerWriter.GetBytesLength(this);
+        // if null set to utf-8
+        encoding = encoding ?? Encoding.UTF8;
+
+        var len = HeaderWriter.GetBytesLength(this, encoding);
+        return len;
     }
 
     /// <summary>
diff --git a/tests/NATS.Client.CoreUnit.Tests/NatsHeaderTest.cs b/tests/NATS.Client.CoreUnit.Tests/NatsHeaderTest.cs
index a918a79d..f938451c 100644
--- a/tests/NATS.Client.CoreUnit.Tests/NatsHeaderTest.cs
+++ b/tests/NATS.Client.CoreUnit.Tests/NatsHeaderTest.cs
@@ -206,8 +206,8 @@ public void GetBytesLengthTest()
             ["a-long-header-key"] = "value",
             ["key"] = "a-long-header-value",
         };
-        var writer = new HeaderWriter(Encoding.UTF8);
-        var bytesLength = headers.GetBytesLength(writer);
+        var encoding = Encoding.UTF8;
+        var bytesLength = headers.GetBytesLength(encoding);
 
         var text = "NATS/1.0\r\nk1: v1\r\nk2: v2-0\r\nk2: v2-1\r\na-long-header-key: value\r\nkey: a-long-header-value\r\n\r\n";
         var expected = new ReadOnlySequence<byte>(Encoding.UTF8.GetBytes(text));