feature #102 [AIBundle] Add #[IsGrantedTool] for tool access control (valtzu)

This PR was merged into the main branch.

Discussion
----------

[AIBundle] Add `#[IsGrantedTool]` for tool access control

| Q             | A
| ------------- | ---
| Bug fix?      | no
| New feature?  | yes
| Docs?         | yes
| Issues        | Fix php-llm/llm-chain#360
| License       | MIT

Add `#[IsGrantedTool]` attribute for tool access control with similar behavior as `#[IsGranted]` in `symfony/security-http`.

Moved from php-llm/llm-chain#382

Commits
-------

40761da [AIBundle] Add `#[IsGrantedTool]` for tool access control

Author: chr-hertel

src/ai-bundle/composer.json

@@ -25,7 +25,9 @@
     },
     "require-dev": {
         "phpstan/phpstan": "^2.1",
-        "phpunit/phpunit": "^11.5"
+        "phpunit/phpunit": "^11.5",
+        "symfony/expression-language": "^6.4 || ^7.0",
+        "symfony/security-core": "^6.4 || ^7.0"
     },
     "config": {
         "sort-packages": true

src/ai-bundle/config/services.php

@@ -23,6 +23,7 @@
 use Symfony\AI\Agent\Toolbox\ToolResultConverter;
 use Symfony\AI\AIBundle\Profiler\DataCollector;
 use Symfony\AI\AIBundle\Profiler\TraceableToolbox;
+use Symfony\AI\AIBundle\Security\EventListener\IsGrantedToolAttributeListener;
 
 return static function (ContainerConfigurator $container): void {
     $container->services()
@@ -69,6 +70,8 @@
                 '$toolbox' => service(ToolboxInterface::class),
                 '$eventDispatcher' => service('event_dispatcher')->nullOnInvalid(),
             ])
+        ->set('symfony_ai.security.is_granted_attribute_listener', IsGrantedToolAttributeListener::class)
+            ->tag('kernel.event_listener')
 
         // profiler
         ->set(DataCollector::class)

src/ai-bundle/doc/index.rst

@@ -180,6 +180,27 @@ To inject only specific tools, list them in the configuration:
                 tools:
                     - 'Symfony\AI\Agent\Toolbox\Tool\SimilaritySearch'
 
+To restrict the access to a tool, you can use the ``IsGrantedTool`` attribute, which
+works similar to ``IsGranted`` attribute in `symfony/security-http`. For this to work,
+make sure you have `symfony/security-core` installed in your project.
+
+::
+
+    use Symfony\AI\Agent\Attribute\IsGrantedTool;
+
+    #[IsGrantedTool('ROLE_ADMIN')]
+    #[AsTool('company_name', 'Provides the name of your company')]
+    final class CompanyName
+    {
+        public function __invoke(): string
+        {
+            return 'ACME Corp.';
+        }
+    }
+The attribute ``IsGrantedTool`` can be added on class- or method-level - even multiple
+times. If multiple attributes apply to one tool call, a logical AND is used and all access
+decisions have to grant access.
+
 Profiler
 --------
 

src/ai-bundle/phpstan.dist.neon

@@ -2,3 +2,12 @@ parameters:
     level: 6
     paths:
         - src/
+    ignoreErrors:
+        -
+            message: '#\\AuthorizationCheckerInterface::isGranted\(\) invoked with 3 parameters, 1-2 required#'
+            path: src/*
+            reportUnmatched: false # only needed for Symfony <= 7.4 versions
+        -
+            message: '#method_exists\(\) with Symfony\\Component\\Security\\Core\\Exception\\AccessDeniedException and ''setAccessDecision'' will always evaluate to true#'
+            path: src/*
+            reportUnmatched: false # only needed for Symfony < 7.3 versions

src/ai-bundle/src/Security/Attribute/IsGrantedTool.php

@@ -0,0 +1,40 @@
+<?php
+
+/*
+ * This file is part of the Symfony package.
+ *
+ * (c) Fabien Potencier <[email protected]>
+ *
+ * For the full copyright and license information, please view the LICENSE
+ * file that was distributed with this source code.
+ */
+
+namespace Symfony\AI\AIBundle\Security\Attribute;
+
+use Symfony\AI\Platform\Tool\Tool;
+use Symfony\Component\ExpressionLanguage\Expression;
+
+/**
+ * Checks if user has permission to access to some tool resource using security roles and voters.
+ *
+ * @see https://symfony.com/doc/current/security.html#roles
+ *
+ * @author Valtteri R <[email protected]>
+ */
+#[\Attribute(\Attribute::IS_REPEATABLE | \Attribute::TARGET_CLASS | \Attribute::TARGET_METHOD)]
+final class IsGrantedTool
+{
+    /**
+     * @param string|Expression                                                             $attribute     The attribute that will be checked against a given authentication token and optional subject
+     * @param array<mixed>|string|Expression|\Closure(array<string,mixed>, Tool):mixed|null $subject       An optional subject - e.g. the current object being voted on
+     * @param string|null                                                                   $message       A custom message when access is not granted
+     * @param int|null                                                                      $exceptionCode If set, will add the exception code to thrown exception
+     */
+    public function __construct(
+        public string|Expression $attribute,
+        public array|string|Expression|\Closure|null $subject = null,
+        public ?string $message = null,
+        public ?int $exceptionCode = null,
+    ) {
+    }
+}

