diff --git a/.gitignore b/.gitignore index 5fd6dfec..cd902dad 100644 --- a/.gitignore +++ b/.gitignore @@ -82,3 +82,5 @@ node_modules/ docs/.vuepress/.cache/ docs/.vuepress/.temp/ docs/.vuepress/dist/ + +.human/ diff --git a/python/tests/test_listener.py b/python/tests/test_listener.py index 5fcf2c0e..8e4b1b20 100644 --- a/python/tests/test_listener.py +++ b/python/tests/test_listener.py @@ -206,6 +206,46 @@ async def trigger() -> None: await listener.shutdown() +@pytest.mark.usefixtures("create_table_for_listener_tests") +async def test_listener_mixed_case_channel( + psql_pool: ConnectionPool, + listener_table_name: str, +) -> None: + """Test that a channel name is subscribed to verbatim, without case folding. + + An unquoted ``LISTEN MixedCase`` is folded by the backend to ``mixedcase``, + so the delivered notification would carry a channel name that never matches + the key the callback was registered under. + """ + channel = "MixedCaseChannel" + listener = psql_pool.listener() + await listener.add_callback( + channel=channel, + callback=construct_insert_callback( + listener_table_name=listener_table_name, + ), + ) + await listener.startup() + listener.listen() + + await wait_until_listening(listener, channel) + + connection = await psql_pool.connection() + try: + await connection.execute(f"NOTIFY \"{channel}\", '{TEST_PAYLOAD}'") + finally: + connection.close() + + rows = await wait_for_callback( + psql_pool=psql_pool, + listener_table_name=listener_table_name, + ) + assert rows[0]["channel"] == channel + assert rows[0]["payload"] == TEST_PAYLOAD + + await listener.shutdown() + + @pytest.mark.usefixtures("create_table_for_listener_tests") async def test_listener_abort( psql_pool: ConnectionPool, diff --git a/python/tests/test_transaction.py b/python/tests/test_transaction.py index 81dc7e2c..6f8f9036 100644 --- a/python/tests/test_transaction.py +++ b/python/tests/test_transaction.py @@ -209,6 +209,55 @@ async def test_transaction_release_savepoint( await transaction.create_savepoint(sp_name_1) +async def test_transaction_savepoint_name_is_quoted( + psql_pool: ConnectionPool, + table_name: str, +) -> None: + """Test that a savepoint name cannot smuggle in a second statement.""" + connection = await psql_pool.connection() + transaction = connection.transaction() + await transaction.begin() + + savepoint_name = f"sp1; DROP TABLE {table_name}" + await transaction.create_savepoint(savepoint_name=savepoint_name) + + result = await transaction.execute( + "SELECT to_regclass($1) IS NOT NULL AS exists", + parameters=[table_name], + ) + assert result.result()[0]["exists"] + + # The quoted name still identifies a real savepoint. + await transaction.rollback_savepoint(savepoint_name=savepoint_name) + await transaction.release_savepoint(savepoint_name=savepoint_name) + + await transaction.commit() + + +async def test_transaction_savepoint_name_with_quotes( + psql_pool: ConnectionPool, + table_name: str, +) -> None: + """Test that a savepoint name containing a double quote is escaped.""" + connection = await psql_pool.connection() + transaction = connection.transaction() + await transaction.begin() + + rows_before = await count_rows_in_test_table(table_name, transaction) + + savepoint_name = 'we"ird name' + await transaction.create_savepoint(savepoint_name=savepoint_name) + await transaction.execute( + f"INSERT INTO {table_name} VALUES ($1, $2)", + parameters=[100, "test_name"], + ) + await transaction.rollback_savepoint(savepoint_name=savepoint_name) + + assert await count_rows_in_test_table(table_name, transaction) == rows_before + + await transaction.commit() + + async def test_transaction_cursor( psql_pool: ConnectionPool, table_name: str, diff --git a/src/driver/listener/core.rs b/src/driver/listener/core.rs index 28fcbef2..b0f61fcd 100644 --- a/src/driver/listener/core.rs +++ b/src/driver/listener/core.rs @@ -21,6 +21,7 @@ use crate::{ utils::{build_tls, is_coroutine_function, ConfiguredTLS}, }, exceptions::rust_errors::{PSQLPyResult, RustPSQLDriverError}, + format_helpers::quote_ident, options::SslMode, runtime::{rustdriver_future, tokio_runtime}, }; @@ -382,6 +383,11 @@ async fn dispatch_callback( /// and executed. Re-subscribing is idempotent, so a redundant `LISTEN` is /// harmless; the `UNLISTEN` half is what stops a cleared channel from delivering. /// +/// Channel names are quoted as identifiers. Besides closing the injection hole, +/// this is what makes a mixed-case channel work at all: unquoted `LISTEN MyChan` +/// is folded by the backend to `mychan`, so the incoming notification would carry +/// a channel name that never matches the key the callback was registered under. +/// /// Lock order is `client` -> `is_listened` -> `channel_callbacks` -> /// `applied_channels`. `mark_subscriptions_dirty` only ever takes `is_listened` /// (never while holding `channel_callbacks`), so the two cannot deadlock. @@ -408,10 +414,10 @@ async fn execute_listen( let mut reconcile_query = String::new(); for channel in applied.difference(&desired) { - reconcile_query.push_str(format!("UNLISTEN {channel};").as_str()); + reconcile_query.push_str(format!("UNLISTEN {};", quote_ident(channel)).as_str()); } for channel in desired.difference(&applied) { - reconcile_query.push_str(format!("LISTEN {channel};").as_str()); + reconcile_query.push_str(format!("LISTEN {};", quote_ident(channel)).as_str()); } if !reconcile_query.is_empty() { diff --git a/src/driver/transaction.rs b/src/driver/transaction.rs index 5d9ff9cf..c30e4a83 100644 --- a/src/driver/transaction.rs +++ b/src/driver/transaction.rs @@ -15,6 +15,7 @@ use crate::{ traits::{CloseTransaction, Connection, StartTransaction as _}, }, exceptions::rust_errors::{PSQLPyResult, RustPSQLDriverError}, + format_helpers::quote_ident, options::{IsolationLevel, ReadVariant}, query_result::{PSQLDriverPyQueryResult, PSQLDriverSinglePyQueryResult}, }; @@ -257,6 +258,9 @@ impl Transaction { /// Create new savepoint in a transaction. /// + /// Savepoint name is quoted as an identifier, so it is safe to pass + /// a name that comes from the outside. + /// /// # Errors /// Can return error if there is a problem with DB communication. pub async fn create_savepoint(&mut self, savepoint_name: String) -> PSQLPyResult<()> { @@ -266,7 +270,7 @@ impl Transaction { let read_conn_g = conn.read().await; read_conn_g - .batch_execute(format!("SAVEPOINT {savepoint_name}").as_str()) + .batch_execute(format!("SAVEPOINT {}", quote_ident(&savepoint_name)).as_str()) .await?; Ok(()) @@ -274,6 +278,9 @@ impl Transaction { /// Release a savepoint in a transaction. /// + /// Savepoint name is quoted as an identifier, so it is safe to pass + /// a name that comes from the outside. + /// /// # Errors /// Can return error if there is a problem with DB communication. pub async fn release_savepoint(&mut self, savepoint_name: String) -> PSQLPyResult<()> { @@ -283,7 +290,7 @@ impl Transaction { let read_conn_g = conn.read().await; read_conn_g - .batch_execute(format!("RELEASE SAVEPOINT {savepoint_name}").as_str()) + .batch_execute(format!("RELEASE SAVEPOINT {}", quote_ident(&savepoint_name)).as_str()) .await?; Ok(()) @@ -291,6 +298,9 @@ impl Transaction { /// Rollback to a savepoint in a transaction. /// + /// Savepoint name is quoted as an identifier, so it is safe to pass + /// a name that comes from the outside. + /// /// # Errors /// Can return error if there is a problem with DB communication. pub async fn rollback_savepoint(&mut self, savepoint_name: String) -> PSQLPyResult<()> { @@ -300,7 +310,9 @@ impl Transaction { let read_conn_g = conn.read().await; read_conn_g - .batch_execute(format!("ROLLBACK TO SAVEPOINT {savepoint_name}").as_str()) + .batch_execute( + format!("ROLLBACK TO SAVEPOINT {}", quote_ident(&savepoint_name)).as_str(), + ) .await?; Ok(()) diff --git a/src/format_helpers.rs b/src/format_helpers.rs index 4aed5d91..98cd49b1 100644 --- a/src/format_helpers.rs +++ b/src/format_helpers.rs @@ -7,3 +7,33 @@ pub fn quote_ident(ident: &str) -> String { pub fn quote_literal(string: &str) -> String { format!("'{}'", string.replace('\'', "''")) } + +#[cfg(test)] +mod tests { + use super::{quote_ident, quote_literal}; + + #[test] + fn quote_ident_wraps_in_double_quotes() { + assert_eq!(quote_ident("users"), "\"users\""); + } + + #[test] + fn quote_ident_doubles_inner_quotes() { + assert_eq!(quote_ident("we\"ird"), "\"we\"\"ird\""); + } + + #[test] + fn quote_ident_neutralizes_statement_separator() { + // Everything after the name stays inside the identifier and cannot + // start a new statement. + assert_eq!( + quote_ident("sp1; DROP TABLE users"), + "\"sp1; DROP TABLE users\"" + ); + } + + #[test] + fn quote_literal_doubles_single_quotes() { + assert_eq!(quote_literal("O'Reilly"), "'O''Reilly'"); + } +}