From 7c5d8cf244924ad468892a5081650705c272a87b Mon Sep 17 00:00:00 2001
From: Alejandro Celaya <alejandro@alejandrocelaya.com>
Date: Tue, 9 Aug 2016 09:13:39 +0200
Subject: [PATCH] Fixed VisitsTracker to take into account the X-Forwarded-For
 header in case the server is behind a load balabncer or proxy

---
 module/Core/src/Service/VisitsTracker.php     | 25 +++++++++++--------
 .../Core/test/Service/VisitsTrackerTest.php   | 23 +++++++++++++++++
 2 files changed, 37 insertions(+), 11 deletions(-)

diff --git a/module/Core/src/Service/VisitsTracker.php b/module/Core/src/Service/VisitsTracker.php
index 6ae38e83..94647086 100644
--- a/module/Core/src/Service/VisitsTracker.php
+++ b/module/Core/src/Service/VisitsTracker.php
@@ -36,8 +36,6 @@ class VisitsTracker implements VisitsTrackerInterface
      */
     public function track($shortCode, ServerRequestInterface $request)
     {
-        $visitorData = $request->getServerParams();
-
         /** @var ShortUrl $shortUrl */
         $shortUrl = $this->em->getRepository(ShortUrl::class)->findOneBy([
             'shortCode' => $shortCode,
@@ -45,22 +43,27 @@ class VisitsTracker implements VisitsTrackerInterface
 
         $visit = new Visit();
         $visit->setShortUrl($shortUrl)
-              ->setUserAgent($this->getArrayValue($visitorData, 'HTTP_USER_AGENT'))
-              ->setReferer($this->getArrayValue($visitorData, 'HTTP_REFERER'))
-              ->setRemoteAddr($this->getArrayValue($visitorData, 'REMOTE_ADDR'));
+              ->setUserAgent($request->getHeaderLine('User-Agent'))
+              ->setReferer($request->getHeaderLine('Referer'))
+              ->setRemoteAddr($this->findOutRemoteAddr($request));
         $this->em->persist($visit);
         $this->em->flush();
     }
 
     /**
-     * @param array $array
-     * @param $key
-     * @param null $default
-     * @return mixed|null
+     * @param ServerRequestInterface $request
+     * @return string
      */
-    protected function getArrayValue(array $array, $key, $default = null)
+    protected function findOutRemoteAddr(ServerRequestInterface $request)
     {
-        return isset($array[$key]) ? $array[$key] : $default;
+        $forwardedFor = $request->getHeaderLine('X-Forwarded-For');
+        if (empty($forwardedFor)) {
+            $serverParams = $request->getServerParams();
+            return isset($serverParams['REMOTE_ADDR']) ? $serverParams['REMOTE_ADDR'] : null;
+        }
+
+        $ips = explode(',', $forwardedFor);
+        return $ips[0];
     }
 
     /**
diff --git a/module/Core/test/Service/VisitsTrackerTest.php b/module/Core/test/Service/VisitsTrackerTest.php
index 12dcf013..557b715f 100644
--- a/module/Core/test/Service/VisitsTrackerTest.php
+++ b/module/Core/test/Service/VisitsTrackerTest.php
@@ -45,6 +45,29 @@ class VisitsTrackerTest extends TestCase
         $this->visitsTracker->track($shortCode, ServerRequestFactory::fromGlobals());
     }
 
+    /**
+     * @test
+     */
+    public function trackUsesForwardedForHeaderIfPresent()
+    {
+        $shortCode = '123ABC';
+        $test = $this;
+        $repo = $this->prophesize(EntityRepository::class);
+        $repo->findOneBy(['shortCode' => $shortCode])->willReturn(new ShortUrl());
+
+        $this->em->getRepository(ShortUrl::class)->willReturn($repo->reveal())->shouldBeCalledTimes(1);
+        $this->em->persist(Argument::any())->will(function ($args) use ($test) {
+            /** @var Visit $visit */
+            $visit = $args[0];
+            $test->assertEquals('4.3.2.1', $visit->getRemoteAddr());
+        })->shouldBeCalledTimes(1);
+        $this->em->flush()->shouldBeCalledTimes(1);
+
+        $this->visitsTracker->track($shortCode, ServerRequestFactory::fromGlobals(
+            ['REMOTE_ADDR' => '1.2.3.4']
+        )->withHeader('X-Forwarded-For', '4.3.2.1,99.99.99.99'));
+    }
+
     /**
      * @test
      */