fix: handle connection's ssl check on init (#15)

Author: sweb

src/intercept.rs

@@ -24,6 +24,23 @@ fn single_text_response(column_name: &str, value: &str) -> PgWireResult<Vec<Resp
     ))])
 }
 
+/// Build a single-column, single-row BOOL text response (`t`/`f` on the wire).
+/// Used for the `pg_stat_ssl` probe, whose `ssl` column is a boolean; a
+/// VARCHAR response would carry the wrong type oid in the RowDescription and
+/// Npgsql decodes the column strictly by that oid.
+fn single_bool_response(column_name: &str, value: bool) -> PgWireResult<Vec<Response>> {
+    let fields = Arc::new(vec![text_field(column_name, Type::BOOL)]);
+
+    let mut encoder = DataRowEncoder::new(Arc::clone(&fields));
+    encoder.encode_field(&value)?;
+    let row = encoder.take_row();
+
+    Ok(vec![Response::Query(QueryResponse::new(
+        fields,
+        stream::iter(vec![Ok(row)]),
+    ))])
+}
+
 /// If the query can be answered locally instead of forwarded to Trino,
 /// build the response and return it. Returns `None` for queries that
 /// should pass through to Trino.
@@ -37,11 +54,15 @@ fn single_text_response(column_name: &str, value: &str) -> PgWireResult<Vec<Resp
 /// `calls_function`) handle the cases where prefix matching would
 /// misroute a user query that *contains* a catalog name as a literal or
 /// column reference.
