From d1b8f0590e47a18837ea6b7ef80da38e18ea0dc8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=A9r=C3=A9my=20Deruss=C3=A9?= <jeremy@derusse.com>
Date: Mon, 22 Apr 2024 10:41:29 +0200
Subject: [PATCH] Fix compatibility with twig 3.9 and yielding

---
 src/Twig/Node/CSPNode.php | 42 +++++++++++++++++++++++++++++----------
 1 file changed, 31 insertions(+), 11 deletions(-)

diff --git a/src/Twig/Node/CSPNode.php b/src/Twig/Node/CSPNode.php
index 35e82497..7d571eac 100644
--- a/src/Twig/Node/CSPNode.php
+++ b/src/Twig/Node/CSPNode.php
@@ -22,9 +22,12 @@
  * file that was distributed with this source code.
  */
 
+use Twig\Attribute\YieldReady;
 use Twig\Compiler;
+use Twig\Node\CaptureNode;
 use Twig\Node\Node;
 
+#[YieldReady]
 final class CSPNode extends Node
 {
     private ?string $sha;
@@ -32,6 +35,11 @@ final class CSPNode extends Node
 
     public function __construct(Node $body, int $lineno, string $tag, string $directive, ?string $sha = null)
     {
+        if (class_exists(CaptureNode::class)) {
+            $body = new CaptureNode($body, $lineno, $tag);
+            $body->setAttribute('raw', true);
+        }
+
         parent::__construct(['body' => $body], [], $lineno, $tag);
         $this->sha = $sha;
         $this->directive = $directive;
@@ -39,23 +47,35 @@ public function __construct(Node $body, int $lineno, string $tag, string $direct
 
     public function compile(Compiler $compiler): void
     {
-        $body = $this->getNode('body');
+        if (class_exists(CaptureNode::class)) {
+            $compiler
+                ->addDebugInfo($this)
+                ->indent()
+                ->write("\$content = ")
+                ->subcompile($this->getNode('body'))
+                ->raw("\n")
+                ->outdent()
+            ;
+        } else {
+            $compiler
+                ->addDebugInfo($this)
+                ->indent()
+                ->write("ob_start();\n")
+                ->subcompile($this->getNode('body'))
+                ->outdent()
+                ->write("\$output = ob_get_clean();\n")
+            ;
+        }
 
         if (null !== $this->sha) {
-            $output = "\$this->env->getRuntime('Nelmio\SecurityBundle\Twig\CSPRuntime')->getListener()->addSha('{$this->directive}', '{$this->sha}');\necho ob_get_clean();\n";
+            $compiler->write("\$this->env->getRuntime('Nelmio\SecurityBundle\Twig\CSPRuntime')->getListener()->addSha('{$this->directive}', '{$this->sha}');\n");
         } elseif ('script-src' === $this->directive) {
-            $output = "\$script = ob_get_clean();\n\$this->env->getRuntime('Nelmio\SecurityBundle\Twig\CSPRuntime')->getListener()->addScript(\$script);\necho \$script;\n";
+            $compiler->write("\$this->env->getRuntime('Nelmio\SecurityBundle\Twig\CSPRuntime')->getListener()->addScript(\$content);\n");
         } elseif ('style-src' === $this->directive) {
-            $output = "\$style = ob_get_clean();\n\$this->env->getRuntime('Nelmio\SecurityBundle\Twig\CSPRuntime')->getListener()->addStyle(\$style);\necho \$style;\n";
+            $compiler->write("\$this->env->getRuntime('Nelmio\SecurityBundle\Twig\CSPRuntime')->getListener()->addStyle(\$content);\n");
         } else {
             throw new \InvalidArgumentException(sprintf('Unable to compile for directive "%s"', $this->directive));
         }
-
-        $compiler
-            ->addDebugInfo($this)
-            ->write("ob_start();\n")
-            ->subcompile($body)
-            ->write($output)
-        ;
+        $compiler->write("echo \$content;\n");
     }
 }