From 8b93f4407720551ae0b29d3db17e16fd238e3357 Mon Sep 17 00:00:00 2001
From: Philip Wood <philip.m.wood@gmail.com>
Date: Tue, 18 Apr 2017 11:05:37 +0100
Subject: [PATCH] Fix request id middleware

---
 .../Middleware/OcelotMiddlewareExtensions.cs  |  2 +-
 .../Middleware/RequestIdMiddleware.cs         | 44 +++++++++++--------
 2 files changed, 26 insertions(+), 20 deletions(-)

diff --git a/src/Ocelot/Middleware/OcelotMiddlewareExtensions.cs b/src/Ocelot/Middleware/OcelotMiddlewareExtensions.cs
index e35c2d11..105be8dd 100644
--- a/src/Ocelot/Middleware/OcelotMiddlewareExtensions.cs
+++ b/src/Ocelot/Middleware/OcelotMiddlewareExtensions.cs
@@ -72,7 +72,7 @@ namespace Ocelot.Middleware
             builder.UseRateLimiting();
 
             // Now we can look for the requestId
-            builder.UseRequestIdMiddleware(); //obsolete?
+            builder.UseRequestIdMiddleware();
 
             // Allow pre authentication logic. The idea being people might want to run something custom before what is built in.
             builder.UseIfNotNull(middlewareConfiguration.PreAuthenticationMiddleware);
diff --git a/src/Ocelot/RequestId/Middleware/RequestIdMiddleware.cs b/src/Ocelot/RequestId/Middleware/RequestIdMiddleware.cs
index 168c377a..08222d9d 100644
--- a/src/Ocelot/RequestId/Middleware/RequestIdMiddleware.cs
+++ b/src/Ocelot/RequestId/Middleware/RequestIdMiddleware.cs
@@ -41,34 +41,40 @@ namespace Ocelot.RequestId.Middleware
 
         private void SetOcelotRequestId(HttpContext context)
         {
-            var key = DefaultRequestIdKey.Value;
-
-            if (DownstreamRoute.ReRoute.RequestIdKey != null)
-            {
-                key = DownstreamRoute.ReRoute.RequestIdKey;
-            }
+            // if get request ID is set on upstream request then retrieve it
+            var key = DownstreamRoute.ReRoute.RequestIdKey ?? DefaultRequestIdKey.Value;
             
-            StringValues requestIds;
-
-            if (context.Request.Headers.TryGetValue(key, out requestIds))
+            StringValues upstreamRequestIds;
+            if (context.Request.Headers.TryGetValue(key, out upstreamRequestIds))
             {
-                var requestId = requestIds.First();
-                var downstreamRequestHeaders = DownstreamRequest.Headers;
+                context.TraceIdentifier = upstreamRequestIds.First();
+            }
 
-                if (!string.IsNullOrEmpty(requestId) && 
-                    !HeaderExists(key, downstreamRequestHeaders))
-                {
-                    downstreamRequestHeaders.Add(key, requestId);
-                }
+            // set request ID on downstream request, if required
+            var requestId = new RequestId(DownstreamRoute?.ReRoute?.RequestIdKey, context.TraceIdentifier);
 
-                context.TraceIdentifier = requestId;
+            if (ShouldAddRequestId(requestId, DownstreamRequest.Headers))
+            {
+                AddRequestIdHeader(requestId, DownstreamRequest);
             }
         }
 
-        private bool HeaderExists(string headerKey, HttpRequestHeaders headers)
+        private bool ShouldAddRequestId(RequestId requestId, HttpRequestHeaders headers)
+        {
+            return !string.IsNullOrEmpty(requestId?.RequestIdKey)
+                   && !string.IsNullOrEmpty(requestId.RequestIdValue)
+                   && !RequestIdInHeaders(requestId, headers);
+        }
+
+        private bool RequestIdInHeaders(RequestId requestId, HttpRequestHeaders headers)
         {
             IEnumerable<string> value;
-            return headers.TryGetValues(headerKey, out value);
+            return headers.TryGetValues(requestId.RequestIdKey, out value);
+        }
+
+        private void AddRequestIdHeader(RequestId requestId, HttpRequestMessage httpRequestMessage)
+        {
+            httpRequestMessage.Headers.Add(requestId.RequestIdKey, requestId.RequestIdValue);
         }
     }
 }
\ No newline at end of file