pgwire/
pg_server.rs

1// Copyright 2022 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::collections::HashMap;
16use std::future::Future;
17use std::str::FromStr;
18use std::sync::Arc;
19use std::time::Instant;
20
21use bytes::Bytes;
22use jsonwebtoken::{Algorithm, DecodingKey, Validation, decode, decode_header};
23use parking_lot::Mutex;
24use risingwave_common::config::HbaEntry;
25use risingwave_common::id::DatabaseId;
26use risingwave_common::types::DataType;
27use risingwave_common::util::runtime::BackgroundShutdownRuntime;
28use risingwave_common::util::tokio_util::sync::CancellationToken;
29use risingwave_sqlparser::ast::Statement;
30use serde::Deserialize;
31use thiserror_ext::AsReport;
32
33use crate::error::{PsqlError, PsqlResult};
34use crate::ldap_auth::LdapAuthenticator;
35use crate::net::{AddressRef, Listener, TcpKeepalive};
36use crate::pg_field_descriptor::PgFieldDescriptor;
37use crate::pg_message::TransactionStatus;
38use crate::pg_protocol::{ConnectionContext, PgByteStream, PgProtocol};
39use crate::pg_response::{PgResponse, ValuesStream};
40use crate::types::Format;
41
42pub type BoxedError = Box<dyn std::error::Error + Send + Sync>;
43type ProcessId = i32;
44type SecretKey = i32;
45pub type SessionId = (ProcessId, SecretKey);
46
47/// The interface for a database system behind pgwire protocol.
48/// We can mock it for testing purpose.
49pub trait SessionManager: Send + Sync + 'static {
50    type Error: Into<BoxedError>;
51    type Session: Session<Error = Self::Error>;
52
53    /// In the process of auto schema change, we need a dummy session to access
54    /// catalog information in frontend and build a replace plan for the table.
55    fn create_dummy_session(
56        &self,
57        database_id: DatabaseId,
58    ) -> Result<Arc<Self::Session>, Self::Error>;
59
60    fn connect(
61        &self,
62        database: &str,
63        user_name: &str,
64        peer_addr: AddressRef,
65    ) -> Result<Arc<Self::Session>, Self::Error>;
66
67    fn cancel_queries_in_session(&self, session_id: SessionId);
68
69    fn cancel_creating_jobs_in_session(&self, session_id: SessionId);
70
71    fn end_session(&self, session: &Self::Session);
72
73    /// Run some cleanup tasks before the server shutdown.
74    fn shutdown(&self) -> impl Future<Output = ()> + Send {
75        async {}
76    }
77}
78
79/// A psql connection. Each connection binds with a database. Switching database will need to
80/// recreate another connection.
81pub trait Session: Send + Sync {
82    type Error: Into<BoxedError>;
83    type ValuesStream: ValuesStream;
84    type PreparedStatement: Send + Clone + 'static;
85    type Portal: Send + Clone + std::fmt::Display + 'static;
86
87    /// The str sql can not use the unparse from AST: There is some problem when dealing with create
88    /// view, see <https://github.com/risingwavelabs/risingwave/issues/6801>.
89    fn run_one_query(
90        self: Arc<Self>,
91        stmt: Statement,
92        format: Format,
93    ) -> impl Future<Output = Result<PgResponse<Self::ValuesStream>, Self::Error>> + Send;
94
95    fn parse(
96        self: Arc<Self>,
97        sql: Option<Statement>,
98        params_types: Vec<Option<DataType>>,
99    ) -> impl Future<Output = Result<Self::PreparedStatement, Self::Error>> + Send;
100
101    /// Receive the next notice message to send to the client.
102    ///
103    /// This function should be cancellation-safe.
104    fn next_notice(self: &Arc<Self>) -> impl Future<Output = String> + Send;
105
106    fn bind(
107        self: Arc<Self>,
108        prepare_statement: Self::PreparedStatement,
109        params: Vec<Option<Bytes>>,
110        param_formats: Vec<Format>,
111        result_formats: Vec<Format>,
112    ) -> Result<Self::Portal, Self::Error>;
113
114    fn execute(
115        self: Arc<Self>,
116        portal: Self::Portal,
117    ) -> impl Future<Output = Result<PgResponse<Self::ValuesStream>, Self::Error>> + Send;
118
119    fn describe_statement(
120        self: Arc<Self>,
121        prepare_statement: Self::PreparedStatement,
122    ) -> Result<(Vec<DataType>, Vec<PgFieldDescriptor>), Self::Error>;
123
124    fn describe_portal(
125        self: Arc<Self>,
126        portal: Self::Portal,
127    ) -> Result<Vec<PgFieldDescriptor>, Self::Error>;
128
129    fn user_authenticator(&self) -> &UserAuthenticator;
130
131    fn id(&self) -> SessionId;
132
133    fn get_config(&self, key: &str) -> Result<String, Self::Error>;
134
135    fn set_config(&self, key: &str, value: String) -> Result<String, Self::Error>;
136
137    fn transaction_status(&self) -> TransactionStatus;
138
139    fn init_exec_context(&self, sql: Arc<str>) -> ExecContextGuard;
140
141    fn check_idle_in_transaction_timeout(&self) -> PsqlResult<()>;
142
143    fn user(&self) -> String;
144}
145
146/// Each session could run different SQLs multiple times.
147/// `ExecContext` represents the lifetime of a running SQL in the current session.
148pub struct ExecContext {
149    pub running_sql: Arc<str>,
150    /// The instant of the running sql
151    pub last_instant: Instant,
152    /// A reference used to update when `ExecContext` is dropped
153    pub last_idle_instant: Arc<Mutex<Option<Instant>>>,
154}
155
156/// `ExecContextGuard` holds a `Arc` pointer. Once `ExecContextGuard` is dropped,
157/// the inner `Arc<ExecContext>` should not be referred anymore, so that its `Weak` reference (used in `SessionImpl`) will be the same lifecycle of the running sql execution context.
158pub struct ExecContextGuard(#[expect(dead_code)] Arc<ExecContext>);
159
160impl ExecContextGuard {
161    pub fn new(exec_context: Arc<ExecContext>) -> Self {
162        Self(exec_context)
163    }
164}
165
166impl Drop for ExecContext {
167    fn drop(&mut self) {
168        *self.last_idle_instant.lock() = Some(Instant::now());
169    }
170}
171
172#[derive(Debug, Clone)]
173pub enum UserAuthenticator {
174    // No need to authenticate.
175    None,
176    // raw password in clear-text form.
177    ClearText(Vec<u8>),
178    // password encrypted with random salt.
179    Md5WithSalt {
180        encrypted_password: Vec<u8>,
181        salt: [u8; 4],
182    },
183    OAuth {
184        metadata: HashMap<String, String>,
185        cluster_id: String,
186    },
187    Ldap(String, HbaEntry),
188}
189
190/// A JWK Set is a JSON object that represents a set of JWKs.
191/// The JSON object MUST have a "keys" member, with its value being an array of JWKs.
192/// See <https://www.rfc-editor.org/rfc/rfc7517.html#section-5> for more details.
193#[derive(Debug, Deserialize)]
194struct Jwks {
195    keys: Vec<Jwk>,
196}
197
198/// A JSON Web Key (JWK) is a JSON object that represents a cryptographic key.
199/// See <https://www.rfc-editor.org/rfc/rfc7517.html#section-4> for more details.
200#[derive(Debug, Deserialize)]
201struct Jwk {
202    kty: Option<String>, // Key Type
203    kid: Option<String>, // Key ID (OPTIONAL per RFC 7517 section 4.5)
204    alg: Option<String>, // Algorithm (OPTIONAL per RFC 7517 section 4.4)
205    n: Option<String>,   // RSA modulus
206    e: Option<String>,   // RSA exponent
207}
208
209/// Algorithms we accept for JWT signature verification.
210///
211/// Restricted to RSA-family algorithms because the only `DecodingKey` we build
212/// is from RSA components (`n`, `e`). Pinning the algorithm to a server-side
213/// allow-list also prevents the classic alg-confusion attack: a token with
214/// `alg: "none"` (no signature) or `alg: "HS256"` forged using the RSA public
215/// key as the HMAC secret cannot select a verification algorithm outside this
216/// set.
217const ALLOWED_JWT_ALGORITHMS: &[Algorithm] = &[
218    Algorithm::RS256,
219    Algorithm::RS384,
220    Algorithm::RS512,
221    Algorithm::PS256,
222    Algorithm::PS384,
223    Algorithm::PS512,
224];
225const RSA_JWK_KEY_TYPE: &str = "RSA";
226/// Optional OAuth user option overriding the expected JWT `aud`.
227const OAUTH_AUDIENCE_KEY: &str = "audience";
228
229async fn validate_jwt(
230    jwt: &str,
231    jwks_url: &str,
232    issuer: &str,
233    audience: &str,
234    metadata: &HashMap<String, String>,
235) -> Result<bool, BoxedError> {
236    let jwks: Jwks = reqwest::get(jwks_url).await?.json().await?;
237    validate_jwt_with_jwks(jwt, &jwks, issuer, audience, metadata)
238}
239
240fn audience_from_cluster_id(cluster_id: &str) -> String {
241    format!("urn:risingwave:cluster:{}", cluster_id)
242}
243
244fn resolve_oauth_audience(
245    configured: Option<String>,
246    cluster_id: &str,
247) -> Result<String, BoxedError> {
248    let Some(configured_audience) = configured else {
249        return Ok(audience_from_cluster_id(cluster_id));
250    };
251
252    let audience = configured_audience.trim();
253    if audience.is_empty() {
254        return Err(format!("OAuth option `{OAUTH_AUDIENCE_KEY}` cannot be empty").into());
255    }
256    Ok(audience.to_owned())
257}
258
259fn validate_jwt_with_jwks(
260    jwt: &str,
261    jwks: &Jwks,
262    issuer: &str,
263    audience: &str,
264    metadata: &HashMap<String, String>,
265) -> Result<bool, BoxedError> {
266    let header = decode_header(jwt)?;
267
268    // 1. Retrieve the kid from the header to find compatible JWKs in the JWK Set.
269    let kid = header.kid.ok_or("JWT header missing 'kid' field")?;
270
271    // 2. Decide which algorithm to use.
272    //
273    // Per RFC 7517 §4.4 the JWK `alg` member is OPTIONAL. When the JWK pins an
274    // `alg`, the JWT header MUST match it; when it doesn't, we fall back to
275    // the header's `alg` but only after checking it against a server-side
276    // allow-list. The allow-list is what ultimately blocks alg-confusion: an
277    // attacker-chosen `alg` from the token header alone must never be trusted
278    // to select the verification algorithm.
279    if !ALLOWED_JWT_ALGORITHMS.contains(&header.alg) {
280        return Err(format!("JWT alg {:?} is not allowed", header.alg).into());
281    }
282
283    // A JWK Set can contain unrelated key types, keys without an optional `kid`,
284    // and equivalent keys of different types that share a `kid`. Scan all
285    // candidates instead of letting the first same-`kid` key shadow a usable
286    // RSA key later in the set.
287    let jwk = jwks
288        .keys
289        .iter()
290        .filter(|jwk| jwk.kid.as_deref() == Some(kid.as_str()))
291        .find(|jwk| {
292            jwk.kty.as_deref() == Some(RSA_JWK_KEY_TYPE)
293                && jwk.n.is_some()
294                && jwk.e.is_some()
295                && match jwk.alg.as_deref() {
296                    Some(alg) => Algorithm::from_str(alg).is_ok_and(|alg| alg == header.alg),
297                    None => true,
298                }
299        })
300        .ok_or_else(|| format!("No compatible RSA key found in JWKS for kid: '{}'", kid))?;
301
302    // 3. Decode the JWT and validate the claims.
303    let n = jwk.n.as_deref().expect("RSA candidate must have a modulus");
304    let e = jwk
305        .e
306        .as_deref()
307        .expect("RSA candidate must have an exponent");
308    let decoding_key = DecodingKey::from_rsa_components(n, e)?;
309    let mut validation = Validation::new(header.alg);
310    validation.set_issuer(&[issuer]);
311    validation.set_audience(&[audience]);
312    validation.set_required_spec_claims(&["exp", "iss", "aud"]);
313    let token_data = decode::<HashMap<String, serde_json::Value>>(jwt, &decoding_key, &validation)?;
314
315    // 4. Check if the metadata in the token matches.
316    if !metadata.iter().all(
317        |(k, v)| matches!(token_data.claims.get(k), Some(serde_json::Value::String(s)) if s == v),
318    ) {
319        return Err("metadata in jwt does not match with metadata declared with user".into());
320    }
321    Ok(true)
322}
323
324impl UserAuthenticator {
325    pub async fn authenticate(&self, password: &[u8]) -> PsqlResult<()> {
326        let success = match self {
327            UserAuthenticator::None => true,
328            UserAuthenticator::ClearText(text) => password == text,
329            UserAuthenticator::Md5WithSalt {
330                encrypted_password, ..
331            } => encrypted_password == password,
332            UserAuthenticator::OAuth {
333                metadata,
334                cluster_id,
335            } => {
336                let mut metadata = metadata.clone();
337                let jwks_url = metadata.remove("jwks_url").unwrap();
338                let issuer = metadata.remove("issuer").unwrap();
339                let audience =
340                    resolve_oauth_audience(metadata.remove(OAUTH_AUDIENCE_KEY), cluster_id)
341                        .map_err(PsqlError::StartupError)?;
342                validate_jwt(
343                    &String::from_utf8_lossy(password),
344                    &jwks_url,
345                    &issuer,
346                    &audience,
347                    &metadata,
348                )
349                .await
350                .map_err(PsqlError::StartupError)?
351            }
352            UserAuthenticator::Ldap(user_name, hba_entry) => {
353                let ldap_auth = LdapAuthenticator::new(hba_entry)?;
354                // Convert password to string, defaulting to empty if not valid UTF-8
355                let password_str = String::from_utf8_lossy(password).into_owned();
356                ldap_auth.authenticate(user_name, &password_str).await?
357            }
358        };
359        if !success {
360            return Err(PsqlError::PasswordError);
361        }
362        Ok(())
363    }
364}
365
366/// Binds a Tcp or Unix listener at `addr`. Spawn a coroutine to serve every new connection.
367///
368/// Returns when the `shutdown` token is triggered.
369pub async fn pg_serve(
370    addr: &str,
371    tcp_keepalive: TcpKeepalive,
372    session_mgr: Arc<impl SessionManager>,
373    context: ConnectionContext,
374    shutdown: CancellationToken,
375) -> Result<(), BoxedError> {
376    let listener = Listener::bind(addr).await?;
377    tracing::info!(addr, "server started");
378
379    let acceptor_runtime = BackgroundShutdownRuntime::from({
380        let mut builder = tokio::runtime::Builder::new_multi_thread();
381        builder.worker_threads(1);
382        builder
383            .thread_name("rw-acceptor")
384            .enable_all()
385            .build()
386            .unwrap()
387    });
388
389    #[cfg(not(madsim))]
390    let worker_runtime = tokio::runtime::Handle::current();
391    #[cfg(madsim)]
392    let worker_runtime = tokio::runtime::Builder::new_multi_thread().build().unwrap();
393    let session_mgr_clone = session_mgr.clone();
394    let f = async move {
395        loop {
396            let conn_ret = listener.accept(&tcp_keepalive).await;
397            match conn_ret {
398                Ok((stream, peer_addr)) => {
399                    tracing::info!(%peer_addr, "accept connection");
400                    worker_runtime.spawn(handle_connection(
401                        stream,
402                        session_mgr_clone.clone(),
403                        Arc::new(peer_addr),
404                        context.clone(),
405                    ));
406                }
407
408                Err(e) => {
409                    tracing::error!(error = %e.as_report(), "failed to accept connection",);
410                }
411            }
412        }
413    };
414    acceptor_runtime.spawn(f);
415
416    // Wait for the shutdown signal.
417    shutdown.cancelled().await;
418
419    // Stop accepting new connections.
420    drop(acceptor_runtime);
421    // Shutdown session manager, typically close all existing sessions.
422    session_mgr.shutdown().await;
423
424    Ok(())
425}
426
427pub async fn handle_connection<S, SM>(
428    stream: S,
429    session_mgr: Arc<SM>,
430    peer_addr: AddressRef,
431    context: ConnectionContext,
432) where
433    S: PgByteStream,
434    SM: SessionManager,
435{
436    PgProtocol::new(stream, session_mgr, peer_addr, context)
437        .run()
438        .await;
439}
440#[cfg(test)]
441mod tests {
442    use std::sync::Arc;
443    use std::time::{Duration, Instant};
444
445    use bytes::Bytes;
446    use futures::StreamExt;
447    use futures::stream::BoxStream;
448    use risingwave_common::id::DatabaseId;
449    use risingwave_common::types::DataType;
450    use risingwave_common::util::tokio_util::sync::CancellationToken;
451    use risingwave_sqlparser::ast::Statement;
452    use tokio_postgres::NoTls;
453
454    use crate::error::PsqlResult;
455    use crate::memory_manager::MessageMemoryManager;
456    use crate::pg_field_descriptor::PgFieldDescriptor;
457    use crate::pg_message::TransactionStatus;
458    use crate::pg_protocol::ConnectionContext;
459    use crate::pg_response::{PgResponse, RowSetResult, StatementType};
460    use crate::pg_server::{
461        BoxedError, ExecContext, ExecContextGuard, Session, SessionId, SessionManager,
462        UserAuthenticator, pg_serve,
463    };
464    use crate::types;
465    use crate::types::Row;
466
467    struct MockSessionManager {}
468    struct MockSession {}
469
470    const STREAMING_TEST_QUERY: &str = "SELECT 'pgwire_streaming_test'";
471
472    impl SessionManager for MockSessionManager {
473        type Error = BoxedError;
474        type Session = MockSession;
475
476        fn create_dummy_session(
477            &self,
478            _database_id: DatabaseId,
479        ) -> Result<Arc<Self::Session>, Self::Error> {
480            unimplemented!()
481        }
482
483        fn connect(
484            &self,
485            _database: &str,
486            _user_name: &str,
487            _peer_addr: crate::net::AddressRef,
488        ) -> Result<Arc<Self::Session>, Self::Error> {
489            Ok(Arc::new(MockSession {}))
490        }
491
492        fn cancel_queries_in_session(&self, _session_id: SessionId) {
493            todo!()
494        }
495
496        fn cancel_creating_jobs_in_session(&self, _session_id: SessionId) {
497            todo!()
498        }
499
500        fn end_session(&self, _session: &Self::Session) {}
501    }
502
503    impl Session for MockSession {
504        type Error = BoxedError;
505        type Portal = String;
506        type PreparedStatement = String;
507        type ValuesStream = BoxStream<'static, RowSetResult>;
508
509        async fn run_one_query(
510            self: Arc<Self>,
511            _stmt: Statement,
512            _format: types::Format,
513        ) -> Result<PgResponse<BoxStream<'static, RowSetResult>>, Self::Error> {
514            Ok(PgResponse::builder(StatementType::SELECT)
515                .values(
516                    futures::stream::iter(vec![Ok(vec![Row::new(vec![Some(Bytes::new())])])])
517                        .boxed(),
518                    vec![
519                        // 1043 is the oid of varchar type.
520                        // -1 is the type len of varchar type.
521                        PgFieldDescriptor::new("".to_owned(), 1043, -1);
522                        1
523                    ],
524                )
525                .into())
526        }
527
528        async fn parse(
529            self: Arc<Self>,
530            sql: Option<Statement>,
531            _params_types: Vec<Option<DataType>>,
532        ) -> Result<String, Self::Error> {
533            Ok(sql.map(|stmt| stmt.to_string()).unwrap_or_default())
534        }
535
536        fn bind(
537            self: Arc<Self>,
538            prepare_statement: String,
539            _params: Vec<Option<Bytes>>,
540            _param_formats: Vec<types::Format>,
541            _result_formats: Vec<types::Format>,
542        ) -> Result<String, Self::Error> {
543            Ok(prepare_statement)
544        }
545
546        async fn execute(
547            self: Arc<Self>,
548            portal: String,
549        ) -> Result<PgResponse<BoxStream<'static, RowSetResult>>, Self::Error> {
550            if portal == STREAMING_TEST_QUERY {
551                let first_row = futures::stream::once(async {
552                    Ok(vec![Row::new(vec![Some(Bytes::from(vec![
553                        b'x';
554                        128 * 1024
555                    ]))])])
556                });
557                let remaining_rows = futures::stream::pending();
558                return Ok(PgResponse::builder(StatementType::SELECT)
559                    .values(
560                        first_row.chain(remaining_rows).boxed(),
561                        vec![PgFieldDescriptor::new("".to_owned(), 1043, -1)],
562                    )
563                    .into());
564            }
565
566            Ok(PgResponse::builder(StatementType::SELECT)
567                .values(
568                    futures::stream::iter(vec![Ok(vec![Row::new(vec![Some(Bytes::new())])])])
569                        .boxed(),
570                    vec![
571                    // 1043 is the oid of varchar type.
572                    // -1 is the type len of varchar type.
573                    PgFieldDescriptor::new("".to_owned(), 1043, -1);
574                    1
575                ],
576                )
577                .into())
578        }
579
580        fn describe_statement(
581            self: Arc<Self>,
582            _statement: String,
583        ) -> Result<(Vec<DataType>, Vec<PgFieldDescriptor>), Self::Error> {
584            Ok((
585                vec![],
586                vec![PgFieldDescriptor::new("".to_owned(), 1043, -1)],
587            ))
588        }
589
590        fn describe_portal(
591            self: Arc<Self>,
592            _portal: String,
593        ) -> Result<Vec<PgFieldDescriptor>, Self::Error> {
594            Ok(vec![PgFieldDescriptor::new("".to_owned(), 1043, -1)])
595        }
596
597        fn user_authenticator(&self) -> &UserAuthenticator {
598            &UserAuthenticator::None
599        }
600
601        fn id(&self) -> SessionId {
602            (0, 0)
603        }
604
605        fn get_config(&self, key: &str) -> Result<String, Self::Error> {
606            match key {
607                "timezone" => Ok("UTC".to_owned()),
608                _ => Err(format!("Unknown config key: {key}").into()),
609            }
610        }
611
612        fn set_config(&self, _key: &str, _value: String) -> Result<String, Self::Error> {
613            Ok("".to_owned())
614        }
615
616        async fn next_notice(self: &Arc<Self>) -> String {
617            std::future::pending().await
618        }
619
620        fn transaction_status(&self) -> TransactionStatus {
621            TransactionStatus::Idle
622        }
623
624        fn init_exec_context(&self, sql: Arc<str>) -> ExecContextGuard {
625            let exec_context = Arc::new(ExecContext {
626                running_sql: sql,
627                last_instant: Instant::now(),
628                last_idle_instant: Default::default(),
629            });
630            ExecContextGuard::new(exec_context)
631        }
632
633        fn check_idle_in_transaction_timeout(&self) -> PsqlResult<()> {
634            Ok(())
635        }
636
637        fn user(&self) -> String {
638            "mock".to_owned()
639        }
640    }
641
642    async fn do_test_query(bind_addr: impl Into<String>, pg_config: impl Into<String>) {
643        let bind_addr = bind_addr.into();
644        let pg_config = pg_config.into();
645
646        let session_mgr = MockSessionManager {};
647        tokio::spawn(async move {
648            pg_serve(
649                &bind_addr,
650                socket2::TcpKeepalive::new(),
651                Arc::new(session_mgr),
652                ConnectionContext {
653                    tls_config: None,
654                    sql_option_redaction: None,
655                    message_memory_manager: MessageMemoryManager::new(u64::MAX, u64::MAX, u64::MAX)
656                        .into(),
657                    stream_flush_threshold_bytes: 64 * 1024,
658                },
659                CancellationToken::new(), // dummy
660            )
661            .await
662        });
663        // wait for server to start
664        tokio::time::sleep(std::time::Duration::from_millis(100)).await;
665
666        // Connect to the database.
667        let (client, connection) = tokio_postgres::connect(&pg_config, NoTls).await.unwrap();
668
669        // The connection object performs the actual communication with the database,
670        // so spawn it off to run on its own.
671        tokio::spawn(async move {
672            if let Err(e) = connection.await {
673                eprintln!("connection error: {}", e);
674            }
675        });
676
677        let rows = client
678            .simple_query("SELECT ''")
679            .await
680            .expect("Error executing query");
681        // Row + CommandComplete
682        assert_eq!(rows.len(), 2);
683
684        let rows = client
685            .query("SELECT ''", &[])
686            .await
687            .expect("Error executing query");
688        assert_eq!(rows.len(), 1);
689    }
690
691    #[tokio::test]
692    async fn test_query_tcp() {
693        do_test_query("127.0.0.1:10000", "host=localhost port=10000").await;
694    }
695
696    #[cfg(not(madsim))]
697    #[tokio::test]
698    async fn test_query_unix() {
699        let port: i16 = 10000;
700        let dir = tempfile::TempDir::new().unwrap();
701        let sock = dir.path().join(format!(".s.PGSQL.{port}"));
702
703        do_test_query(
704            format!("unix:{}", sock.to_str().unwrap()),
705            format!("host={} port={}", dir.path().to_str().unwrap(), port),
706        )
707        .await;
708    }
709
710    #[cfg(not(madsim))]
711    #[tokio::test]
712    async fn test_large_query_row_is_streamed_before_query_finishes() {
713        let port: i16 = 10001;
714        let dir = tempfile::TempDir::new().unwrap();
715        let sock = dir.path().join(format!(".s.PGSQL.{port}"));
716        let bind_addr = format!("unix:{}", sock.to_str().unwrap());
717        let pg_config = format!("host={} port={port}", dir.path().to_str().unwrap());
718        let cancellation = CancellationToken::new();
719
720        let server_cancellation = cancellation.clone();
721        tokio::spawn(async move {
722            pg_serve(
723                &bind_addr,
724                socket2::TcpKeepalive::new(),
725                Arc::new(MockSessionManager {}),
726                ConnectionContext {
727                    tls_config: None,
728                    sql_option_redaction: None,
729                    message_memory_manager: MessageMemoryManager::new(u64::MAX, u64::MAX, u64::MAX)
730                        .into(),
731                    stream_flush_threshold_bytes: 64 * 1024,
732                },
733                server_cancellation,
734            )
735            .await
736        });
737        tokio::time::sleep(Duration::from_millis(100)).await;
738
739        let (client, connection) = tokio_postgres::connect(&pg_config, NoTls).await.unwrap();
740        let connection_handle = tokio::spawn(connection);
741        let params: &[&str] = &[];
742        let rows = client
743            .query_raw(STREAMING_TEST_QUERY, params)
744            .await
745            .unwrap();
746        futures::pin_mut!(rows);
747
748        let first_row = tokio::time::timeout(Duration::from_secs(5), rows.next())
749            .await
750            .expect("first row should be sent before the result stream finishes")
751            .unwrap()
752            .unwrap();
753        assert_eq!(first_row.get::<_, &str>(0).len(), 128 * 1024);
754
755        drop(client);
756        cancellation.cancel();
757        connection_handle.abort();
758    }
759
760    mod jwt_validation_tests {
761        use std::collections::HashMap;
762        use std::time::{SystemTime, UNIX_EPOCH};
763
764        use base64::Engine;
765        use jsonwebtoken::{Algorithm, EncodingKey, Header};
766        use openssl::pkey::{Private, Public};
767        use openssl::rsa::Rsa;
768        use serde_json::json;
769
770        use crate::pg_server::{Jwk, Jwks, resolve_oauth_audience, validate_jwt_with_jwks};
771
772        fn create_test_rsa_keys() -> (Rsa<Private>, Rsa<Public>) {
773            let private_key = Rsa::generate(2048).expect("failed to generate a key");
774            let public_key = Rsa::from_public_components(
775                private_key.n().to_owned().unwrap(),
776                private_key.e().to_owned().unwrap(),
777            )
778            .unwrap();
779            (private_key, public_key)
780        }
781
782        fn create_test_jwks(public_key: &Rsa<Public>, kid: &str, alg: Option<&str>) -> Jwks {
783            let n =
784                base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(public_key.n().to_vec());
785            let e =
786                base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(public_key.e().to_vec());
787
788            Jwks {
789                keys: vec![Jwk {
790                    kty: Some("RSA".to_owned()),
791                    kid: Some(kid.to_owned()),
792                    alg: alg.map(ToOwned::to_owned),
793                    n: Some(n),
794                    e: Some(e),
795                }],
796            }
797        }
798
799        fn create_jwt_token(
800            private_key: &Rsa<Private>,
801            kid: &str,
802            algorithm: Algorithm,
803            issuer: &str,
804            audience: Option<&str>,
805            exp: u64,
806            additional_claims: HashMap<String, serde_json::Value>,
807        ) -> String {
808            let mut header = Header::new(algorithm);
809            header.kid = Some(kid.to_owned());
810
811            let mut claims = json!({
812                "iss": issuer,
813                "exp": exp,
814            });
815
816            if let Some(aud) = audience {
817                claims["aud"] = json!(aud);
818            }
819
820            for (key, value) in additional_claims {
821                claims[key] = value;
822            }
823
824            let encoding_key =
825                EncodingKey::from_rsa_pem(&private_key.private_key_to_pem().unwrap()).unwrap();
826
827            jsonwebtoken::encode(&header, &claims, &encoding_key).unwrap()
828        }
829
830        fn get_future_timestamp() -> u64 {
831            SystemTime::now()
832                .duration_since(UNIX_EPOCH)
833                .unwrap()
834                .as_secs()
835                + 3600 // 1 hour from now
836        }
837
838        fn get_past_timestamp() -> u64 {
839            SystemTime::now()
840                .duration_since(UNIX_EPOCH)
841                .unwrap()
842                .as_secs()
843                - 3600 // 1 hour ago
844        }
845
846        #[test]
847        fn test_jwt_with_invalid_audience() {
848            let (private_key, public_key) = create_test_rsa_keys();
849            let jwks = create_test_jwks(&public_key, "test-kid", Some("RS256"));
850
851            let metadata = HashMap::new();
852
853            let jwt = create_jwt_token(
854                &private_key,
855                "test-kid",
856                Algorithm::RS256,
857                "https://test-issuer.com",
858                Some("urn:risingwave:cluster:wrong-cluster-id"),
859                get_future_timestamp(),
860                HashMap::new(),
861            );
862
863            let result = validate_jwt_with_jwks(
864                &jwt,
865                &jwks,
866                "https://test-issuer.com",
867                "urn:risingwave:cluster:test-cluster-id",
868                &metadata,
869            );
870
871            let error = result.unwrap_err();
872            assert!(error.to_string().contains("InvalidAudience"));
873        }
874
875        #[test]
876        fn test_jwt_with_configured_audience() {
877            let (private_key, public_key) = create_test_rsa_keys();
878            let jwks = create_test_jwks(&public_key, "test-kid", Some("RS256"));
879
880            let mut additional_claims = HashMap::new();
881            additional_claims.insert("aud".to_owned(), json!(["urn:sn:cloud:o-for6u"]));
882            let jwt = create_jwt_token(
883                &private_key,
884                "test-kid",
885                Algorithm::RS256,
886                "https://test-issuer.com",
887                None,
888                get_future_timestamp(),
889                additional_claims,
890            );
891
892            let result = validate_jwt_with_jwks(
893                &jwt,
894                &jwks,
895                "https://test-issuer.com",
896                "urn:sn:cloud:o-for6u",
897                &HashMap::new(),
898            );
899
900            assert!(result.unwrap());
901        }
902
903        #[test]
904        fn test_empty_configured_audience() {
905            let error =
906                resolve_oauth_audience(Some("  ".to_owned()), "test-cluster-id").unwrap_err();
907            assert!(error.to_string().contains("cannot be empty"));
908        }
909
910        #[test]
911        fn test_default_audience_from_cluster_id() {
912            let audience = resolve_oauth_audience(None, "test-cluster-id").unwrap();
913            assert_eq!(audience, "urn:risingwave:cluster:test-cluster-id");
914        }
915
916        #[test]
917        fn test_jwt_with_missing_audience() {
918            let (private_key, public_key) = create_test_rsa_keys();
919            let jwks = create_test_jwks(&public_key, "test-kid", Some("RS256"));
920
921            let metadata = HashMap::new();
922
923            let jwt = create_jwt_token(
924                &private_key,
925                "test-kid",
926                Algorithm::RS256,
927                "https://test-issuer.com",
928                None, // No audience claim
929                get_future_timestamp(),
930                HashMap::new(),
931            );
932
933            let result = validate_jwt_with_jwks(
934                &jwt,
935                &jwks,
936                "https://test-issuer.com",
937                "urn:risingwave:cluster:test-cluster-id",
938                &metadata,
939            );
940
941            let error = result.unwrap_err();
942            assert!(error.to_string().contains("Missing required claim: aud"));
943        }
944
945        #[test]
946        fn test_jwt_with_invalid_issuer() {
947            let (private_key, public_key) = create_test_rsa_keys();
948            let jwks = create_test_jwks(&public_key, "test-kid", Some("RS256"));
949
950            let metadata = HashMap::new();
951
952            let jwt = create_jwt_token(
953                &private_key,
954                "test-kid",
955                Algorithm::RS256,
956                "https://wrong-issuer.com",
957                Some("urn:risingwave:cluster:test-cluster-id"),
958                get_future_timestamp(),
959                HashMap::new(),
960            );
961
962            let result = validate_jwt_with_jwks(
963                &jwt,
964                &jwks,
965                "https://test-issuer.com",
966                "urn:risingwave:cluster:test-cluster-id",
967                &metadata,
968            );
969
970            let error = result.unwrap_err();
971            assert!(error.to_string().contains("InvalidIssuer"));
972        }
973
974        #[test]
975        fn test_jwt_with_kid_not_found_in_jwks() {
976            let (private_key, public_key) = create_test_rsa_keys();
977            let jwks = create_test_jwks(&public_key, "different-kid", Some("RS256"));
978
979            let metadata = HashMap::new();
980
981            let jwt = create_jwt_token(
982                &private_key,
983                "missing-kid",
984                Algorithm::RS256,
985                "https://test-issuer.com",
986                Some("urn:risingwave:cluster:test-cluster-id"),
987                get_future_timestamp(),
988                HashMap::new(),
989            );
990
991            let result = validate_jwt_with_jwks(
992                &jwt,
993                &jwks,
994                "https://test-issuer.com",
995                "urn:risingwave:cluster:test-cluster-id",
996                &metadata,
997            );
998
999            let error = result.unwrap_err();
1000            assert!(
1001                error
1002                    .to_string()
1003                    .contains("No compatible RSA key found in JWKS for kid: 'missing-kid'")
1004            );
1005        }
1006
1007        #[test]
1008        fn test_jwt_with_empty_jwks_reports_no_matching_key() {
1009            let (private_key, _) = create_test_rsa_keys();
1010            let jwks = Jwks { keys: vec![] };
1011
1012            let jwt = create_jwt_token(
1013                &private_key,
1014                "missing-kid",
1015                Algorithm::RS256,
1016                "https://test-issuer.com",
1017                Some("urn:risingwave:cluster:test-cluster-id"),
1018                get_future_timestamp(),
1019                HashMap::new(),
1020            );
1021
1022            let result = validate_jwt_with_jwks(
1023                &jwt,
1024                &jwks,
1025                "https://test-issuer.com",
1026                "urn:risingwave:cluster:test-cluster-id",
1027                &HashMap::new(),
1028            );
1029
1030            let error = result.unwrap_err();
1031            assert!(
1032                error
1033                    .to_string()
1034                    .contains("No compatible RSA key found in JWKS for kid: 'missing-kid'")
1035            );
1036        }
1037
1038        #[test]
1039        fn test_jwt_with_expired_token() {
1040            let (private_key, public_key) = create_test_rsa_keys();
1041            let jwks = create_test_jwks(&public_key, "test-kid", Some("RS256"));
1042
1043            let metadata = HashMap::new();
1044
1045            let jwt = create_jwt_token(
1046                &private_key,
1047                "test-kid",
1048                Algorithm::RS256,
1049                "https://test-issuer.com",
1050                Some("urn:risingwave:cluster:test-cluster-id"),
1051                get_past_timestamp(), // Expired token
1052                HashMap::new(),
1053            );
1054
1055            let result = validate_jwt_with_jwks(
1056                &jwt,
1057                &jwks,
1058                "https://test-issuer.com",
1059                "urn:risingwave:cluster:test-cluster-id",
1060                &metadata,
1061            );
1062
1063            let error = result.unwrap_err();
1064            assert!(error.to_string().contains("ExpiredSignature"));
1065        }
1066
1067        #[test]
1068        fn test_jwt_with_invalid_signature() {
1069            let (_, public_key) = create_test_rsa_keys();
1070            let (wrong_private_key, _) = create_test_rsa_keys(); // Different key pair
1071            let jwks = create_test_jwks(&public_key, "test-kid", Some("RS256"));
1072
1073            let metadata = HashMap::new();
1074
1075            // Sign with wrong private key
1076            let jwt = create_jwt_token(
1077                &wrong_private_key,
1078                "test-kid",
1079                Algorithm::RS256,
1080                "https://test-issuer.com",
1081                Some("urn:risingwave:cluster:test-cluster-id"),
1082                get_future_timestamp(),
1083                HashMap::new(),
1084            );
1085
1086            let result = validate_jwt_with_jwks(
1087                &jwt,
1088                &jwks,
1089                "https://test-issuer.com",
1090                "urn:risingwave:cluster:test-cluster-id",
1091                &metadata,
1092            );
1093
1094            let error = result.unwrap_err();
1095            assert!(error.to_string().contains("InvalidSignature"));
1096        }
1097
1098        #[test]
1099        fn test_metadata_validation_success() {
1100            let (private_key, public_key) = create_test_rsa_keys();
1101            let jwks = create_test_jwks(&public_key, "test-kid", Some("RS256"));
1102
1103            let mut metadata = HashMap::new();
1104            metadata.insert("role".to_owned(), "admin".to_owned());
1105            metadata.insert("department".to_owned(), "security".to_owned());
1106
1107            let mut claims = HashMap::new();
1108            claims.insert("role".to_owned(), json!("admin"));
1109            claims.insert("department".to_owned(), json!("security"));
1110            claims.insert("extra_claim".to_owned(), json!("ignored")); // Extra claims are fine
1111
1112            let jwt = create_jwt_token(
1113                &private_key,
1114                "test-kid",
1115                Algorithm::RS256,
1116                "https://test-issuer.com",
1117                Some("urn:risingwave:cluster:test-cluster-id"),
1118                get_future_timestamp(),
1119                claims,
1120            );
1121
1122            let result = validate_jwt_with_jwks(
1123                &jwt,
1124                &jwks,
1125                "https://test-issuer.com",
1126                "urn:risingwave:cluster:test-cluster-id",
1127                &metadata,
1128            );
1129
1130            assert!(result.unwrap());
1131        }
1132
1133        #[test]
1134        fn test_metadata_validation_failure() {
1135            let (private_key, public_key) = create_test_rsa_keys();
1136            let jwks = create_test_jwks(&public_key, "test-kid", Some("RS256"));
1137
1138            let mut metadata = HashMap::new();
1139            metadata.insert("role".to_owned(), "admin".to_owned());
1140            metadata.insert("department".to_owned(), "security".to_owned());
1141
1142            let mut claims = HashMap::new();
1143            claims.insert("role".to_owned(), json!("user")); // Wrong role
1144            claims.insert("department".to_owned(), json!("security"));
1145
1146            let jwt = create_jwt_token(
1147                &private_key,
1148                "test-kid",
1149                Algorithm::RS256,
1150                "https://test-issuer.com",
1151                Some("urn:risingwave:cluster:test-cluster-id"),
1152                get_future_timestamp(),
1153                claims,
1154            );
1155
1156            let result = validate_jwt_with_jwks(
1157                &jwt,
1158                &jwks,
1159                "https://test-issuer.com",
1160                "urn:risingwave:cluster:test-cluster-id",
1161                &metadata,
1162            );
1163
1164            let error = result.unwrap_err();
1165            assert_eq!(
1166                error.to_string(),
1167                "metadata in jwt does not match with metadata declared with user"
1168            );
1169        }
1170
1171        #[test]
1172        fn test_jwt_with_jwk_missing_alg_succeeds() {
1173            let (private_key, public_key) = create_test_rsa_keys();
1174            let jwks = create_test_jwks(&public_key, "test-kid", None);
1175
1176            let jwt = create_jwt_token(
1177                &private_key,
1178                "test-kid",
1179                Algorithm::RS256,
1180                "https://test-issuer.com",
1181                Some("urn:risingwave:cluster:test-cluster-id"),
1182                get_future_timestamp(),
1183                HashMap::new(),
1184            );
1185
1186            let result = validate_jwt_with_jwks(
1187                &jwt,
1188                &jwks,
1189                "https://test-issuer.com",
1190                "urn:risingwave:cluster:test-cluster-id",
1191                &HashMap::new(),
1192            );
1193
1194            assert!(result.unwrap());
1195        }
1196
1197        #[test]
1198        fn test_jwt_with_mixed_jwks_ignores_non_matching_ec_key() {
1199            let (private_key, public_key) = create_test_rsa_keys();
1200            let mut jwks = create_test_jwks(&public_key, "rsa-kid", Some("RS256"));
1201            jwks.keys.push(Jwk {
1202                kty: Some("EC".to_owned()),
1203                kid: Some("ec-kid".to_owned()),
1204                alg: Some("ES256".to_owned()),
1205                n: None,
1206                e: None,
1207            });
1208
1209            let jwt = create_jwt_token(
1210                &private_key,
1211                "rsa-kid",
1212                Algorithm::RS256,
1213                "https://test-issuer.com",
1214                Some("urn:risingwave:cluster:test-cluster-id"),
1215                get_future_timestamp(),
1216                HashMap::new(),
1217            );
1218
1219            let result = validate_jwt_with_jwks(
1220                &jwt,
1221                &jwks,
1222                "https://test-issuer.com",
1223                "urn:risingwave:cluster:test-cluster-id",
1224                &HashMap::new(),
1225            );
1226
1227            assert!(result.unwrap());
1228        }
1229
1230        #[test]
1231        fn test_jwks_deserializes_ec_key_without_kid() {
1232            let jwks: Jwks = serde_json::from_value(json!({
1233                "keys": [
1234                    {
1235                        "alg": "ES256",
1236                        "kty": "EC",
1237                        "crv": "P-256",
1238                        "x": "x-coordinate",
1239                        "y": "y-coordinate"
1240                    },
1241                    {
1242                        "kid": "rsa-kid",
1243                        "alg": "RS256",
1244                        "kty": "RSA",
1245                        "n": "modulus",
1246                        "e": "AQAB"
1247                    }
1248                ]
1249            }))
1250            .unwrap();
1251
1252            assert_eq!(jwks.keys.len(), 2);
1253            assert!(jwks.keys[0].kid.is_none());
1254            assert_eq!(jwks.keys[1].kid.as_deref(), Some("rsa-kid"));
1255            assert_eq!(jwks.keys[1].n.as_deref(), Some("modulus"));
1256        }
1257
1258        #[test]
1259        fn test_jwt_uses_compatible_rsa_key_when_ec_key_has_same_kid() {
1260            let (private_key, public_key) = create_test_rsa_keys();
1261            let mut jwks = create_test_jwks(&public_key, "shared-kid", Some("RS256"));
1262            jwks.keys.insert(
1263                0,
1264                Jwk {
1265                    kty: Some("EC".to_owned()),
1266                    kid: Some("shared-kid".to_owned()),
1267                    alg: Some("ES256".to_owned()),
1268                    n: None,
1269                    e: None,
1270                },
1271            );
1272
1273            let jwt = create_jwt_token(
1274                &private_key,
1275                "shared-kid",
1276                Algorithm::RS256,
1277                "https://test-issuer.com",
1278                Some("urn:risingwave:cluster:test-cluster-id"),
1279                get_future_timestamp(),
1280                HashMap::new(),
1281            );
1282
1283            let result = validate_jwt_with_jwks(
1284                &jwt,
1285                &jwks,
1286                "https://test-issuer.com",
1287                "urn:risingwave:cluster:test-cluster-id",
1288                &HashMap::new(),
1289            );
1290
1291            assert!(result.unwrap());
1292        }
1293
1294        #[test]
1295        fn test_jwt_with_jwk_missing_alg_rejects_disallowed_header_alg() {
1296            let (_, public_key) = create_test_rsa_keys();
1297            let jwks = create_test_jwks(&public_key, "test-kid", None);
1298
1299            // Craft a token whose header claims HS256. The allow-list check
1300            // must reject it before the signature is ever verified — this is
1301            // the defence against the classic "alg=HS256 forged with the RSA
1302            // public key as the HMAC secret" confusion attack.
1303            let mut header = Header::new(Algorithm::HS256);
1304            header.kid = Some("test-kid".to_owned());
1305            let claims = json!({
1306                "iss": "https://test-issuer.com",
1307                "aud": "urn:risingwave:cluster:test-cluster-id",
1308                "exp": get_future_timestamp(),
1309            });
1310            let jwt = jsonwebtoken::encode(
1311                &header,
1312                &claims,
1313                &EncodingKey::from_secret(b"attacker-chosen"),
1314            )
1315            .unwrap();
1316
1317            let result = validate_jwt_with_jwks(
1318                &jwt,
1319                &jwks,
1320                "https://test-issuer.com",
1321                "urn:risingwave:cluster:test-cluster-id",
1322                &HashMap::new(),
1323            );
1324
1325            let error = result.unwrap_err();
1326            assert!(
1327                error.to_string().contains("is not allowed"),
1328                "unexpected error: {}",
1329                error
1330            );
1331        }
1332
1333        #[test]
1334        fn test_jwt_alg_mismatch_between_header_and_jwk() {
1335            let (private_key, public_key) = create_test_rsa_keys();
1336            // JWK pins RS384; token header declares RS256 — must be rejected.
1337            let jwks = create_test_jwks(&public_key, "test-kid", Some("RS384"));
1338
1339            let jwt = create_jwt_token(
1340                &private_key,
1341                "test-kid",
1342                Algorithm::RS256,
1343                "https://test-issuer.com",
1344                Some("urn:risingwave:cluster:test-cluster-id"),
1345                get_future_timestamp(),
1346                HashMap::new(),
1347            );
1348
1349            let result = validate_jwt_with_jwks(
1350                &jwt,
1351                &jwks,
1352                "https://test-issuer.com",
1353                "urn:risingwave:cluster:test-cluster-id",
1354                &HashMap::new(),
1355            );
1356
1357            let error = result.unwrap_err();
1358            assert_eq!(
1359                error.to_string(),
1360                "No compatible RSA key found in JWKS for kid: 'test-kid'"
1361            );
1362        }
1363    }
1364}
pgwire/pg_server.rs

pgwire/
pg_server.rs
