From 705dc2ec390eacfdbbad267b30c7745f01de0999 Mon Sep 17 00:00:00 2001
From: Alejandro Celaya <alejandro@alejandrocelaya.com>
Date: Wed, 13 Nov 2019 21:04:44 +0100
Subject: [PATCH] Added forward of query string from short URLs to long one

---
 .../src/Action/AbstractTrackingAction.php     | 19 ++++++-
 .../Core/test/Action/RedirectActionTest.php   | 51 +++++++++----------
 2 files changed, 41 insertions(+), 29 deletions(-)

diff --git a/module/Core/src/Action/AbstractTrackingAction.php b/module/Core/src/Action/AbstractTrackingAction.php
index 464f3361..9bacaf39 100644
--- a/module/Core/src/Action/AbstractTrackingAction.php
+++ b/module/Core/src/Action/AbstractTrackingAction.php
@@ -10,14 +10,19 @@ use Psr\Http\Server\MiddlewareInterface;
 use Psr\Http\Server\RequestHandlerInterface;
 use Psr\Log\LoggerInterface;
 use Psr\Log\NullLogger;
+use Shlinkio\Shlink\Core\Entity\ShortUrl;
 use Shlinkio\Shlink\Core\Exception\EntityDoesNotExistException;
 use Shlinkio\Shlink\Core\Exception\InvalidShortCodeException;
 use Shlinkio\Shlink\Core\Model\Visitor;
 use Shlinkio\Shlink\Core\Options\AppOptions;
 use Shlinkio\Shlink\Core\Service\UrlShortenerInterface;
 use Shlinkio\Shlink\Core\Service\VisitsTrackerInterface;
+use Zend\Diactoros\Uri;
 
 use function array_key_exists;
+use function array_merge;
+use function GuzzleHttp\Psr7\parse_query;
+use function http_build_query;
 
 abstract class AbstractTrackingAction implements MiddlewareInterface
 {
@@ -66,13 +71,25 @@ abstract class AbstractTrackingAction implements MiddlewareInterface
                 $this->visitTracker->track($shortCode, Visitor::fromRequest($request));
             }
 
-            return $this->createSuccessResp($url->getLongUrl());
+            return $this->createSuccessResp($this->buildUrlToRedirectTo($url, $query, $disableTrackParam));
         } catch (InvalidShortCodeException | EntityDoesNotExistException $e) {
             $this->logger->warning('An error occurred while tracking short code. {e}', ['e' => $e]);
             return $this->createErrorResp($request, $handler);
         }
     }
 
+    private function buildUrlToRedirectTo(ShortUrl $shortUrl, array $currentQuery, ?string $disableTrackParam): string
+    {
+        $uri = new Uri($shortUrl->getLongUrl());
+        $hardcodedQuery = parse_query($uri->getQuery());
+        if ($disableTrackParam !== null) {
+            unset($currentQuery[$disableTrackParam]);
+        }
+        $mergedQuery = array_merge($hardcodedQuery, $currentQuery);
+
+        return (string) $uri->withQuery(http_build_query($mergedQuery));
+    }
+
     abstract protected function createSuccessResp(string $longUrl): ResponseInterface;
 
     abstract protected function createErrorResp(
diff --git a/module/Core/test/Action/RedirectActionTest.php b/module/Core/test/Action/RedirectActionTest.php
index eec71dbd..a65927eb 100644
--- a/module/Core/test/Action/RedirectActionTest.php
+++ b/module/Core/test/Action/RedirectActionTest.php
@@ -17,6 +17,8 @@ use Shlinkio\Shlink\Core\Service\VisitsTracker;
 use Zend\Diactoros\Response;
 use Zend\Diactoros\ServerRequest;
 
+use function array_key_exists;
+
 class RedirectActionTest extends TestCase
 {
     /** @var RedirectAction */
@@ -38,23 +40,36 @@ class RedirectActionTest extends TestCase
         );
     }
 
-    /** @test */
-    public function redirectionIsPerformedToLongUrl(): void
+    /**
+     * @test
+     * @dataProvider provideQueries
+     */
+    public function redirectionIsPerformedToLongUrl(string $expectedUrl, array $query): void
     {
         $shortCode = 'abc123';
-        $expectedUrl = 'http://domain.com/foo/bar';
-        $shortUrl = new ShortUrl($expectedUrl);
-        $this->urlShortener->shortCodeToUrl($shortCode, '')->willReturn($shortUrl)
-                                                           ->shouldBeCalledOnce();
-        $this->visitTracker->track(Argument::cetera())->shouldBeCalledOnce();
+        $shortUrl = new ShortUrl('http://domain.com/foo/bar?some=thing');
+        $shortCodeToUrl = $this->urlShortener->shortCodeToUrl($shortCode, '')->willReturn($shortUrl);
+        $track = $this->visitTracker->track(Argument::cetera())->will(function () {
+        });
 
-        $request = (new ServerRequest())->withAttribute('shortCode', $shortCode);
+        $request = (new ServerRequest())->withAttribute('shortCode', $shortCode)->withQueryParams($query);
         $response = $this->action->process($request, $this->prophesize(RequestHandlerInterface::class)->reveal());
 
         $this->assertInstanceOf(Response\RedirectResponse::class, $response);
         $this->assertEquals(302, $response->getStatusCode());
         $this->assertTrue($response->hasHeader('Location'));
         $this->assertEquals($expectedUrl, $response->getHeaderLine('Location'));
+        $shortCodeToUrl->shouldHaveBeenCalledOnce();
+        $track->shouldHaveBeenCalledTimes(array_key_exists('foobar', $query) ? 0 : 1);
+    }
+
+    public function provideQueries(): iterable
+    {
+        yield ['http://domain.com/foo/bar?some=thing', []];
+        yield ['http://domain.com/foo/bar?some=thing', ['foobar' => 'notrack']];
+        yield ['http://domain.com/foo/bar?some=thing&foo=bar', ['foo' => 'bar']];
+        yield ['http://domain.com/foo/bar?some=overwritten&foo=bar', ['foo' => 'bar', 'some' => 'overwritten']];
+        yield ['http://domain.com/foo/bar?some=overwritten', ['foobar' => 'notrack', 'some' => 'overwritten']];
     }
 
     /** @test */
@@ -73,24 +88,4 @@ class RedirectActionTest extends TestCase
 
         $handle->shouldHaveBeenCalledOnce();
     }
-
-    /** @test */
-    public function visitIsNotTrackedIfDisableParamIsProvided(): void
-    {
-        $shortCode = 'abc123';
-        $expectedUrl = 'http://domain.com/foo/bar';
-        $shortUrl = new ShortUrl($expectedUrl);
-        $this->urlShortener->shortCodeToUrl($shortCode, '')->willReturn($shortUrl)
-                                                           ->shouldBeCalledOnce();
-        $this->visitTracker->track(Argument::cetera())->shouldNotBeCalled();
-
-        $request = (new ServerRequest())->withAttribute('shortCode', $shortCode)
-                                                      ->withQueryParams(['foobar' => true]);
-        $response = $this->action->process($request, $this->prophesize(RequestHandlerInterface::class)->reveal());
-
-        $this->assertInstanceOf(Response\RedirectResponse::class, $response);
-        $this->assertEquals(302, $response->getStatusCode());
-        $this->assertTrue($response->hasHeader('Location'));
-        $this->assertEquals($expectedUrl, $response->getHeaderLine('Location'));
-    }
 }