feat: Add #[IsGrantedTool] for tool access control

composer.json

@@ -58,9 +58,11 @@
         "symfony/dom-crawler": "^6.4 || ^7.1",
         "symfony/dotenv": "^6.4 || ^7.1",
         "symfony/event-dispatcher": "^6.4 || ^7.1",
+        "symfony/expression-language": "^6.4 || ^7.1",
         "symfony/finder": "^6.4 || ^7.1",
         "symfony/http-foundation": "^6.4 || ^7.1",
         "symfony/process": "^6.4 || ^7.1",
+        "symfony/security-core": "^6.4 || ^7.1",
         "symfony/var-dumper": "^6.4 || ^7.1"
     },
     "suggest": {
@@ -73,7 +75,9 @@
         "mrmysql/youtube-transcript": "For using the YouTube transcription tool.",
         "probots-io/pinecone-php": "For using the Pinecone as retrieval vector store.",
         "symfony/dom-crawler": "For using the Crawler tool.",
+        "symfony/expression-language": "For using Expressions with #[IsGranted] attribute",
         "symfony/http-foundation": "For using the SessionStore as message store.",
+        "symfony/security-core": "For using #[IsGranted] attribute for tool access control",
         "psr/cache": "For using the CacheStore as message store."
     },
     "config": {

src/Chain/Toolbox/Security/Attribute/IsGrantedTool.php

@@ -0,0 +1,31 @@
+<?php
+
+namespace PhpLlm\LlmChain\Chain\Toolbox\Security\Attribute;
+
+use PhpLlm\LlmChain\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/Chain/Toolbox/Security/EventListener/IsGrantedToolAttributeListener.php

@@ -0,0 +1,105 @@
+<?php
+
+namespace PhpLlm\LlmChain\Chain\Toolbox\Security\EventListener;
+
+use PhpLlm\LlmChain\Chain\Toolbox\Event\ToolCallArgumentsResolved;
+use PhpLlm\LlmChain\Chain\Toolbox\Security\Attribute\IsGrantedTool;
+use Symfony\Component\EventDispatcher\Attribute\AsEventListener;
+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]>
+ */
+#[AsEventListener]
+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 ?: $accessDecision->getMessage();
+
+                $e = new AccessDeniedException($message, code: $attribute->exceptionCode ?? 403);
+                $e->setAttributes([$attribute->attribute]);
+                $e->setSubject($subject);
+                if ($accessDecision) {
+                    $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];
+    }
+}

tests/Chain/Toolbox/Security/IsGrantedAttributeListenerTest.php

@@ -0,0 +1,83 @@
+<?php
+
+declare(strict_types=1);
+
+namespace PhpLlm\LlmChain\Tests\Chain\Toolbox\Security;
+
+use PhpLlm\LlmChain\Chain\Toolbox\Event\ToolCallArgumentsResolved;
+use PhpLlm\LlmChain\Chain\Toolbox\Security\EventListener\IsGrantedToolAttributeListener;
+use PhpLlm\LlmChain\Platform\Tool\ExecutionReference;
+use PhpLlm\LlmChain\Platform\Tool\Tool;
+use PhpLlm\LlmChain\Tests\Fixture\Tool\ToolWithIsGrantedOnClass;
+use PhpLlm\LlmChain\Tests\Fixture\Tool\ToolWithIsGrantedOnMethod;
+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\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 IsGrantedAttributeListenerTest 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]));
+    }
+}

tests/Fixture/Tool/ToolWithIsGrantedOnClass.php

@@ -0,0 +1,14 @@
+<?php
+
+namespace PhpLlm\LlmChain\Tests\Fixture\Tool;
+
+use PhpLlm\LlmChain\Chain\Toolbox\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
+    {
+    }
+}

tests/Fixture/Tool/ToolWithIsGrantedOnMethod.php

@@ -0,0 +1,27 @@
+<?php
+
+namespace PhpLlm\LlmChain\Tests\Fixture\Tool;
+
+use PhpLlm\LlmChain\Chain\Toolbox\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;
+    }
+}