From 8a9a597590cb6affab1680217e454c8ad2705c94 Mon Sep 17 00:00:00 2001
From: snakorse <512979011@qq.com>
Date: Thu, 28 Jan 2021 10:20:54 +0800
Subject: [PATCH] feature: Support collect headers & aggregate collected
 cookies into on span tag. (#381)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-authored-by: 郭刚平 <gangping.guo@xiaobao100.com>
---
 sample/SkyApm.Sample.Frontend/skyapm.json     |  3 +-
 src/SkyApm.Abstractions/Common/Tags.cs        |  4 ++
 .../Config/HostingDiagnosticConfig.cs         |  9 ++-
 .../DefaultHostingDiagnosticHandler.cs        | 57 ++++++++++++++++---
 4 files changed, 63 insertions(+), 10 deletions(-)

diff --git a/sample/SkyApm.Sample.Frontend/skyapm.json b/sample/SkyApm.Sample.Frontend/skyapm.json
index 0469e0f3..aaed758b 100644
--- a/sample/SkyApm.Sample.Frontend/skyapm.json
+++ b/sample/SkyApm.Sample.Frontend/skyapm.json
@@ -28,7 +28,8 @@
     },
     "Component": {
       "AspNetCore": {
-        "AutoTagCookies": ["c-b"]
+        "CollectCookies": [ "c-b" ],
+        "CollectHeaders": ["User-Agent"]
       },
       "HttpClient": {
         "StopHeaderPropagationPaths": [ "**/localhost:5002/api/values/stoppropagation" ]
diff --git a/src/SkyApm.Abstractions/Common/Tags.cs b/src/SkyApm.Abstractions/Common/Tags.cs
index b3333ef7..fc24c4ce 100644
--- a/src/SkyApm.Abstractions/Common/Tags.cs
+++ b/src/SkyApm.Abstractions/Common/Tags.cs
@@ -27,6 +27,10 @@ public static class Tags
 
         public static readonly string HTTP_METHOD = "http.method";
 
+        public static readonly string HTTP_COOKIES = "http.cookies";
+
+        public static readonly string HTTP_HEADERS = "http.headers";
+
         public static readonly string STATUS_CODE = "status_code";
 
         public static readonly string DB_TYPE = "db.type";
diff --git a/src/SkyApm.Diagnostics.AspNetCore/Config/HostingDiagnosticConfig.cs b/src/SkyApm.Diagnostics.AspNetCore/Config/HostingDiagnosticConfig.cs
index 93e85a97..457f3732 100644
--- a/src/SkyApm.Diagnostics.AspNetCore/Config/HostingDiagnosticConfig.cs
+++ b/src/SkyApm.Diagnostics.AspNetCore/Config/HostingDiagnosticConfig.cs
@@ -9,8 +9,13 @@ namespace SkyApm.Diagnostics.AspNetCore.Config
     public class HostingDiagnosticConfig
     {
         /// <summary>
-        /// Auto collect specific cookies as span tags.
+        /// Auto collect specific cookies as span tag.
         /// </summary>
-        public List<string> AutoTagCookies { get; set; }
+        public List<string> CollectCookies { get; set; }
+
+        /// <summary>
+        /// Auto collect specific headers as span tag
+        /// </summary>
+        public List<string> CollectHeaders { get; set; }
     }
 }
diff --git a/src/SkyApm.Diagnostics.AspNetCore/Handlers/DefaultHostingDiagnosticHandler.cs b/src/SkyApm.Diagnostics.AspNetCore/Handlers/DefaultHostingDiagnosticHandler.cs
index 08b2cf71..97eddbc4 100644
--- a/src/SkyApm.Diagnostics.AspNetCore/Handlers/DefaultHostingDiagnosticHandler.cs
+++ b/src/SkyApm.Diagnostics.AspNetCore/Handlers/DefaultHostingDiagnosticHandler.cs
@@ -18,12 +18,15 @@
 
 using Microsoft.AspNetCore.Http;
 using Microsoft.AspNetCore.Http.Extensions;
+using Microsoft.Extensions.Primitives;
 using SkyApm.AspNetCore.Diagnostics;
 using SkyApm.Common;
 using SkyApm.Config;
 using SkyApm.Diagnostics.AspNetCore.Config;
 using SkyApm.Tracing;
 using SkyApm.Tracing.Segments;
+using System.Collections.Generic;
+using System.Text;
 
 namespace SkyApm.Diagnostics.AspNetCore.Handlers
 {
@@ -52,14 +55,18 @@ public void BeginRequest(ITracingContext tracingContext, HttpContext httpContext
             context.Span.AddTag(Tags.PATH, httpContext.Request.Path);
             context.Span.AddTag(Tags.HTTP_METHOD, httpContext.Request.Method);
 
-            if(_config.AutoTagCookies?.Count > 0)
+            if(_config.CollectCookies?.Count > 0)
             {
-                foreach (var key in _config.AutoTagCookies)
-                {
-                    if (!httpContext.Request.Cookies.TryGetValue(key, out string value))
-                        continue;
-                    context.Span.AddTag("cookie." + key, value);
-                }
+                var cookies = CollectCookies(httpContext, _config.CollectCookies);
+                if (!string.IsNullOrEmpty(cookies))
+                    context.Span.AddTag(Tags.HTTP_COOKIES, cookies);
+            }
+
+            if(_config.CollectHeaders?.Count > 0)
+            {
+                var headers = CollectHeaders(httpContext, _config.CollectHeaders);
+                if (!string.IsNullOrEmpty(headers))
+                    context.Span.AddTag(Tags.HTTP_HEADERS, headers);
             }
         }
 
@@ -73,5 +80,41 @@ public void EndRequest(SegmentContext segmentContext, HttpContext httpContext)
 
             segmentContext.Span.AddTag(Tags.STATUS_CODE, statusCode);
         }
+
+        private string CollectCookies(HttpContext httpContext, IEnumerable<string> keys)
+        {
+            var sb = new StringBuilder();
+            foreach (var key in keys)
+            {
+                if (!httpContext.Request.Cookies.TryGetValue(key, out string value))
+                    continue;
+
+                if(sb.Length > 0)
+                    sb.Append("; ");
+
+                sb.Append(key);
+                sb.Append('=');
+                sb.Append(value);
+            }
+            return sb.ToString();
+        }
+
+        private string CollectHeaders(HttpContext httpContext, IEnumerable<string> keys)
+        {
+            var sb = new StringBuilder();
+            foreach (var key in keys)
+            {
+                if (!httpContext.Request.Headers.TryGetValue(key, out StringValues value))
+                    continue;
+
+                if(sb.Length > 0)
+                    sb.Append('\n');
+
+                sb.Append(key);
+                sb.Append(": ");
+                sb.Append(value);
+            }
+            return sb.ToString();
+        }
     }
 }
\ No newline at end of file