Merge pull request #12 from run-as-root/fix/security-audit-remediations

fix(security): comprehensive security audit remediations

Author: DavidLambauer

Controller/Adminhtml/Assistant/Chat.php

@@ -7,6 +7,7 @@
 use Magento\Backend\App\Action;
 use Magento\Backend\App\Action\Context;
 use Magento\Framework\App\Action\HttpPostActionInterface;
+use Magento\Framework\App\CacheInterface;
 use Magento\Framework\Controller\Result\Json;
 use Magento\Framework\Controller\Result\JsonFactory;
 use Psr\Log\LoggerInterface;
@@ -17,12 +18,16 @@ class Chat extends Action implements HttpPostActionInterface
 {
     public const ADMIN_RESOURCE = 'RunAsRoot_TypeSense::ai_assistant';
 
+    private const RATE_LIMIT_MAX_REQUESTS = 100;
+    private const RATE_LIMIT_WINDOW_SECONDS = 3600;
+
     public function __construct(
         Context $context,
         private readonly JsonFactory $jsonFactory,
         private readonly AgentLoop $agentLoop,
         private readonly LoggerInterface $logger,
         private readonly TypeSenseConfigInterface $config,
+        private readonly CacheInterface $cache,
     ) {
         parent::__construct($context);
     }
@@ -35,6 +40,11 @@ public function execute(): Json
             return $result->setData(['success' => false, 'error' => 'AI Assistant is not enabled.']);
         }
 
+        $rateLimitError = $this->checkRateLimit();
+        if ($rateLimitError !== null) {
+            return $result->setData(['success' => false, 'error' => $rateLimitError]);
+        }
+
         try {
             $query = (string) $this->getRequest()->getParam('query', '');
             $historyJson = (string) $this->getRequest()->getParam('history', '');
@@ -47,7 +57,7 @@ public function execute(): Json
             if ($historyJson !== '') {
                 $decoded = json_decode($historyJson, true);
                 if (is_array($decoded)) {
-                    $history = $decoded;
+                    $history = $this->sanitizeHistory($decoded);
                 }
             }
 
@@ -56,7 +66,7 @@ public function execute(): Json
             return $result->setData([
                 'success' => true,
                 'answer' => $response['answer'],
-                'messages' => $response['messages'],
+                'messages' => $this->filterResponseMessages($response['messages']),
             ]);
         } catch (\Exception $e) {
             $this->logger->error('Admin AI Assistant error: ' . $e->getMessage());
@@ -67,4 +77,84 @@ public function execute(): Json
             ]);
         }
     }
+
+    private function checkRateLimit(): ?string
+    {
+        $adminId = (string) $this->_auth->getUser()->getId();
+        $cacheKey = 'ai_assistant_rate_' . $adminId;
+
+        $count = (int) $this->cache->load($cacheKey);
+
+        if ($count >= self::RATE_LIMIT_MAX_REQUESTS) {
+            $this->logger->warning('AI Assistant rate limit exceeded for admin ' . $adminId);
+            return 'Rate limit exceeded. Maximum ' . self::RATE_LIMIT_MAX_REQUESTS
+                . ' requests per hour. Please try again later.';
+        }
+
+        $this->cache->save(
+            (string) ($count + 1),
+            $cacheKey,
+            [],
+            self::RATE_LIMIT_WINDOW_SECONDS
+        );
+
+        return null;
+    }
+
+    /**
+     * Filter response messages to prevent leaking system prompt and tool internals.
+     *
+     * @param array<int, array<string, mixed>> $messages
+     * @return array<int, array{role: string, content: string}>
+     */
+    private function filterResponseMessages(array $messages): array
+    {
+        $exposedRoles = ['user', 'assistant'];
+        $filtered = [];
+
+        foreach ($messages as $message) {
+            $role = (string) ($message['role'] ?? '');
+            if (!in_array($role, $exposedRoles, true)) {
+                continue;
+            }
+
+            $filtered[] = [
+                'role' => $role,
+                'content' => (string) ($message['content'] ?? ''),
+            ];
+        }
+
+        return $filtered;
+    }
+
+    /**
+     * Sanitize conversation history from client input.
+     * Only allow user and assistant roles. Strip tool_calls and system messages.
+     *
+     * @param array<int, mixed> $history
+     * @return array<int, array{role: string, content: string}>
+     */
+    private function sanitizeHistory(array $history): array
+    {
+        $allowedRoles = ['user', 'assistant'];
+        $sanitized = [];
+
+        foreach ($history as $message) {
+            if (!is_array($message)) {
+                continue;
+            }
+
+            $role = (string) ($message['role'] ?? '');
+            if (!in_array($role, $allowedRoles, true)) {
+                continue;
+            }
+
+            $sanitized[] = [
+                'role' => $role,
+                'content' => (string) ($message['content'] ?? ''),
+            ];
+        }
+
+        return $sanitized;
+    }
 }