+///
+/// `client_is_secure` is the TLS state of the PG client connection, used to
+/// answer the `pg_stat_ssl` probe truthfully (see below).
 pub fn try_intercept(
     query: &str,
     parsed_query: &ParsedQuery,
     catalog: &str,
     schema: &str,
+    client_is_secure: bool,
 ) -> Option<PgWireResult<Vec<Response>>> {
     let trimmed = query.trim();
     if trimmed.is_empty() {
@@ -95,6 +116,17 @@ pub fn try_intercept(
         return Some(resp);
     }
 
+    // Npgsql (Power BI) and pgjdbc probe the connection's encryption state
+    // right after connecting with `SELECT ssl FROM pg_stat_ssl WHERE pid =
+    // pg_backend_pid()`. Trino has no `pg_stat_ssl`, so forwarding it aborts
+    // the session before the client runs a single user query. Answer locally
+    // with the real TLS state of this connection; the `WHERE pid = ...` clause
+    // is moot because a backend only ever sees its own row in pg_stat_ssl.
+    if parsed_query.references_table("pg_stat_ssl") {
+        tracing::trace!(query = trimmed, "Intercept: pg_stat_ssl");
+        return Some(single_bool_response("ssl", client_is_secure));
+    }
+
     if let Some(resp) = crate::catalog::handle_catalog_query(parsed_query) {
         tracing::trace!(query = trimmed, "Intercept: pg_catalog");
         return Some(resp);
@@ -188,15 +220,15 @@ mod tests {
     fn assert_intercepted(query: &str) {
         let parsed_query = ParsedQuery::new(query);
         assert!(
-            try_intercept(query, &parsed_query, "test_catalog", "test_schema").is_some(),
+            try_intercept(query, &parsed_query, "test_catalog", "test_schema", false).is_some(),
             "expected query to be intercepted: {query}"
         );
     }
 
     fn assert_not_intercepted(query: &str) {
         let parsed_query = ParsedQuery::new(query);
         assert!(
-            try_intercept(query, &parsed_query, "test_catalog", "test_schema").is_none(),
+            try_intercept(query, &parsed_query, "test_catalog", "test_schema", false).is_none(),
             "expected query to NOT be intercepted: {query}"
         );
     }
@@ -246,7 +278,7 @@ mod tests {
 
         for &(query, expected) in cases {
             let parsed_query = ParsedQuery::new(query);
-            let result = try_intercept(query, &parsed_query, "test_catalog", "test_schema")
+            let result = try_intercept(query, &parsed_query, "test_catalog", "test_schema", false)
                 .unwrap_or_else(|| panic!("SHOW not intercepted: {query}"));
             assert!(result.is_ok(), "SHOW returned error for: {query}");
 
@@ -280,6 +312,45 @@ mod tests {
         assert_intercepted("SELECT current_setting('server_version')");
     }
 
+    /// The Npgsql/pgjdbc SSL probe is answered locally instead of forwarded
+    /// to Trino (which has no `pg_stat_ssl` and would abort the connection).
+    #[test]
+    fn pg_stat_ssl_probe_intercepted() {
+        assert_intercepted("SELECT ssl FROM pg_stat_ssl WHERE pid = pg_backend_pid()");
+        assert_intercepted("select ssl from pg_stat_ssl");
+    }
+
+    /// The `ssl` column reports the connection's real TLS state and is typed
+    /// BOOL (not VARCHAR) so the RowDescription oid matches what the driver
+    /// decodes.
+    #[test]
+    fn pg_stat_ssl_reports_tls_state_as_bool() {
+        let query = "SELECT ssl FROM pg_stat_ssl WHERE pid = pg_backend_pid()";
+        let parsed_query = ParsedQuery::new(query);
+
+        for is_secure in [true, false] {
+            let resp = try_intercept(query, &parsed_query, "cat", "sch", is_secure)
+                .expect("pg_stat_ssl must be intercepted")
+                .expect("intercept must not error");
+            assert_eq!(resp.len(), 1);
+            match &resp[0] {
+                Response::Query(qr) => {
+                    assert_eq!(qr.row_schema.len(), 1);
+                    assert_eq!(qr.row_schema[0].name(), "ssl");
+                    assert_eq!(qr.row_schema[0].datatype(), &Type::BOOL);
+                }
+                other => panic!("expected Query response, got: {other:?}"),
+            }
+        }
+    }
+
+    /// Regression: a user table literally named `pg_stat_ssl` in a literal
+    /// must not trip the probe intercept.
+    #[test]
+    fn pg_stat_ssl_in_literal_not_intercepted() {
+        assert_not_intercepted("SELECT 'pg_stat_ssl' AS sentinel");
+    }
+
     #[test]
     fn regular_queries_not_intercepted() {
         assert_not_intercepted("SELECT 1");

src/query_extended.rs

@@ -91,6 +91,8 @@ impl ExtendedQueryHandler for GatewayExtendedQueryHandler {
         let query = &portal.statement.statement;
         tracing::debug!(query, "Extended query execute");
 
+        let client_is_secure = client.is_secure();
+
         let conn_state = client
             .session_extensions()
             .get::<ConnectionState>()
@@ -118,6 +120,7 @@ impl ExtendedQueryHandler for GatewayExtendedQueryHandler {
             &conn_state.config,
             Some(&conn_state.active_query_id),
             Some(&portal.result_column_format),
+            client_is_secure,
         )
         .await?;
         let response = responses
@@ -182,6 +185,8 @@ impl ExtendedQueryHandler for GatewayExtendedQueryHandler {
             return Ok(DescribeStatementResponse::no_data());
         }
 
+        let client_is_secure = client.is_secure();
+
         let conn_state = client
             .session_extensions()
             .get::<ConnectionState>()
@@ -201,6 +206,7 @@ impl ExtendedQueryHandler for GatewayExtendedQueryHandler {
             &conn_state.config,
             Some(&conn_state.active_query_id),
             None,
+            client_is_secure,
         )
         .await?;
         let response = responses
@@ -230,6 +236,8 @@ impl ExtendedQueryHandler for GatewayExtendedQueryHandler {
         let query = &portal.statement.statement;
         tracing::debug!(query, "Extended query describe portal");
 
+        let client_is_secure = client.is_secure();
+
         let conn_state = client
             .session_extensions()
             .get::<ConnectionState>()
@@ -241,6 +249,7 @@ impl ExtendedQueryHandler for GatewayExtendedQueryHandler {
             &conn_state.config,
             Some(&conn_state.active_query_id),
             Some(&portal.result_column_format),
+            client_is_secure,
         )
         .await?;
         let response = responses

src/query_pipeline.rs

@@ -49,6 +49,7 @@ pub(crate) async fn process_query(
     config: &Arc<Config>,
     active_query_id: Option<&ActiveQueryId>,
     result_format: Option<&Format>,
+    client_is_secure: bool,
 ) -> PgWireResult<Vec<Response>> {
     tracing::trace!(query, "Pipeline: enter");
 
@@ -60,15 +61,23 @@ pub(crate) async fn process_query(
             config,
             active_query_id,
             result_format,
+            client_is_secure,
         )
         .await;
     }
 
     tracing::trace!(count = pieces.len(), "Pipeline: multi-statement input");
     let mut out = Vec::with_capacity(pieces.len());
     for stmt in &pieces {
-        match process_single_statement(stmt, trino_client, config, active_query_id, result_format)
-            .await
+        match process_single_statement(
+            stmt,
+            trino_client,
+            config,
+            active_query_id,
+            result_format,
+            client_is_secure,
+        )
+        .await
         {
             Ok(mut responses) => out.append(&mut responses),
             // User-visible errors (e.g. a Trino syntax error on statement N
@@ -107,6 +116,7 @@ async fn process_single_statement(
     config: &Arc<Config>,
     active_query_id: Option<&ActiveQueryId>,
     result_format: Option<&Format>,
+    client_is_secure: bool,
 ) -> PgWireResult<Vec<Response>> {
     // The query is parsed up to three times: once here (for routing
     // checks), once by the multi-statement splitter in the public
@@ -122,6 +132,7 @@ async fn process_single_statement(
         &parsed_query,
         &config.trino_catalog,
         &config.trino_schema,
+        client_is_secure,
     ) {
         tracing::trace!("Pipeline: static intercept matched");
         return result;

src/query_simple.rs

@@ -26,6 +26,10 @@ impl SimpleQueryHandler for GatewayQueryHandler {
     {
         tracing::debug!(query, "Simple query received");
 
+        // Captured before borrowing session state, so the `pg_stat_ssl`
+        // intercept can report this connection's real TLS state.
+        let client_is_secure = client.is_secure();
+
         let conn_state = client
             .session_extensions()
             .get::<ConnectionState>()
@@ -39,6 +43,7 @@ impl SimpleQueryHandler for GatewayQueryHandler {
             &conn_state.config,
             Some(&conn_state.active_query_id),
             None,
+            client_is_secure,
         )
         .await;
         match &result {