From 789bfab7218da396329ca0966ee5a2160db8d087 Mon Sep 17 00:00:00 2001 From: Bob Pankratz Date: Mon, 13 Jul 2026 18:18:29 -0500 Subject: [PATCH] Validate dataset primary key identifiers --- .../InMemoryDatasetStore.swift | 3 +++ .../PostgreSQLDatasetStore.swift | 3 +++ .../SQLiteDatasetStore.swift | 3 +++ .../DatasetStoreTests.swift | 16 ++++++++++++++ .../PersistenceKit/rust/src/dataset_store.rs | 3 +++ .../kits/PersistenceKit/rust/src/postgres.rs | 3 +++ .../kits/PersistenceKit/rust/src/sqlite.rs | 22 +++++++++++++++++++ 7 files changed, 53 insertions(+) diff --git a/packages/kits/PersistenceKit/Sources/PersistenceKitInMemory/InMemoryDatasetStore.swift b/packages/kits/PersistenceKit/Sources/PersistenceKitInMemory/InMemoryDatasetStore.swift index 91fbd32..eb4a48c 100644 --- a/packages/kits/PersistenceKit/Sources/PersistenceKitInMemory/InMemoryDatasetStore.swift +++ b/packages/kits/PersistenceKit/Sources/PersistenceKitInMemory/InMemoryDatasetStore.swift @@ -63,6 +63,9 @@ final class InMemoryDatasetStore: DatasetStore, Sendable { for col in schema.columns { try validateDatasetColumnIdentifier(col.name) } + if let primaryKeyColumn = schema.primaryKeyColumn { + try validateDatasetColumnIdentifier(primaryKeyColumn) + } for idx in indexes { try validateDatasetColumnIdentifier(idx.column) } diff --git a/packages/kits/PersistenceKit/Sources/PersistenceKitPostgreSQL/PostgreSQLDatasetStore.swift b/packages/kits/PersistenceKit/Sources/PersistenceKitPostgreSQL/PostgreSQLDatasetStore.swift index e5c1f22..0b6e465 100644 --- a/packages/kits/PersistenceKit/Sources/PersistenceKitPostgreSQL/PostgreSQLDatasetStore.swift +++ b/packages/kits/PersistenceKit/Sources/PersistenceKitPostgreSQL/PostgreSQLDatasetStore.swift @@ -87,6 +87,9 @@ final class PostgreSQLDatasetStore: DatasetStore, Sendable { for col in schema.columns { try validateDatasetColumnIdentifier(col.name) } + if let primaryKeyColumn = schema.primaryKeyColumn { + try validateDatasetColumnIdentifier(primaryKeyColumn) + } for idx in indexes { try validateDatasetColumnIdentifier(idx.column) } diff --git a/packages/kits/PersistenceKit/Sources/PersistenceKitSQLite/SQLiteDatasetStore.swift b/packages/kits/PersistenceKit/Sources/PersistenceKitSQLite/SQLiteDatasetStore.swift index 8d200d1..213ab93 100644 --- a/packages/kits/PersistenceKit/Sources/PersistenceKitSQLite/SQLiteDatasetStore.swift +++ b/packages/kits/PersistenceKit/Sources/PersistenceKitSQLite/SQLiteDatasetStore.swift @@ -63,6 +63,9 @@ final class SQLiteDatasetStore: DatasetStore, Sendable { for col in schema.columns { try validateDatasetColumnIdentifier(col.name) } + if let primaryKeyColumn = schema.primaryKeyColumn { + try validateDatasetColumnIdentifier(primaryKeyColumn) + } for idx in indexes { try validateDatasetColumnIdentifier(idx.column) } diff --git a/packages/kits/PersistenceKit/Tests/PersistenceKitDatasetTests/DatasetStoreTests.swift b/packages/kits/PersistenceKit/Tests/PersistenceKitDatasetTests/DatasetStoreTests.swift index 51dee75..562a5d3 100644 --- a/packages/kits/PersistenceKit/Tests/PersistenceKitDatasetTests/DatasetStoreTests.swift +++ b/packages/kits/PersistenceKit/Tests/PersistenceKitDatasetTests/DatasetStoreTests.swift @@ -181,6 +181,22 @@ struct DatasetStoreTests { } } + /// SQL-injection attempt: double-quote in primary-key column name. + @Test func invalidIdentifier_primaryKeyColumn_sqlInjection_doubleQuote() async { + let ds = makeInMemory().datasetStore + let badName = "id\")); DROP TABLE dataset; --" + let schema = DatasetSchema( + columns: [ + ColumnDeclaration(name: "id", type: .integer, nullable: false), + ColumnDeclaration(name: "value", type: .text, nullable: true) + ], + primaryKeyColumn: badName + ) + await #expect(throws: StorageError.self) { + try await ds.createDataset(id: UUID(), schema: schema, indexes: []) + } + } + /// Rejected column name in `appendRows` key must fail the whole /// operation before any rows are written. @Test func invalidIdentifier_inAppendRows() async throws { diff --git a/packages/kits/PersistenceKit/rust/src/dataset_store.rs b/packages/kits/PersistenceKit/rust/src/dataset_store.rs index 07ee6d9..5edbd69 100644 --- a/packages/kits/PersistenceKit/rust/src/dataset_store.rs +++ b/packages/kits/PersistenceKit/rust/src/dataset_store.rs @@ -334,6 +334,9 @@ impl DatasetStore for InMemoryDatasetStore { for col in &schema.columns { validate_dataset_column_identifier(&col.name)?; } + if let Some(pk) = &schema.primary_key_column { + validate_dataset_column_identifier(pk)?; + } for idx in indexes { validate_dataset_column_identifier(&idx.column)?; } diff --git a/packages/kits/PersistenceKit/rust/src/postgres.rs b/packages/kits/PersistenceKit/rust/src/postgres.rs index 3a9d032..f0d0809 100644 --- a/packages/kits/PersistenceKit/rust/src/postgres.rs +++ b/packages/kits/PersistenceKit/rust/src/postgres.rs @@ -2833,6 +2833,9 @@ impl DatasetStore for PgDatasetStore { for col in &schema.columns { validate_dataset_column_identifier(&col.name)?; } + if let Some(pk) = &schema.primary_key_column { + validate_dataset_column_identifier(pk)?; + } for idx in indexes { validate_dataset_column_identifier(&idx.column)?; } diff --git a/packages/kits/PersistenceKit/rust/src/sqlite.rs b/packages/kits/PersistenceKit/rust/src/sqlite.rs index e8b14d2..0aaef2b 100644 --- a/packages/kits/PersistenceKit/rust/src/sqlite.rs +++ b/packages/kits/PersistenceKit/rust/src/sqlite.rs @@ -2449,6 +2449,9 @@ impl DatasetStore for SqliteDatasetStoreShim { for col in &schema.columns { validate_dataset_column_identifier(&col.name)?; } + if let Some(pk) = &schema.primary_key_column { + validate_dataset_column_identifier(pk)?; + } for idx in indexes { validate_dataset_column_identifier(&idx.column)?; } @@ -4140,6 +4143,25 @@ mod sql_identifier_injection_tests { ); } + #[test] + fn create_dataset_rejects_primary_key_column_with_double_quote() { + let storage = injection_test_storage(); + let ds = Storage::dataset_store(&storage).expect("dataset store"); + let schema = DatasetSchema { + columns: vec![ + ColumnDeclaration::new("id", ColumnType::Int), + ColumnDeclaration::new("value", ColumnType::Text), + ], + primary_key_column: Some("id\")); DROP TABLE items; --".to_string()), + }; + + let err = ds.create_dataset(Uuid::new_v4(), &schema, &[]).unwrap_err(); + assert!( + matches!(err, StorageError::InvalidIdentifier { .. }), + "expected InvalidIdentifier, got: {err:?}" + ); + } + /// A semicolon in a column name allows statement stacking — classic /// injection vector. Reject before the SQL string is built. #[test]