From db1a9c254797363dd10052833cc02ea9d840489f Mon Sep 17 00:00:00 2001
From: Martin Herndl <martin@herndl.org>
Date: Wed, 1 Sep 2021 10:22:05 +0200
Subject: [PATCH] [HttpKernel] Fix `TimeDataCollector`

In order for the `TimeDataCollector` to work properly the `X-Debug-Token` from the response header needs to be used which was removed in #42331 to fix #36623. Fixes #42804
---
 .../Debug/TraceableEventDispatcher.php        | 11 +++++---
 .../Debug/TraceableEventDispatcherTest.php    | 26 ++++++++++++++++---
 2 files changed, 31 insertions(+), 6 deletions(-)

diff --git a/src/Symfony/Component/HttpKernel/Debug/TraceableEventDispatcher.php b/src/Symfony/Component/HttpKernel/Debug/TraceableEventDispatcher.php
index 832bfb58d0637..46d17e0ba604e 100644
--- a/src/Symfony/Component/HttpKernel/Debug/TraceableEventDispatcher.php
+++ b/src/Symfony/Component/HttpKernel/Debug/TraceableEventDispatcher.php
@@ -41,7 +41,7 @@ protected function beforeDispatch(string $eventName, $event)
                 }
                 break;
             case KernelEvents::TERMINATE:
-                $sectionId = $event->getRequest()->attributes->get('_stopwatch_token');
+                $sectionId = $this->determineStopwatchToken($event);
                 if (null === $sectionId) {
                     break;
                 }
@@ -68,7 +68,7 @@ protected function afterDispatch(string $eventName, $event)
                 $this->stopwatch->start('controller', 'section');
                 break;
             case KernelEvents::RESPONSE:
-                $sectionId = $event->getRequest()->attributes->get('_stopwatch_token');
+                $sectionId = $this->determineStopwatchToken($event);
                 if (null === $sectionId) {
                     break;
                 }
@@ -77,7 +77,7 @@ protected function afterDispatch(string $eventName, $event)
             case KernelEvents::TERMINATE:
                 // In the special case described in the `preDispatch` method above, the `$token` section
                 // does not exist, then closing it throws an exception which must be caught.
-                $sectionId = $event->getRequest()->attributes->get('_stopwatch_token');
+                $sectionId = $this->determineStopwatchToken($event);
                 if (null === $sectionId) {
                     break;
                 }
@@ -88,4 +88,9 @@ protected function afterDispatch(string $eventName, $event)
                 break;
         }
     }
+
+    private function determineStopwatchToken($event)
+    {
+        return $event->getResponse()->headers->get('X-Debug-Token') ?? $event->getRequest()->attributes->get('_stopwatch_token');
+    }
 }
diff --git a/src/Symfony/Component/HttpKernel/Tests/Debug/TraceableEventDispatcherTest.php b/src/Symfony/Component/HttpKernel/Tests/Debug/TraceableEventDispatcherTest.php
index 90ea544b1abe2..6e5cd538d4167 100644
--- a/src/Symfony/Component/HttpKernel/Tests/Debug/TraceableEventDispatcherTest.php
+++ b/src/Symfony/Component/HttpKernel/Tests/Debug/TraceableEventDispatcherTest.php
@@ -45,6 +45,26 @@ public function testStopwatchSections()
         ], array_keys($events));
     }
 
+    public function testStopwatchSectionsWithProfilerToken()
+    {
+        $dispatcher = new TraceableEventDispatcher(new EventDispatcher(), $stopwatch = new Stopwatch());
+        $kernel = $this->getHttpKernel($dispatcher, new Response('', 200, ['X-Debug-Token' => '292e1e']));
+        $request = Request::create('/');
+        $response = $kernel->handle($request);
+        $kernel->terminate($request, $response);
+
+        $events = $stopwatch->getSectionEvents($response->headers->get('X-Debug-Token'));
+        $this->assertEquals([
+            '__section__',
+            'kernel.request',
+            'kernel.controller',
+            'kernel.controller_arguments',
+            'controller',
+            'kernel.response',
+            'kernel.terminate',
+        ], array_keys($events));
+    }
+
     public function testStopwatchCheckControllerOnRequestEvent()
     {
         $stopwatch = $this->getMockBuilder(Stopwatch::class)
@@ -110,11 +130,11 @@ public function testListenerCanRemoveItselfWhenExecuted()
         $this->assertCount(1, $eventDispatcher->getListeners('foo'), 'expected listener1 to be removed');
     }
 
-    protected function getHttpKernel($dispatcher)
+    protected function getHttpKernel($dispatcher, Response $response = null)
     {
         $controllerResolver = $this->createMock(ControllerResolverInterface::class);
-        $controllerResolver->expects($this->once())->method('getController')->willReturn(function () {
-            return new Response();
+        $controllerResolver->expects($this->once())->method('getController')->willReturn(function () use ($response) {
+            return $response ?? new Response();
         });
         $argumentResolver = $this->createMock(ArgumentResolverInterface::class);
         $argumentResolver->expects($this->once())->method('getArguments')->willReturn([]);