diff --git a/CHANGELOG.md b/CHANGELOG.md index fa54cf273..2fd881936 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,8 @@ # Changelog - 2.5.13 + - Fix SQL Table offset syntax to work with Postgres + - Prevent SQL Table injection attack - Pipeline add missing traided guard on score method - 2.5.12 diff --git a/src/Extractors/SQLTable.php b/src/Extractors/SQLTable.php index 50359addf..e68d97080 100644 --- a/src/Extractors/SQLTable.php +++ b/src/Extractors/SQLTable.php @@ -11,6 +11,7 @@ use function Rubix\ML\iterator_first; use function count; use function array_keys; +use function preg_match; /** * SQL Table @@ -25,6 +26,8 @@ */ class SQLTable implements Extractor { + protected const TABLE_NAME_PATTERN = '/^[A-Za-z_][A-Za-z0-9_]*(\.[A-Za-z_][A-Za-z0-9_]*)?$/'; + /** * The PDO connection to the database. * @@ -60,13 +63,17 @@ public function __construct(PDO $connection, string $table, int $batchSize = 256 throw new InvalidArgumentException('Table name cannot be empty.'); } + if (!preg_match(self::TABLE_NAME_PATTERN, $table)) { + throw new InvalidArgumentException("Table name '{$table}' is not a valid identifier."); + } + if ($batchSize < 1) { throw new InvalidArgumentException('Batch size must be' . " greater than 0, $batchSize given."); } $this->connection = $connection; - $this->table = $connection->quote($table); + $this->table = $table; $this->batchSize = $batchSize; } @@ -87,7 +94,7 @@ public function header() : array */ public function getIterator() : Traversable { - $query = "SELECT * FROM {$this->table} LIMIT :offset, {$this->batchSize}"; + $query = "SELECT * FROM {$this->table} LIMIT {$this->batchSize} OFFSET :offset"; $statement = $this->connection->prepare($query); diff --git a/tests/Extractors/SQLTableTest.php b/tests/Extractors/SQLTableTest.php index 38832c8fb..f54de6ea0 100644 --- a/tests/Extractors/SQLTableTest.php +++ b/tests/Extractors/SQLTableTest.php @@ -4,6 +4,7 @@ use Rubix\ML\Extractors\SQLTable; use Rubix\ML\Extractors\Extractor; +use Rubix\ML\Exceptions\InvalidArgumentException; use PHPUnit\Framework\TestCase; use IteratorAggregate; use Traversable; @@ -68,4 +69,16 @@ public function extract() : void $this->assertEquals($expected, $header); } + + /** + * @test + */ + public function rejectInvalidIdentifier() : void + { + $connection = new PDO('sqlite::memory:'); + + $this->expectException(InvalidArgumentException::class); + + new SQLTable($connection, "pets'; DROP TABLE users; --"); + } }