open Header.Writer class and create GetBytesLength method

nats-io · Nov 9, 2024 · df43364 · df43364
1 parent 7c00b41
commit df43364
Show file tree

Hide file tree

Showing 2 changed files with 53 additions and 11 deletions.
diff --git a/...NATS.Client.Core/Internal/HeaderWriter.cs → src/NATS.Client.Core/HeaderWriter.cs b/...NATS.Client.Core/Internal/HeaderWriter.cs → src/NATS.Client.Core/HeaderWriter.cs
@@ -1,11 +1,11 @@
 using System.Buffers;
-using System.IO.Pipelines;
 using System.Text;
 using NATS.Client.Core.Commands;
+using NATS.Client.Core.Internal;
 
-namespace NATS.Client.Core.Internal;
+namespace NATS.Client.Core;
 
-internal class HeaderWriter
+public class HeaderWriter
 {
     private const byte ByteCr = (byte)'\r';
     private const byte ByteLf = (byte)'\n';
@@ -21,6 +21,37 @@ internal class HeaderWriter
 
     private static ReadOnlySpan<byte> ColonSpace => new[] { ByteColon, ByteSpace };
 
+    public long GetBytesLength(NatsHeaders headers)
+    {
+        var len = CommandConstants.NatsHeaders10NewLine.Length;
+        foreach (var kv in headers)
+        {
+            foreach (var value in kv.Value)
+            {
+                if (value != null)
+                {
+                    // key length
+                    var keyLength = _encoding.GetByteCount(kv.Key);
+                    len += keyLength;
+
+                    // colon space length
+                    len += ColonSpace.Length;
+
+                    // value length
+                    var valueLength = _encoding.GetByteCount(value);
+                    len += valueLength;
+
+                    // CrLf length
+                    len += CrLf.Length;
+                }
+            }
+        }
+
+        // CrLf length for empty headers
+        len += CrLf.Length;
+        return len;
+    }
+
     internal long Write(IBufferWriter<byte> bufferWriter, NatsHeaders headers)
     {
         bufferWriter.WriteSpan(CommandConstants.NatsHeaders10NewLine);
@@ -53,9 +84,7 @@ internal long Write(IBufferWriter<byte> bufferWriter, NatsHeaders headers)
                     var valueSpan = bufferWriter.GetSpan(valueLength);
                     _encoding.GetBytes(value, valueSpan);
                     if (!ValidateValue(valueSpan.Slice(0, valueLength)))
-                    {
                         throw new NatsException($"Invalid header value for key '{kv.Key}': contains CRLF");
-                    }
 
                     bufferWriter.Advance(valueLength);
                     len += valueLength;
@@ -81,9 +110,7 @@ private static bool ValidateKey(ReadOnlySpan<byte> span)
         foreach (var b in span)
         {
             if (b <= ByteSpace || b == ByteColon || b >= ByteDel)
-            {
                 return false;
-            }
         }
 
         return true;
@@ -96,15 +123,11 @@ private static bool ValidateValue(ReadOnlySpan<byte> span)
         {
             var pos = span.IndexOf(ByteCr);
             if (pos == -1 || pos == span.Length - 1)
-            {
                 return true;
-            }
 
             pos += 1;
             if (span[pos] == ByteLf)
-            {
                 return false;
-            }
 
             span = span[pos..];
         }

diff --git a/tests/NATS.Client.CoreUnit.Tests/NatsHeaderTest.cs b/tests/NATS.Client.CoreUnit.Tests/NatsHeaderTest.cs
@@ -195,4 +195,23 @@ public void GetLastValueTests(string text, bool expectedResult, string? expected
         Assert.Equal(result, expectedResult);
         Assert.Equal(lastValue, expectedLastValue);
     }
+
+    [Fact]
+    public void GetBytesLengthTest()
+    {
+        var headers = new NatsHeaders
+        {
+            ["k1"] = "v1",
+            ["k2"] = new[] { "v2-0", "v2-1" },
+            ["a-long-header-key"] = "value",
+            ["key"] = "a-long-header-value",
+        };
+        var writer = new HeaderWriter(Encoding.UTF8);
+        var bytesLength = writer.GetBytesLength(headers);
+
+        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));
+
+        Assert.Equal(expected.Length, bytesLength);
+    }
 }