Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
77 changes: 74 additions & 3 deletions src/intercept.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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() {
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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}"
);
}
Expand Down Expand Up @@ -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}");

Expand Down Expand Up @@ -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");
Expand Down
9 changes: 9 additions & 0 deletions src/query_extended.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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>()
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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>()
Expand All @@ -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
Expand Down Expand Up @@ -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>()
Expand All @@ -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
Expand Down
15 changes: 13 additions & 2 deletions src/query_pipeline.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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");

Expand All @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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;
Expand Down
5 changes: 5 additions & 0 deletions src/query_simple.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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>()
Expand All @@ -39,6 +43,7 @@ impl SimpleQueryHandler for GatewayQueryHandler {
&conn_state.config,
Some(&conn_state.active_query_id),
None,
client_is_secure,
)
.await;
match &result {
Expand Down
Loading