From cbb8f6edc180d48bdd08bd1adb4dd0ad83b9f138 Mon Sep 17 00:00:00 2001
From: Fabien Potencier <fabien@potencier.org>
Date: Mon, 11 Mar 2019 08:52:26 +0100
Subject: [PATCH] tweaked sandbox generated code

---
 src/Node/CheckSecurityNode.php         | 3 ++-
 src/Node/SandboxNode.php               | 7 +++----
 src/NodeVisitor/SandboxNodeVisitor.php | 2 +-
 src/Template.php                       | 1 +
 test/Twig/Tests/Node/SandboxTest.php   | 7 +++----
 5 files changed, 10 insertions(+), 10 deletions(-)

diff --git a/src/Node/CheckSecurityNode.php b/src/Node/CheckSecurityNode.php
index f300f080f..cf0a7a13d 100644
--- a/src/Node/CheckSecurityNode.php
+++ b/src/Node/CheckSecurityNode.php
@@ -45,12 +45,13 @@ class CheckSecurityNode extends Node
         }
 
         $compiler
+            ->write("\$this->sandbox = \$this->env->getExtension('\Twig\Extension\SandboxExtension');\n")
             ->write('$tags = ')->repr(array_filter($tags))->raw(";\n")
             ->write('$filters = ')->repr(array_filter($filters))->raw(";\n")
             ->write('$functions = ')->repr(array_filter($functions))->raw(";\n\n")
             ->write("try {\n")
             ->indent()
-            ->write("\$this->env->getExtension('\Twig\Extension\SandboxExtension')->checkSecurity(\n")
+            ->write("\$this->sandbox->checkSecurity(\n")
             ->indent()
             ->write(!$tags ? "[],\n" : "['".implode("', '", array_keys($tags))."'],\n")
             ->write(!$filters ? "[],\n" : "['".implode("', '", array_keys($filters))."'],\n")
diff --git a/src/Node/SandboxNode.php b/src/Node/SandboxNode.php
index af1a5f284..2d644c3a1 100644
--- a/src/Node/SandboxNode.php
+++ b/src/Node/SandboxNode.php
@@ -29,16 +29,15 @@ class SandboxNode extends Node
     {
         $compiler
             ->addDebugInfo($this)
-            ->write("\$sandbox = \$this->env->getExtension('\Twig\Extension\SandboxExtension');\n")
-            ->write("if (!\$alreadySandboxed = \$sandbox->isSandboxed()) {\n")
+            ->write("if (!\$alreadySandboxed = \$this->sandbox->isSandboxed()) {\n")
             ->indent()
-            ->write("\$sandbox->enableSandbox();\n")
+            ->write("\$this->sandbox->enableSandbox();\n")
             ->outdent()
             ->write("}\n")
             ->subcompile($this->getNode('body'))
             ->write("if (!\$alreadySandboxed) {\n")
             ->indent()
-            ->write("\$sandbox->disableSandbox();\n")
+            ->write("\$this->sandbox->disableSandbox();\n")
             ->outdent()
             ->write("}\n")
         ;
diff --git a/src/NodeVisitor/SandboxNodeVisitor.php b/src/NodeVisitor/SandboxNodeVisitor.php
index 5d8605480..726d87501 100644
--- a/src/NodeVisitor/SandboxNodeVisitor.php
+++ b/src/NodeVisitor/SandboxNodeVisitor.php
@@ -77,7 +77,7 @@ class SandboxNodeVisitor extends AbstractNodeVisitor
         if ($node instanceof ModuleNode) {
             $this->inAModule = false;
 
-            $node->setNode('display_start', new Node([new CheckSecurityNode($this->filters, $this->tags, $this->functions), $node->getNode('display_start')]));
+            $node->setNode('constructor_end', new Node([new CheckSecurityNode($this->filters, $this->tags, $this->functions), $node->getNode('display_start')]));
         }
 
         return $node;
diff --git a/src/Template.php b/src/Template.php
index d8ee4b53e..ee65580e4 100644
--- a/src/Template.php
+++ b/src/Template.php
@@ -39,6 +39,7 @@ abstract class Template implements \Twig_TemplateInterface
     protected $env;
     protected $blocks = [];
     protected $traits = [];
+    protected $sandbox;
 
     public function __construct(Environment $env)
     {
diff --git a/test/Twig/Tests/Node/SandboxTest.php b/test/Twig/Tests/Node/SandboxTest.php
index 9b9f7cc76..7dd3e08f0 100644
--- a/test/Twig/Tests/Node/SandboxTest.php
+++ b/test/Twig/Tests/Node/SandboxTest.php
@@ -32,13 +32,12 @@ class Twig_Tests_Node_SandboxTest extends NodeTestCase
 
         $tests[] = [$node, <<<EOF
 // line 1
-\$sandbox = \$this->env->getExtension('\Twig\Extension\SandboxExtension');
-if (!\$alreadySandboxed = \$sandbox->isSandboxed()) {
-    \$sandbox->enableSandbox();
+if (!\$alreadySandboxed = \$this->sandbox->isSandboxed()) {
+    \$this->sandbox->enableSandbox();
 }
 echo "foo";
 if (!\$alreadySandboxed) {
-    \$sandbox->disableSandbox();
+    \$this->sandbox->disableSandbox();
 }
 EOF
         ];