src/ai-bundle/src/Security/EventListener/IsGrantedToolAttributeListener.php

@@ -0,0 +1,112 @@
+<?php
+
+/*
+ * This file is part of the Symfony package.
+ *
+ * (c) Fabien Potencier <[email protected]>
+ *
+ * For the full copyright and license information, please view the LICENSE
+ * file that was distributed with this source code.
+ */
+
+namespace Symfony\AI\AIBundle\Security\EventListener;
+
+use Symfony\AI\Agent\Toolbox\Event\ToolCallArgumentsResolved;
+use Symfony\AI\AIBundle\Security\Attribute\IsGrantedTool;
+use Symfony\Component\ExpressionLanguage\Expression;
+use Symfony\Component\ExpressionLanguage\ExpressionLanguage;
+use Symfony\Component\Security\Core\Authorization\AccessDecision;
+use Symfony\Component\Security\Core\Authorization\AuthorizationCheckerInterface;
+use Symfony\Component\Security\Core\Exception\AccessDeniedException;
+use Symfony\Component\Security\Core\Exception\RuntimeException;
+
+/**
+ * Checks {@see IsGrantedTool} attributes on tools just before they are called.
+ *
+ * @author Valtteri R <[email protected]>
+ */
+class IsGrantedToolAttributeListener
+{
+    public function __construct(
+        private readonly AuthorizationCheckerInterface $authChecker,
+        private ?ExpressionLanguage $expressionLanguage = null,
+    ) {
+    }
+
+    public function __invoke(ToolCallArgumentsResolved $event): void
+    {
+        $tool = $event->tool;
+        $class = new \ReflectionClass($tool);
+        $method = $class->getMethod($event->metadata->reference->method);
+        $classAttributes = $class->getAttributes(IsGrantedTool::class);
+        $methodAttributes = $method->getAttributes(IsGrantedTool::class);
+
+        if (!$classAttributes && !$methodAttributes) {
+            return;
+        }
+
+        $arguments = $event->arguments;
+
+        foreach (array_merge($classAttributes, $methodAttributes) as $attr) {
+            /** @var IsGrantedTool $attribute */
+            $attribute = $attr->newInstance();
+            $subject = null;
+
+            if ($subjectRef = $attribute->subject) {
+                if (\is_array($subjectRef)) {
+                    foreach ($subjectRef as $refKey => $ref) {
+                        $subject[\is_string($refKey) ? $refKey : (string) $ref] = $this->getIsGrantedSubject($ref, $tool, $arguments);
+                    }
+                } else {
+                    $subject = $this->getIsGrantedSubject($subjectRef, $tool, $arguments);
+                }
+            }
+
+            $accessDecision = null;
+            // bc layer
+            if (class_exists(AccessDecision::class)) {
+                $accessDecision = new AccessDecision();
+                $accessDecision->isGranted = false;
+                $decision = &$accessDecision->isGranted;
+            }
+
+            if (!$decision = $this->authChecker->isGranted($attribute->attribute, $subject, $accessDecision)) {
+                $message = $attribute->message ?: (class_exists(AccessDecision::class, false) ? $accessDecision->getMessage() : 'Access Denied.');
+
+                $e = new AccessDeniedException($message, code: $attribute->exceptionCode ?? 403);
+                $e->setAttributes([$attribute->attribute]);
+                $e->setSubject($subject);
+                if ($accessDecision && method_exists($e, 'setAccessDecision')) {
+                    $e->setAccessDecision($accessDecision);
+                }
+
+                throw $e;
+            }
+        }
+    }
+
+    /**
+     * @param array<string, mixed> $arguments
+     */
+    private function getIsGrantedSubject(string|Expression|\Closure $subjectRef, object $tool, array $arguments): mixed
+    {
+        if ($subjectRef instanceof \Closure) {
+            return $subjectRef($arguments, $tool);
+        }
+
+        if ($subjectRef instanceof Expression) {
+            $this->expressionLanguage ??= new ExpressionLanguage();
+
+            return $this->expressionLanguage->evaluate($subjectRef, [
+                'tool' => $tool,
+                'args' => $arguments,
+            ]);
+        }
+
+        if (!\array_key_exists($subjectRef, $arguments)) {
+            throw new RuntimeException(\sprintf('Could not find the subject "%s" for the #[IsGranted] attribute. Try adding a "$%s" argument to your tool method.', $subjectRef, $subjectRef));
+        }
+
+        return $arguments[$subjectRef];
+    }
+}

