Fix/quote identifiers

.gitignore

@@ -82,3 +82,5 @@ node_modules/
 docs/.vuepress/.cache/
 docs/.vuepress/.temp/
 docs/.vuepress/dist/
+
+.human/

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,

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,

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() {

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,14 +270,17 @@ 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(())
     }
 
     /// 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,14 +290,17 @@ 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(())
     }
 
     /// 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(())

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'");
+    }
+}