Configure headers in response setup

Author: mabar

src/Bridges/HttpDI/HttpExtension.php

@@ -10,8 +10,6 @@
 namespace Nette\Bridges\HttpDI;
 
 use Nette;
-use Nette\PhpGenerator\Helpers;
-
 
 /**
  * HTTP extension for Nette DI.
@@ -71,13 +69,13 @@ public function loadConfiguration()
 	}
 
 
-	public function afterCompile(Nette\PhpGenerator\ClassType $class)
+	public function beforeCompile()
 	{
 		if ($this->cliMode) {
 			return;
 		}
 
-		$initialize = $class->getMethod('initialize');
+		$builder = $this->getContainerBuilder();
 		$config = $this->getConfig();
 		$headers = array_map('strval', $config['headers']);
 
@@ -91,16 +89,14 @@ public function afterCompile(Nette\PhpGenerator\ClassType $class)
 			$headers['X-Frame-Options'] = $frames;
 		}
 
-		$code = [];
 		foreach (['csp', 'cspReportOnly'] as $key) {
 			if (empty($config[$key])) {
 				continue;
 			}
 			$value = self::buildPolicy($config[$key]);
 			if (strpos($value, "'nonce'")) {
-				$code[0] = '$cspNonce = base64_encode(random_bytes(16));';
 				$value = Nette\DI\ContainerBuilder::literal(
-					'str_replace(?, ? . $cspNonce, ?)',
+					'str_replace(?, ? . (isset($cspNonce) \? $cspNonce : $cspNonce = base64_encode(random_bytes(16))), ?)',
 					["'nonce", "'nonce-", $value]
 				);
 			}
@@ -111,16 +107,15 @@ public function afterCompile(Nette\PhpGenerator\ClassType $class)
 			$headers['Feature-Policy'] = self::buildPolicy($config['featurePolicy']);
 		}
 
-		$code[] = Helpers::formatArgs('$response = $this->getService(?);', [$this->prefix('response')]);
 		foreach ($headers as $key => $value) {
 			if ($value !== '') {
-				$code[] = Helpers::formatArgs('$response->setHeader(?, ?);', [$key, $value]);
+				$builder->getDefinition($this->prefix('response'))
+					->addSetup('?->setHeader(?, ?);', ['@self', $key, $value]);
 			}
 		}
 
-		$code[] = Helpers::formatArgs('$response->setCookie(...?);', [['nette-samesite', '1', 0, '/', null, null, true, 'Strict']]);
-
-		$initialize->addBody("(function () {\n\t" . implode("\n\t", $code) . "\n})();");
+		$builder->getDefinition($this->prefix('response'))
+			->addSetup('?->setCookie(...?)', ['@self', ['nette-samesite', '1', 0, '/', null, null, true, 'Strict']]);
 	}