src/ai-bundle/tests/Fixture/Tool/ToolWithIsGrantedOnClass.php

@@ -0,0 +1,23 @@
+<?php
+
+/*
+ * This file is part of the Symfony package.
+ *
+ * (c) Fabien Potencier <[email protected]>
+ *
+ * For the full copyright and license information, please view the LICENSE
+ * file that was distributed with this source code.
+ */
+
+namespace Symfony\AI\AIBundle\Tests\Fixture\Tool;
+
+use Symfony\AI\AIBundle\Security\Attribute\IsGrantedTool;
+use Symfony\Component\ExpressionLanguage\Expression;
+
+#[IsGrantedTool('test:permission', new Expression('args["itemId"] ?? 0'), message: 'No access to ToolWithIsGrantedOnClass tool.')]
+final class ToolWithIsGrantedOnClass
+{
+    public function __invoke(int $itemId): void
+    {
+    }
+}

src/ai-bundle/tests/Fixture/Tool/ToolWithIsGrantedOnMethod.php

@@ -0,0 +1,36 @@
+<?php
+
+/*
+ * This file is part of the Symfony package.
+ *
+ * (c) Fabien Potencier <[email protected]>
+ *
+ * For the full copyright and license information, please view the LICENSE
+ * file that was distributed with this source code.
+ */
+
+namespace Symfony\AI\AIBundle\Tests\Fixture\Tool;
+
+use Symfony\AI\AIBundle\Security\Attribute\IsGrantedTool;
+use Symfony\Component\ExpressionLanguage\Expression;
+
+final class ToolWithIsGrantedOnMethod
+{
+    #[IsGrantedTool('ROLE_USER', message: 'No access to simple tool.')]
+    public function simple(): bool
+    {
+        return true;
+    }
+
+    #[IsGrantedTool('test:permission', 'itemId', message: 'No access to argumentAsSubject tool.')]
+    public function argumentAsSubject(int $itemId): int
+    {
+        return $itemId;
+    }
+
+    #[IsGrantedTool('test:permission', new Expression('args["itemId"]'), message: 'No access to expressionAsSubject tool.')]
+    public function expressionAsSubject(int $itemId): int
+    {
+        return $itemId;
+    }
+}

src/ai-bundle/tests/Security/IsGrantedToolAttributeListenerTest.php

