XFF should be added even if no headers (with tests)

kayakhttp · Mar 14, 2012 · faa9e6a · faa9e6a
1 parent 2929fc3
commit faa9e6a
Show file tree

Hide file tree

Showing 4 changed files with 43 additions and 4 deletions.
diff --git a/Kayak.Tests/Http/HttpServerTransactionDelegateTests.cs b/Kayak.Tests/Http/HttpServerTransactionDelegateTests.cs
@@ -80,6 +80,20 @@ public void SetUp()
             postedActions = new Queue<Action>();
         }
 
+        [Test]
+        public void Adds_xff_if_no_headers()
+        {
+            responseAccumulator.RemoteEndPoint = new IPEndPoint(IPAddress.Parse("1.1.1.1"), 0);
+
+            transactionInput.OnRequest(Request.OneOneNoBodyNoHeaders);
+            transactionInput.OnRequestEnd();
+
+            var expected = Request.OneOhNoBodyNoHeaders;
+			expected.Head.Headers = new Dictionary<string,string>(StringComparer.InvariantCultureIgnoreCase);
+            expected.Head.Headers["X-Forwarded-For"] = "1.1.1.1";
+            requestAccumulator.AssertRequests(new[] { expected });
+        }
+
         [Test]
         public void Adds_xff_if_none()
         {

diff --git a/Kayak.Tests/Http/Transaction/Request.cs b/Kayak.Tests/Http/Transaction/Request.cs
@@ -51,6 +51,24 @@ class Request
             Data = new[] { "hello ", "world!" }
         };
 
+        public static RequestInfo OneOhNoBodyNoHeaders = new RequestInfo()
+        {
+            Head = new HttpRequestHead()
+            {
+                Version = new Version(1, 0),
+                Headers = null
+            }
+        };
+
+        public static RequestInfo OneOneNoBodyNoHeaders = new RequestInfo()
+        {
+            Head = new HttpRequestHead()
+            {
+                Version = new Version(1, 1),
+                Headers = null
+            }
+        };
+
         public static RequestInfo OneOhNoBody = new RequestInfo()
         {
             Head = new HttpRequestHead()

diff --git a/Kayak.Tests/Http/Transaction/TransactionInput.cs b/Kayak.Tests/Http/Transaction/TransactionInput.cs
@@ -28,11 +28,15 @@ public void OnRequest(RequestInfo request)
         bool ShouldKeepAlive(HttpRequestHead head)
         {
             if (head.Version.Major > 0 && head.Version.Minor > 0)
+            {
                 // HTTP/1.1
+                if (head.Headers == null) return true;
                 return !(head.Headers.ContainsKey("connection") && head.Headers["connection"] == "close");
-            else
+            } else {
                 // < HTTP/1.1
+                if (head.Headers == null) return false;
                 return (head.Headers.ContainsKey("connection") && head.Headers["connection"] == "keep-alive");
+            }
         }
 
         public void OnRequestData(string data)

diff --git a/Kayak/Http/HttpServerTransactionDelegate.cs b/Kayak/Http/HttpServerTransactionDelegate.cs
@@ -28,7 +28,7 @@ void QueueSegment(ITransactionSegment segment)
 
         public void OnRequest(IHttpServerTransaction transaction, HttpRequestHead request, bool shouldKeepAlive)
         {
-            AddXFF(request, transaction.RemoteEndPoint);
+            AddXFF(ref request, transaction.RemoteEndPoint);
 
             var expectContinue = request.IsContinueExpected();
             var ignoreResponseBody = request.Method != null && request.Method.ToUpperInvariant() == "HEAD";
@@ -123,10 +123,13 @@ public void OnResponse(HttpResponseHead head, IDataProducer body)
             }
         }
 
-        void AddXFF(HttpRequestHead request, IPEndPoint remoteEndPoint)
+        void AddXFF(ref HttpRequestHead request, IPEndPoint remoteEndPoint)
         {
-            if (remoteEndPoint != null && request.Headers != null)
+            if (remoteEndPoint != null)
             {
+                if (request.Headers == null)
+                    request.Headers = new Dictionary<string, string>(StringComparer.InvariantCultureIgnoreCase);
+
                 if (request.Headers.ContainsKey("X-Forwarded-For"))
                 {
                     request.Headers["X-Forwarded-For"] += "," + remoteEndPoint.Address.ToString();