Controller/Adminhtml/CategoryMerchandiser/Save.php

@@ -7,7 +7,6 @@
 use Magento\Backend\App\Action;
 use Magento\Backend\App\Action\Context;
 use Magento\Framework\Api\SearchCriteriaBuilder;
-use Magento\Framework\App\Action\HttpGetActionInterface;
 use Magento\Framework\App\Action\HttpPostActionInterface;
 use Magento\Framework\Controller\Result\Json;
 use Magento\Framework\Controller\Result\JsonFactory;
@@ -17,7 +16,7 @@
 use RunAsRoot\TypeSense\Model\Curation\CategoryMerchandisingSync;
 use RunAsRoot\TypeSense\Model\Merchandising\CategoryMerchandisingFactory;
 
-class Save extends Action implements HttpPostActionInterface, HttpGetActionInterface
+class Save extends Action implements HttpPostActionInterface
 {
     public const ADMIN_RESOURCE = 'RunAsRoot_TypeSense::overrides';
 
@@ -38,45 +37,6 @@ public function execute(): Json
     {
         $resultJson = $this->jsonFactory->create();
 
-        // Handle GET: return existing rules for a category
-        if ($this->getRequest()->isGet()) {
-            return $this->handleGet($resultJson);
-        }
-
-        return $this->handlePost($resultJson);
-    }
-
-    private function handleGet(Json $resultJson): Json
-    {
-        $categoryId = (int) $this->getRequest()->getParam('category_id', 0);
-
-        if ($categoryId === 0) {
-            return $resultJson->setData(['rules' => []]);
-        }
-
-        $searchCriteria = $this->searchCriteriaBuilder
-            ->addFilter('category_id', $categoryId)
-            ->create();
-
-        $searchResults = $this->repository->getList($searchCriteria);
-        $rules = [];
-
-        foreach ($searchResults->getItems() as $item) {
-            $rules[] = [
-                'product_id' => $item->getProductId(),
-                'position'   => $item->getPosition(),
-                'action'     => $item->getAction(),
-                'name'       => '',
-                'sku'        => '',
-                'image_url'  => '',
-            ];
-        }
-
-        return $resultJson->setData(['rules' => $rules]);
-    }
-
-    private function handlePost(Json $resultJson): Json
-    {
         try {
             $raw = $this->getRequest()->getParam('payload') ?: $this->getRequest()->getContent();
             $payload = json_decode((string) $raw, true, 512, JSON_THROW_ON_ERROR);

Model/Assistant/AgentLoop.php

@@ -10,6 +10,7 @@
 class AgentLoop
 {
     private const MAX_ITERATIONS = 10;
+    private const MAX_EXECUTION_SECONDS = 60;
 
     public function __construct(
         private readonly OpenAiClientFactory $openAiClientFactory,
@@ -41,8 +42,19 @@ public function run(string $userQuery, array $conversationHistory = []): array
 
         $tools = $this->toolRegistry->getToolDefinitions();
         $client = $this->openAiClientFactory->create();
+        $startTime = microtime(true);
 
         for ($i = 0; $i < self::MAX_ITERATIONS; $i++) {
+            if ((microtime(true) - $startTime) > self::MAX_EXECUTION_SECONDS) {
+                $this->logger->warning('AI Assistant agent loop timed out', [
+                    'elapsed_seconds' => round(microtime(true) - $startTime, 1),
+                    'iterations' => $i,
+                ]);
+                return [
+                    'answer' => 'The request took too long to process. Please try a simpler question.',
+                    'messages' => $messages,
+                ];
+            }
             $response = $client->chat()->create([
                 'model' => $this->resolveModel(),
                 'messages' => $messages,

Model/Assistant/Tool/DescribeDatabaseTool.php

@@ -100,7 +100,7 @@ private function describeTable(string $tableName): string
 
             return json_encode(['table' => $tableName, 'columns' => $columns]);
         } catch (\Exception $e) {
-            return json_encode(['error' => 'Could not describe table: ' . $e->getMessage()]);
+            return json_encode(['error' => 'Could not describe table. Please check the table name and try again.']);
         }
     }
 
@@ -143,7 +143,7 @@ private function showRelationships(string $tableName): string
                 'referenced_by' => $incoming,
             ]);
         } catch (\Exception $e) {
-            return json_encode(['error' => 'Could not fetch relationships: ' . $e->getMessage()]);
+            return json_encode(['error' => 'Could not fetch relationships. Please check the table name and try again.']);
         }
     }
 }

Model/Assistant/Tool/ExecuteSqlTool.php

@@ -5,6 +5,7 @@
 namespace RunAsRoot\TypeSense\Model\Assistant\Tool;
 
 use Magento\Framework\App\ResourceConnection;
+use Psr\Log\LoggerInterface;
 
 class ExecuteSqlTool implements ToolInterface
 {
@@ -13,6 +14,7 @@ class ExecuteSqlTool implements ToolInterface
     public function __construct(
         private readonly ResourceConnection $resource,
         private readonly SqlSandbox $sandbox,
+        private readonly LoggerInterface $logger,
     ) {
     }
 
@@ -23,7 +25,7 @@ public function getName(): string
 
     public function getDescription(): string
     {
-        return 'Execute a read-only SELECT query on the Magento MySQL database. Returns up to 100 rows. Use describe_database to explore table/column names first. Blocked tables: admin_user, oauth_token, authorization_role. Blocked columns: password_hash.';
+        return 'Execute a read-only SELECT query on the Magento MySQL database. Returns up to 100 rows. Use describe_database to explore table/column names first. Many tables and columns are blocked for security (admin_user, oauth_token, authorization_role, integration, vault_payment_token, etc). UNION, subqueries into blocked tables, and SQL comments are not allowed.';
     }
 
     public function getParametersSchema(): array
@@ -48,9 +50,15 @@ public function execute(array $arguments): string
             return json_encode(['error' => 'Query cannot be empty.']);
         }
 
+        $this->logger->info('AI Assistant SQL query', ['query' => mb_substr($query, 0, 500)]);
+
         try {
             $this->sandbox->validate($query);
         } catch (\InvalidArgumentException $e) {
+            $this->logger->warning('AI Assistant blocked SQL query', [
+                'query' => mb_substr($query, 0, 500),
+                'reason' => $e->getMessage(),
+            ]);
             return json_encode(['error' => $e->getMessage()]);
         }
 
@@ -73,7 +81,11 @@ public function execute(array $arguments): string
                 'count' => count($rows),
             ]);
         } catch (\Exception $e) {
-            return json_encode(['error' => 'SQL error: ' . $e->getMessage()]);
+            $this->logger->error('AI Assistant SQL execution error', [
+                'query' => mb_substr($query, 0, 500),
+                'error' => $e->getMessage(),
+            ]);
+            return json_encode(['error' => 'Query execution failed. Please check your SQL syntax and try again.']);
         }
     }
 }