KIKIMR-20533 TraceId parser unittests (#617)

4 files changed, 120 insertions, 1 deletions

diff --git a/ydb/library/actors/wilson/ut/ya.make b/ydb/library/actors/wilson/ut/ya.make
new file mode 100644
index 0000000000..210e615d7f
--- /dev/null
+++ b/ydb/library/actors/wilson/ut/ya.make
@@ -0,0 +1,11 @@
+UNITTEST_FOR(ydb/library/actors/wilson)
+
+OWNER(
+    g:kikimr
+)
+
+SRCS(
+    wilson_trace_ut.cpp
+)
+
+END()
diff --git a/ydb/library/actors/wilson/wilson_trace.cpp b/ydb/library/actors/wilson/wilson_trace.cpp
index babcde2e95..ec244856bc 100644
--- a/ydb/library/actors/wilson/wilson_trace.cpp
+++ b/ydb/library/actors/wilson/wilson_trace.cpp
@@ -26,6 +26,9 @@ namespace NWilson {
             return {};
         }
 
+        // Specification only allows lower case letters, but we want to allow upper case
+        // letters for convenience
+        // https://w3c.github.io/trace-context/#traceparent-header-field-values
         auto isHex = [](char c) {
             return (c >= '0' && c <= '9') || (c >= 'A' && c <= 'F') || (c >= 'a' && c <= 'f');
         };
@@ -37,7 +40,7 @@ namespace NWilson {
             return {};
         }
 
-        if (header[traceIdStart - 1] != '-' || header[parentSpanIdStart - 1] != '-' || header[traceIdChars - 1] != '-') {
+        if (header[traceIdStart - 1] != '-' || header[parentSpanIdStart - 1] != '-' || header[traceFlagsStart - 1] != '-') {
             return {};
         }
 
diff --git a/ydb/library/actors/wilson/wilson_trace_ut.cpp b/ydb/library/actors/wilson/wilson_trace_ut.cpp
new file mode 100644
index 0000000000..c35bc3a55d
--- /dev/null
+++ b/ydb/library/actors/wilson/wilson_trace_ut.cpp
@@ -0,0 +1,101 @@
+#include "wilson_trace.h"
+
+#include <library/cpp/testing/unittest/registar.h>
+#include <util/system/byteorder.h>
+
+using namespace NWilson;
+
+using TTrace = std::array<ui64, 2>;
+
+TTrace GetTrace(const TTraceId& traceId) {
+    static_assert(TTraceId::GetTraceIdSize() == sizeof(TTrace));
+    return *reinterpret_cast<const TTrace*>(traceId.GetTraceIdPtr());
+}
+
+ui64 GetSpan(const TTraceId& traceId) {
+    static_assert(TTraceId::GetSpanIdSize() == sizeof(ui64));
+    return *reinterpret_cast<const ui64*>(traceId.GetSpanIdPtr());
+}
+
+bool CheckAllZero(const TTraceId& traceId) {
+    auto trace = GetTrace(traceId);
+    return trace[0] == 0 && trace[1] == 0 && GetSpan(traceId) == 0;
+}
+
+template<class T>
+inline T HostToBig(T val) noexcept {
+    return LittleToBig(HostToLittle(val));
+}
+
+Y_UNIT_TEST_SUITE(TTraceId) {
+    Y_UNIT_TEST(OpenTelemetryHeaderParser) {
+        const auto incorrectHeaders = std::array {
+            "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-011", // Wrong length
+            "0-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01",
+            "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-0",
+
+            "01-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01", // Wrong version
+
+            "0r-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01", // Wrong charset
+            "00-0af7651916xd43dd8448eb211c80319c-b7ad6b7169203331-01",
+            "00-0af7651916cd43dd8448eb211c80319c-b7ad6y7169203331-01",
+            "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-t1",
+
+            "00=0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01", // Wrong dash
+            "00-0af7651916cd43dd8448eb211c80319c+b7ad6b7169203331-01",
+            "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331_01",
+
+            "00-0af7651916cd43dd8448eb211c80319c-0000000000000000-01", // One of the ids is zero
+            "00-00000000000000000000000000000000-b7ad6b7169203331-01",
+
+            "00-0af7651916cd43dd8448eb211c8031-baa7ad6b7169203331-01", // Random errors
+            "abacaba",
+            "ababababab23131231231bababababababcdcdc123123123acdadab",
+        };
+
+        for (auto header : incorrectHeaders) {
+            auto traceId = TTraceId::FromTraceparentHeader(header);
+            UNIT_ASSERT(CheckAllZero(traceId) && !traceId);
+        }
+
+        const auto correctHeaders = std::array {
+            std::tuple {
+                "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01",
+                TTrace {0x0af7651916cd43dd, 0x8448eb211c80319c},
+                0xb7ad6b7169203331,
+            },
+            std::tuple {
+                "00-0AF7651916CD43DD8448EB211C80319C-b7Ad6B7169203331-01",
+                TTrace {0x0af7651916cd43dd, 0x8448eb211c80319c},
+                0xb7ad6b7169203331,
+            },
+            std::tuple {
+                "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-03",
+                TTrace {0x0af7651916cd43dd, 0x8448eb211c80319c},
+                0xb7ad6b7169203331,
+            },
+            std::tuple {
+                "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-00",
+                TTrace {0x0af7651916cd43dd, 0x8448eb211c80319c},
+                0xb7ad6b7169203331,
+            },
+            std::tuple {
+                "00-000101020305080d1522375990e97962-db3d18556dc22ff1-01",
+                TTrace {0x000101020305080d, 0x1522375990e97962},
+                0xdb3d18556dc22ff1,
+            },
+        };
+
+        for (auto [header, expectedTrace, expectedSpan] : correctHeaders) {
+            auto traceId = TTraceId::FromTraceparentHeader(header);
+            auto realTrace = GetTrace(traceId);
+            auto realSpan = GetSpan(traceId);
+            for (auto& part : realTrace) {
+                part = HostToBig(part);
+            }
+            realSpan = HostToBig(realSpan);
+            UNIT_ASSERT_EQUAL(realTrace, expectedTrace);
+            UNIT_ASSERT_EQUAL(realSpan, expectedSpan);
+        }
+    }
+}
diff --git a/ydb/library/actors/wilson/ya.make b/ydb/library/actors/wilson/ya.make
index 2e6552a568..9786754fae 100644
--- a/ydb/library/actors/wilson/ya.make
+++ b/ydb/library/actors/wilson/ya.make
@@ -19,3 +19,7 @@ END()
 RECURSE(
     protos
 )
+
+RECURSE_FOR_TESTS(
+    ut
+)