From 49d2e2b049f85d62b96f890db3cab313afeedb01 Mon Sep 17 00:00:00 2001
From: Luke Rodgers <lr@amp.co>
Date: Mon, 24 Oct 2022 10:33:30 +0100
Subject: [PATCH] Register shutdown function to set header (#15)

Fix header is set on `Allowed memory size of X bytes exhausted` errors
---
 .travis.yml                                   |  1 -
 phpunit.xml.dist                              |  2 +-
 src/Service/CorrelationIdentifier.php         | 32 +++++++++++++++++++
 .../Service/CorrelationIdentifierTest.php     | 27 ++++++++++++++++
 4 files changed, 60 insertions(+), 2 deletions(-)

diff --git a/.travis.yml b/.travis.yml
index 8be6c13..4af3115 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -30,7 +30,6 @@ cache:
 
 before_install:
     - if [ ! "$TRAVIS_PULL_REQUEST" = "false" ]; then git branch; git branch -D "$TRAVIS_BRANCH" || true; git checkout -b "$TRAVIS_BRANCH"; fi
-    - phpenv config-rm xdebug.ini || true
     - composer self-update --2
 install:
     - export COMPOSER_MEMORY_LIMIT=-1
diff --git a/phpunit.xml.dist b/phpunit.xml.dist
index d16c6bd..23670a6 100644
--- a/phpunit.xml.dist
+++ b/phpunit.xml.dist
@@ -1,7 +1,7 @@
 <?xml version="1.0" encoding="UTF-8"?>
 <phpunit xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" bootstrap="dev/phpunit-unit-bootstrap.php" backupGlobals="false"
          backupStaticAttributes="false" colors="true" verbose="true" convertErrorsToExceptions="true"
-         convertNoticesToExceptions="true" convertWarningsToExceptions="true" processIsolation="false"
+         convertNoticesToExceptions="true" convertWarningsToExceptions="true" processIsolation="true"
          stopOnFailure="false" xsi:noNamespaceSchemaLocation="https://schema.phpunit.de/9.3/phpunit.xsd">
     <testsuites>
         <testsuite name="Ampersand Module Unit Test Suite">
diff --git a/src/Service/CorrelationIdentifier.php b/src/Service/CorrelationIdentifier.php
index 71fa70b..308492f 100644
--- a/src/Service/CorrelationIdentifier.php
+++ b/src/Service/CorrelationIdentifier.php
@@ -2,6 +2,7 @@
 declare(strict_types=1);
 namespace Ampersand\LogCorrelationId\Service;
 
+use Ampersand\LogCorrelationId\HttpResponse\HeaderProvider\LogCorrelationIdHeader as Header;
 use Ampersand\LogCorrelationId\Service\CorrelationIdentifier\Storage;
 use Magento\Framework\App\Request\Http as HttpRequest;
 
@@ -23,6 +24,8 @@ public function __construct(string $identifierKey, string $headerInput)
     {
         Storage::setKey($identifierKey);
         $this->headerInput = $headerInput;
+        // phpcs:ignore Magento2.Functions.DiscouragedFunction.Discouraged
+        register_shutdown_function([$this, 'shutDownFunction']);
     }
 
     /**
@@ -70,4 +73,33 @@ public function getIdentifierKey(): string
     {
         return Storage::getKey();
     }
+
+    /**
+     * A shutdown function to ensure the correlation ID header is added for every type of erroring request
+     *
+     * This was added to catch "Allowed memory size of X bytes exhausted" type errors
+     *
+     * @see \Magento\Framework\Webapi\ErrorProcessor::registerShutdownFunction
+     *
+     * @return void
+     */
+    public function shutDownFunction()
+    {
+        // phpcs:ignore Magento2.Functions.DiscouragedFunction.Discouraged
+        if (headers_sent()) {
+            return;
+        }
+
+        // phpcs:ignore Magento2.Functions.DiscouragedFunction.Discouraged
+        $headerAlreadyDefined = array_filter(headers_list(), function ($header) {
+            return (stripos($header, Header::X_LOG_CORRELATION_ID) !== false);
+        });
+
+        if (!empty($headerAlreadyDefined)) {
+            return;
+        }
+
+        // phpcs:ignore Magento2.Functions.DiscouragedFunction.Discouraged
+        header(Header::X_LOG_CORRELATION_ID . ': ' . $this->getIdentifierValue());
+    }
 }
diff --git a/src/Test/Unit/Service/CorrelationIdentifierTest.php b/src/Test/Unit/Service/CorrelationIdentifierTest.php
index cdba520..1e37b47 100644
--- a/src/Test/Unit/Service/CorrelationIdentifierTest.php
+++ b/src/Test/Unit/Service/CorrelationIdentifierTest.php
@@ -2,6 +2,7 @@
 declare(strict_types=1);
 namespace Ampersand\LogCorrelationId\Test\Unit\Service;
 
+use Ampersand\LogCorrelationId\HttpResponse\HeaderProvider\LogCorrelationIdHeader as Header;
 use Ampersand\LogCorrelationId\Service\CorrelationIdentifier;
 use Magento\Framework\App\Request\Http as HttpRequest;
 use PHPUnit\Framework\TestCase;
@@ -53,6 +54,32 @@ public function testGetIdentifier()
         );
     }
 
+    /**
+     */
+    public function testIdentifierHeaderIsSetOnShutdown()
+    {
+        $httpRequest = $this->createMock(HttpRequest::class);
+        $httpRequest->expects($this->any())
+            ->method('getHeader')
+            ->willReturn(false);
+
+        $service = $this->createService();
+        $service->init($httpRequest, true);
+
+        // https://stackoverflow.com/a/39892373/4354325
+        $headerIsNotSet = empty(array_filter(xdebug_get_headers(), function ($header) {
+            return (stripos($header, Header::X_LOG_CORRELATION_ID) !== false);
+        }));
+        $this->assertTrue($headerIsNotSet, 'The correlation header should not yet be set');
+
+        $service->shutDownFunction();
+
+        $headerIsSet = !empty(array_filter(xdebug_get_headers(), function ($header) {
+            return (stripos($header, Header::X_LOG_CORRELATION_ID) !== false);
+        }));
+        $this->assertTrue($headerIsSet, 'The correlation header should be set');
+    }
+
     public function testGetIdentifierFromHeader()
     {
         $httpRequest = $this->createMock(HttpRequest::class);