xray: Use correct types for segment document output (#10834)

Risk Level: Low
Testing: unit tests, manual
Docs Changes: N/A
Release Notes: N/A
Fixes: #10814

source/extensions/tracers/xray/daemon.proto

@@ -5,6 +5,7 @@ syntax = "proto3";
 package source.extensions.tracers.xray.daemon;
 
 import "validate/validate.proto";
+import "google/protobuf/struct.proto";
 
 // see https://docs.aws.amazon.com/xray/latest/devguide/xray-api-segmentdocuments.html
 message Segment {
@@ -14,12 +15,12 @@ message Segment {
   double start_time = 4 [(validate.rules).double = {gt: 0}];
   double end_time = 5 [(validate.rules).double = {gt: 0}];
   string parent_id = 6;
-  map<string, string> annotations = 7;
-  http_annotations http = 8;
+  http_annotations http = 7;
   message http_annotations {
-    map<string, string> request = 1;
-    map<string, string> response = 2;
+    google.protobuf.Struct request = 1;
+    google.protobuf.Struct response = 2;
   }
+  map<string, string> annotations = 8;
 }
 
 message Header {

source/extensions/tracers/xray/tracer.cc

@@ -11,6 +11,7 @@
 #include "common/common/fmt.h"
 #include "common/common/hex.h"
 #include "common/protobuf/message_validator_impl.h"
+#include "common/protobuf/utility.h"
 #include "common/runtime/runtime_impl.h"
 
 #include "source/extensions/tracers/xray/daemon.pb.validate.h"
@@ -76,17 +77,25 @@ void Span::finishSpan() {
   s.set_end_time(
       time_point_cast<SecondsWithFraction>(time_source_.systemTime()).time_since_epoch().count());
   s.set_parent_id(parentId());
-  using KeyValue = Protobuf::Map<std::string, std::string>::value_type;
-  for (const auto& item : custom_annotations_) {
-    s.mutable_annotations()->insert(KeyValue{item.first, item.second});
+
+  // HTTP annotations
+  using StructField = Protobuf::MapPair<std::string, ProtobufWkt::Value>;
+
+  ProtobufWkt::Struct* request = s.mutable_http()->mutable_request();
+  auto* request_fields = request->mutable_fields();
+  for (const auto& field : http_request_annotations_) {
+    request_fields->insert(StructField{field.first, field.second});
   }
 
-  for (const auto& item : http_request_annotations_) {
-    s.mutable_http()->mutable_request()->insert(KeyValue{item.first, item.second});
+  ProtobufWkt::Struct* response = s.mutable_http()->mutable_response();
+  auto* response_fields = response->mutable_fields();
+  for (const auto& field : http_response_annotations_) {
+    response_fields->insert(StructField{field.first, field.second});
   }
 
-  for (const auto& item : http_response_annotations_) {
-    s.mutable_http()->mutable_response()->insert(KeyValue{item.first, item.second});
+  using KeyValue = Protobuf::Map<std::string, std::string>::value_type;
+  for (const auto& item : custom_annotations_) {
+    s.mutable_annotations()->insert(KeyValue{item.first, item.second});
   }
 
   const std::string json = MessageUtil::getJsonStringFromMessage(
@@ -179,20 +188,30 @@ void Span::setTag(absl::string_view name, absl::string_view value) {
   }
 
   if (name == HttpUrl) {
-    http_request_annotations_.emplace(SpanUrl, value);
+    http_request_annotations_.emplace(SpanUrl, ValueUtil::stringValue(std::string(value)));
   } else if (name == HttpMethod) {
-    http_request_annotations_.emplace(SpanMethod, value);
+    http_request_annotations_.emplace(SpanMethod, ValueUtil::stringValue(std::string(value)));
   } else if (name == HttpUserAgent) {
-    http_request_annotations_.emplace(SpanUserAgent, value);
+    http_request_annotations_.emplace(SpanUserAgent, ValueUtil::stringValue(std::string(value)));
   } else if (name == HttpStatusCode) {
-    http_response_annotations_.emplace(SpanStatus, value);
+    uint64_t status_code;
+    if (!absl::SimpleAtoi(value, &status_code)) {
+      ENVOY_LOG(debug, "{} must be a number, given: {}", HttpStatusCode, value);
+      return;
+    }
+    http_response_annotations_.emplace(SpanStatus, ValueUtil::numberValue(status_code));
   } else if (name == HttpResponseSize) {
-    http_response_annotations_.emplace(SpanContentLength, value);
+    uint64_t response_size;
+    if (!absl::SimpleAtoi(value, &response_size)) {
+      ENVOY_LOG(debug, "{} must be a number, given: {}", HttpResponseSize, value);
+      return;
+    }
+    http_response_annotations_.emplace(SpanContentLength, ValueUtil::numberValue(response_size));
   } else if (name == PeerAddress) {
-    http_request_annotations_.emplace(SpanClientIp, value);
+    http_request_annotations_.emplace(SpanClientIp, ValueUtil::stringValue(std::string(value)));
     // In this case, PeerAddress refers to the client's actual IP address, not
     // the address specified in the the HTTP X-Forwarded-For header.
-    http_request_annotations_.emplace(SpanXForwardedFor, "false");
+    http_request_annotations_.emplace(SpanXForwardedFor, ValueUtil::boolValue(false));
   } else {
     custom_annotations_.emplace(name, value);
   }

source/extensions/tracers/xray/tracer.h

@@ -8,6 +8,7 @@
 #include "envoy/tracing/http_tracer.h"
 
 #include "common/common/hex.h"
+#include "common/protobuf/utility.h"
 
 #include "extensions/tracers/xray/daemon_broker.h"
 #include "extensions/tracers/xray/sampling_strategy.h"
@@ -23,7 +24,7 @@ namespace XRay {
 
 constexpr auto XRayTraceHeader = "x-amzn-trace-id";
 
-class Span : public Tracing::Span {
+class Span : public Tracing::Span, Logger::Loggable<Logger::Id::config> {
 public:
   /**
    * Creates a new Span.
@@ -147,8 +148,8 @@ class Span : public Tracing::Span {
   std::string trace_id_;
   std::string parent_segment_id_;
   std::string name_;
-  absl::flat_hash_map<std::string, std::string> http_request_annotations_;
-  absl::flat_hash_map<std::string, std::string> http_response_annotations_;
+  absl::flat_hash_map<std::string, ProtobufWkt::Value> http_request_annotations_;
+  absl::flat_hash_map<std::string, ProtobufWkt::Value> http_response_annotations_;
   absl::flat_hash_map<std::string, std::string> custom_annotations_;
   Envoy::TimeSource& time_source_;
   DaemonBroker& broker_;

test/extensions/tracers/xray/tracer_test.cc

@@ -13,6 +13,7 @@
 #include "test/mocks/server/mocks.h"
 #include "test/mocks/tracing/mocks.h"
 
+#include "absl/strings/str_format.h"
 #include "absl/strings/str_split.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
@@ -46,10 +47,10 @@ TEST_F(XRayTracerTest, SerializeSpanTest) {
   constexpr auto expected_http_method = "POST";
   constexpr auto expected_http_url = "/first/second";
   constexpr auto expected_user_agent = "Mozilla/5.0 (Macintosh; Intel Mac OS X)";
-  constexpr auto expected_status_code = "202";
-  constexpr auto expected_content_length = "1337";
+  constexpr uint32_t expected_status_code = 202;
+  constexpr uint32_t expected_content_length = 1337;
   constexpr auto expected_client_ip = "10.0.0.100";
-  constexpr auto expected_x_forwarded_for = "false";
+  constexpr auto expected_x_forwarded_for = false;
   constexpr auto expected_upstream_address = "10.0.0.200";
 
   auto on_send = [&](const std::string& json) {
@@ -61,13 +62,19 @@ TEST_F(XRayTracerTest, SerializeSpanTest) {
     ASSERT_EQ(1, s.annotations().size());
     ASSERT_TRUE(s.parent_id().empty());
     ASSERT_STREQ(expected_span_name, s.name().c_str());
-    ASSERT_STREQ(expected_http_method, s.http().request().at("method").c_str());
-    ASSERT_STREQ(expected_http_url, s.http().request().at("url").c_str());
-    ASSERT_STREQ(expected_user_agent, s.http().request().at("user_agent").c_str());
-    ASSERT_STREQ(expected_status_code, s.http().response().at("status").c_str());
-    ASSERT_STREQ(expected_content_length, s.http().response().at("content_length").c_str());
-    ASSERT_STREQ(expected_client_ip, s.http().request().at("client_ip").c_str());
-    ASSERT_STREQ(expected_x_forwarded_for, s.http().request().at("x_forwarded_for").c_str());
+    ASSERT_STREQ(expected_http_method,
+                 s.http().request().fields().at("method").string_value().c_str());
+    ASSERT_STREQ(expected_http_url, s.http().request().fields().at("url").string_value().c_str());
+    ASSERT_STREQ(expected_user_agent,
+                 s.http().request().fields().at("user_agent").string_value().c_str());
+    ASSERT_DOUBLE_EQ(expected_status_code,
+                     s.http().response().fields().at("status").number_value());
+    ASSERT_DOUBLE_EQ(expected_content_length,
+                     s.http().response().fields().at("content_length").number_value());
+    ASSERT_STREQ(expected_client_ip,
+                 s.http().request().fields().at("client_ip").string_value().c_str());
+    ASSERT_EQ(expected_x_forwarded_for,
+              s.http().request().fields().at("x_forwarded_for").bool_value());
     ASSERT_STREQ(expected_upstream_address, s.annotations().at("upstream_address").c_str());
   };
 
@@ -78,8 +85,8 @@ TEST_F(XRayTracerTest, SerializeSpanTest) {
   span->setTag("http.method", expected_http_method);
   span->setTag("http.url", expected_http_url);
   span->setTag("user_agent", expected_user_agent);
-  span->setTag("http.status_code", expected_status_code);
-  span->setTag("response_size", expected_content_length);
+  span->setTag("http.status_code", absl::StrFormat("%d", expected_status_code));
+  span->setTag("response_size", absl::StrFormat("%d", expected_content_length));
   span->setTag("peer.address", expected_client_ip);
   span->setTag("upstream_address", expected_upstream_address);
   span->finishSpan();