@@ -0,0 +1,92 @@
+<?php
+
+declare(strict_types=1);
+
+/*
+ * This file is part of the Symfony package.
+ *
+ * (c) Fabien Potencier <[email protected]>
+ *
+ * For the full copyright and license information, please view the LICENSE
+ * file that was distributed with this source code.
+ */
+
+namespace Symfony\AI\AIBundle\Tests\Security;
+
+use PHPUnit\Framework\Attributes\Before;
+use PHPUnit\Framework\Attributes\CoversClass;
+use PHPUnit\Framework\Attributes\Test;
+use PHPUnit\Framework\Attributes\TestWith;
+use PHPUnit\Framework\Attributes\UsesClass;
+use PHPUnit\Framework\MockObject\MockObject;
+use PHPUnit\Framework\TestCase;
+use Symfony\AI\Agent\Toolbox\Event\ToolCallArgumentsResolved;
+use Symfony\AI\AIBundle\Security\EventListener\IsGrantedToolAttributeListener;
+use Symfony\AI\AIBundle\Tests\Fixture\Tool\ToolWithIsGrantedOnClass;
+use Symfony\AI\AIBundle\Tests\Fixture\Tool\ToolWithIsGrantedOnMethod;
+use Symfony\AI\Platform\Tool\ExecutionReference;
+use Symfony\AI\Platform\Tool\Tool;
+use Symfony\Component\EventDispatcher\EventDispatcher;
+use Symfony\Component\EventDispatcher\EventDispatcherInterface;
+use Symfony\Component\ExpressionLanguage\Expression;
+use Symfony\Component\Security\Core\Authorization\AuthorizationCheckerInterface;
+use Symfony\Component\Security\Core\Exception\AccessDeniedException;
+
+#[CoversClass(IsGrantedToolAttributeListener::class)]
+#[UsesClass(EventDispatcher::class)]
+#[UsesClass(ToolCallArgumentsResolved::class)]
+#[UsesClass(Expression::class)]
+#[UsesClass(AccessDeniedException::class)]
+#[UsesClass(Tool::class)]
+#[UsesClass(ExecutionReference::class)]
+class IsGrantedToolAttributeListenerTest extends TestCase
+{
+    private EventDispatcherInterface $dispatcher;
+    private AuthorizationCheckerInterface&MockObject $authChecker;
+
+    #[Before]
+    protected function setupTool(): void
+    {
+        $this->dispatcher = new EventDispatcher();
+        $this->authChecker = $this->createMock(AuthorizationCheckerInterface::class);
+        $this->dispatcher->addListener(ToolCallArgumentsResolved::class, new IsGrantedToolAttributeListener($this->authChecker));
+    }
+
+    #[Test]
+    #[TestWith([new ToolWithIsGrantedOnMethod(), new Tool(new ExecutionReference(ToolWithIsGrantedOnMethod::class, 'simple'), 'simple', '')])]
+    #[TestWith([new ToolWithIsGrantedOnMethod(), new Tool(new ExecutionReference(ToolWithIsGrantedOnMethod::class, 'expressionAsSubject'), 'expressionAsSubject', '')])]
+    #[TestWith([new ToolWithIsGrantedOnClass(), new Tool(new ExecutionReference(ToolWithIsGrantedOnClass::class, '__invoke'), 'ToolWithIsGrantedOnClass', '')])]
+    public function itWillThrowWhenNotGranted(object $tool, Tool $metadata): void
+    {
+        $this->authChecker->expects(self::once())->method('isGranted')->willReturn(false);
+
+        self::expectException(AccessDeniedException::class);
+        self::expectExceptionMessage(\sprintf('No access to %s tool.', $metadata->name));
+        $this->dispatcher->dispatch(new ToolCallArgumentsResolved($tool, $metadata, []));
+    }
+
+    #[Test]
+    #[TestWith([new ToolWithIsGrantedOnMethod(), new Tool(new ExecutionReference(ToolWithIsGrantedOnMethod::class, 'simple'), '', '')], 'method')]
+    public function itWillNotThrowWhenGranted(object $tool, Tool $metadata): void
+    {
+        $this->authChecker->expects(self::once())->method('isGranted')->with('ROLE_USER')->willReturn(true);
+        $this->dispatcher->dispatch(new ToolCallArgumentsResolved($tool, $metadata, []));
+    }
+
+    #[Test]
+    #[TestWith([new ToolWithIsGrantedOnMethod(), new Tool(new ExecutionReference(ToolWithIsGrantedOnMethod::class, 'argumentAsSubject'), '', '')], 'method')]
+    #[TestWith([new ToolWithIsGrantedOnClass(), new Tool(new ExecutionReference(ToolWithIsGrantedOnClass::class, '__invoke'), '', '')], 'class')]
+    public function itWillProvideArgumentAsSubject(object $tool, Tool $metadata): void
+    {
+        $this->authChecker->expects(self::once())->method('isGranted')->with('test:permission', 44)->willReturn(true);
+        $this->dispatcher->dispatch(new ToolCallArgumentsResolved($tool, $metadata, ['itemId' => 44]));
+    }
+
+    #[Test]
+    #[TestWith([new ToolWithIsGrantedOnMethod(), new Tool(new ExecutionReference(ToolWithIsGrantedOnMethod::class, 'expressionAsSubject'), '', '')], 'method')]
+    public function itWillEvaluateSubjectExpression(object $tool, Tool $metadata): void
+    {
+        $this->authChecker->expects(self::once())->method('isGranted')->with('test:permission', 44)->willReturn(true);
+        $this->dispatcher->dispatch(new ToolCallArgumentsResolved($tool, $metadata, ['itemId' => 44]));
+    }
+}