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                    redact_sql_option_keywords: 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                    redact_sql_option_keywords: 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 rsa::pkcs1::EncodeRsaPrivateKey;
767        use rsa::traits::PublicKeyParts;
768        use rsa::{RsaPrivateKey, RsaPublicKey};
769        use serde_json::json;
770
771        use crate::pg_server::{Jwk, Jwks, resolve_oauth_audience, validate_jwt_with_jwks};
772
773        fn create_test_rsa_keys() -> (RsaPrivateKey, RsaPublicKey) {
774            let mut rng = rand::thread_rng();
775            let private_key = RsaPrivateKey::new(&mut rng, 2048).expect("failed to generate a key");
776            let public_key = RsaPublicKey::from(&private_key);
777            (private_key, public_key)
778        }
779
780        fn create_test_jwks(public_key: &RsaPublicKey, kid: &str, alg: Option<&str>) -> Jwks {
781            let n = base64::engine::general_purpose::URL_SAFE_NO_PAD
782                .encode(public_key.n().to_bytes_be());
783            let e = base64::engine::general_purpose::URL_SAFE_NO_PAD
784                .encode(public_key.e().to_bytes_be());
785
786            Jwks {
787                keys: vec![Jwk {
788                    kty: Some("RSA".to_owned()),
789                    kid: Some(kid.to_owned()),
790                    alg: alg.map(ToOwned::to_owned),
791                    n: Some(n),
792                    e: Some(e),
793                }],
794            }
795        }
796
797        fn create_jwt_token(
798            private_key: &RsaPrivateKey,
799            kid: &str,
800            algorithm: Algorithm,
801            issuer: &str,
802            audience: Option<&str>,
803            exp: u64,
804            additional_claims: HashMap<String, serde_json::Value>,
805        ) -> String {
806            let mut header = Header::new(algorithm);
807            header.kid = Some(kid.to_owned());
808
809            let mut claims = json!({
810                "iss": issuer,
811                "exp": exp,
812            });
813
814            if let Some(aud) = audience {
815                claims["aud"] = json!(aud);
816            }
817
818            for (key, value) in additional_claims {
819                claims[key] = value;
820            }
821
822            let encoding_key = EncodingKey::from_rsa_pem(
823                private_key
824                    .to_pkcs1_pem(rsa::pkcs1::LineEnding::LF)
825                    .unwrap()
826                    .as_bytes(),
827            )
828            .unwrap();
829
830            jsonwebtoken::encode(&header, &claims, &encoding_key).unwrap()
831        }
832
833        fn get_future_timestamp() -> u64 {
834            SystemTime::now()
835                .duration_since(UNIX_EPOCH)
836                .unwrap()
837                .as_secs()
838                + 3600 // 1 hour from now
839        }
840
841        fn get_past_timestamp() -> u64 {
842            SystemTime::now()
843                .duration_since(UNIX_EPOCH)
844                .unwrap()
845                .as_secs()
846                - 3600 // 1 hour ago
847        }
848
849        #[test]
850        fn test_jwt_with_invalid_audience() {
851            let (private_key, public_key) = create_test_rsa_keys();
852            let jwks = create_test_jwks(&public_key, "test-kid", Some("RS256"));
853
854            let metadata = HashMap::new();
855
856            let jwt = create_jwt_token(
857                &private_key,
858                "test-kid",
859                Algorithm::RS256,
860                "https://test-issuer.com",
861                Some("urn:risingwave:cluster:wrong-cluster-id"),
862                get_future_timestamp(),
863                HashMap::new(),
864            );
865
866            let result = validate_jwt_with_jwks(
867                &jwt,
868                &jwks,
869                "https://test-issuer.com",
870                "urn:risingwave:cluster:test-cluster-id",
871                &metadata,
872            );
873
874            let error = result.unwrap_err();
875            assert!(error.to_string().contains("InvalidAudience"));
876        }
877
878        #[test]
879        fn test_jwt_with_configured_audience() {
880            let (private_key, public_key) = create_test_rsa_keys();
881            let jwks = create_test_jwks(&public_key, "test-kid", Some("RS256"));
882
883            let mut additional_claims = HashMap::new();
884            additional_claims.insert("aud".to_owned(), json!(["urn:sn:cloud:o-for6u"]));
885            let jwt = create_jwt_token(
886                &private_key,
887                "test-kid",
888                Algorithm::RS256,
889                "https://test-issuer.com",
890                None,
891                get_future_timestamp(),
892                additional_claims,
893            );
894
895            let result = validate_jwt_with_jwks(
896                &jwt,
897                &jwks,
898                "https://test-issuer.com",
899                "urn:sn:cloud:o-for6u",
900                &HashMap::new(),
901            );
902
903            assert!(result.unwrap());
904        }
905
906        #[test]
907        fn test_empty_configured_audience() {
908            let error =
909                resolve_oauth_audience(Some("  ".to_owned()), "test-cluster-id").unwrap_err();
910            assert!(error.to_string().contains("cannot be empty"));
911        }
912
913        #[test]
914        fn test_default_audience_from_cluster_id() {
915            let audience = resolve_oauth_audience(None, "test-cluster-id").unwrap();
916            assert_eq!(audience, "urn:risingwave:cluster:test-cluster-id");
917        }
918
919        #[test]
920        fn test_jwt_with_missing_audience() {
921            let (private_key, public_key) = create_test_rsa_keys();
922            let jwks = create_test_jwks(&public_key, "test-kid", Some("RS256"));
923
924            let metadata = HashMap::new();
925
926            let jwt = create_jwt_token(
927                &private_key,
928                "test-kid",
929                Algorithm::RS256,
930                "https://test-issuer.com",
931                None, // No audience claim
932                get_future_timestamp(),
933                HashMap::new(),
934            );
935
936            let result = validate_jwt_with_jwks(
937                &jwt,
938                &jwks,
939                "https://test-issuer.com",
940                "urn:risingwave:cluster:test-cluster-id",
941                &metadata,
942            );
943
944            let error = result.unwrap_err();
945            assert!(error.to_string().contains("Missing required claim: aud"));
946        }
947
948        #[test]
949        fn test_jwt_with_invalid_issuer() {
950            let (private_key, public_key) = create_test_rsa_keys();
951            let jwks = create_test_jwks(&public_key, "test-kid", Some("RS256"));
952
953            let metadata = HashMap::new();
954
955            let jwt = create_jwt_token(
956                &private_key,
957                "test-kid",
958                Algorithm::RS256,
959                "https://wrong-issuer.com",
960                Some("urn:risingwave:cluster:test-cluster-id"),
961                get_future_timestamp(),
962                HashMap::new(),
963            );
964
965            let result = validate_jwt_with_jwks(
966                &jwt,
967                &jwks,
968                "https://test-issuer.com",
969                "urn:risingwave:cluster:test-cluster-id",
970                &metadata,
971            );
972
973            let error = result.unwrap_err();
974            assert!(error.to_string().contains("InvalidIssuer"));
975        }
976
977        #[test]
978        fn test_jwt_with_kid_not_found_in_jwks() {
979            let (private_key, public_key) = create_test_rsa_keys();
980            let jwks = create_test_jwks(&public_key, "different-kid", Some("RS256"));
981
982            let metadata = HashMap::new();
983
984            let jwt = create_jwt_token(
985                &private_key,
986                "missing-kid",
987                Algorithm::RS256,
988                "https://test-issuer.com",
989                Some("urn:risingwave:cluster:test-cluster-id"),
990                get_future_timestamp(),
991                HashMap::new(),
992            );
993
994            let result = validate_jwt_with_jwks(
995                &jwt,
996                &jwks,
997                "https://test-issuer.com",
998                "urn:risingwave:cluster:test-cluster-id",
999                &metadata,
1000            );
1001
1002            let error = result.unwrap_err();
1003            assert!(
1004                error
1005                    .to_string()
1006                    .contains("No compatible RSA key found in JWKS for kid: 'missing-kid'")
1007            );
1008        }
1009
1010        #[test]
1011        fn test_jwt_with_empty_jwks_reports_no_matching_key() {
1012            let (private_key, _) = create_test_rsa_keys();
1013            let jwks = Jwks { keys: vec![] };
1014
1015            let jwt = create_jwt_token(
1016                &private_key,
1017                "missing-kid",
1018                Algorithm::RS256,
1019                "https://test-issuer.com",
1020                Some("urn:risingwave:cluster:test-cluster-id"),
1021                get_future_timestamp(),
1022                HashMap::new(),
1023            );
1024
1025            let result = validate_jwt_with_jwks(
1026                &jwt,
1027                &jwks,
1028                "https://test-issuer.com",
1029                "urn:risingwave:cluster:test-cluster-id",
1030                &HashMap::new(),
1031            );
1032
1033            let error = result.unwrap_err();
1034            assert!(
1035                error
1036                    .to_string()
1037                    .contains("No compatible RSA key found in JWKS for kid: 'missing-kid'")
1038            );
1039        }
1040
1041        #[test]
1042        fn test_jwt_with_expired_token() {
1043            let (private_key, public_key) = create_test_rsa_keys();
1044            let jwks = create_test_jwks(&public_key, "test-kid", Some("RS256"));
1045
1046            let metadata = HashMap::new();
1047
1048            let jwt = create_jwt_token(
1049                &private_key,
1050                "test-kid",
1051                Algorithm::RS256,
1052                "https://test-issuer.com",
1053                Some("urn:risingwave:cluster:test-cluster-id"),
1054                get_past_timestamp(), // Expired token
1055                HashMap::new(),
1056            );
1057
1058            let result = validate_jwt_with_jwks(
1059                &jwt,
1060                &jwks,
1061                "https://test-issuer.com",
1062                "urn:risingwave:cluster:test-cluster-id",
1063                &metadata,
1064            );
1065
1066            let error = result.unwrap_err();
1067            assert!(error.to_string().contains("ExpiredSignature"));
1068        }
1069
1070        #[test]
1071        fn test_jwt_with_invalid_signature() {
1072            let (_, public_key) = create_test_rsa_keys();
1073            let (wrong_private_key, _) = create_test_rsa_keys(); // Different key pair
1074            let jwks = create_test_jwks(&public_key, "test-kid", Some("RS256"));
1075
1076            let metadata = HashMap::new();
1077
1078            // Sign with wrong private key
1079            let jwt = create_jwt_token(
1080                &wrong_private_key,
1081                "test-kid",
1082                Algorithm::RS256,
1083                "https://test-issuer.com",
1084                Some("urn:risingwave:cluster:test-cluster-id"),
1085                get_future_timestamp(),
1086                HashMap::new(),
1087            );
1088
1089            let result = validate_jwt_with_jwks(
1090                &jwt,
1091                &jwks,
1092                "https://test-issuer.com",
1093                "urn:risingwave:cluster:test-cluster-id",
1094                &metadata,
1095            );
1096
1097            let error = result.unwrap_err();
1098            assert!(error.to_string().contains("InvalidSignature"));
1099        }
1100
1101        #[test]
1102        fn test_metadata_validation_success() {
1103            let (private_key, public_key) = create_test_rsa_keys();
1104            let jwks = create_test_jwks(&public_key, "test-kid", Some("RS256"));
1105
1106            let mut metadata = HashMap::new();
1107            metadata.insert("role".to_owned(), "admin".to_owned());
1108            metadata.insert("department".to_owned(), "security".to_owned());
1109
1110            let mut claims = HashMap::new();
1111            claims.insert("role".to_owned(), json!("admin"));
1112            claims.insert("department".to_owned(), json!("security"));
1113            claims.insert("extra_claim".to_owned(), json!("ignored")); // Extra claims are fine
1114
1115            let jwt = create_jwt_token(
1116                &private_key,
1117                "test-kid",
1118                Algorithm::RS256,
1119                "https://test-issuer.com",
1120                Some("urn:risingwave:cluster:test-cluster-id"),
1121                get_future_timestamp(),
1122                claims,
1123            );
1124
1125            let result = validate_jwt_with_jwks(
1126                &jwt,
1127                &jwks,
1128                "https://test-issuer.com",
1129                "urn:risingwave:cluster:test-cluster-id",
1130                &metadata,
1131            );
1132
1133            assert!(result.unwrap());
1134        }
1135
1136        #[test]
1137        fn test_metadata_validation_failure() {
1138            let (private_key, public_key) = create_test_rsa_keys();
1139            let jwks = create_test_jwks(&public_key, "test-kid", Some("RS256"));
1140
1141            let mut metadata = HashMap::new();
1142            metadata.insert("role".to_owned(), "admin".to_owned());
1143            metadata.insert("department".to_owned(), "security".to_owned());
1144
1145            let mut claims = HashMap::new();
1146            claims.insert("role".to_owned(), json!("user")); // Wrong role
1147            claims.insert("department".to_owned(), json!("security"));
1148
1149            let jwt = create_jwt_token(
1150                &private_key,
1151                "test-kid",
1152                Algorithm::RS256,
1153                "https://test-issuer.com",
1154                Some("urn:risingwave:cluster:test-cluster-id"),
1155                get_future_timestamp(),
1156                claims,
1157            );
1158
1159            let result = validate_jwt_with_jwks(
1160                &jwt,
1161                &jwks,
1162                "https://test-issuer.com",
1163                "urn:risingwave:cluster:test-cluster-id",
1164                &metadata,
1165            );
1166
1167            let error = result.unwrap_err();
1168            assert_eq!(
1169                error.to_string(),
1170                "metadata in jwt does not match with metadata declared with user"
1171            );
1172        }
1173
1174        #[test]
1175        fn test_jwt_with_jwk_missing_alg_succeeds() {
1176            let (private_key, public_key) = create_test_rsa_keys();
1177            let jwks = create_test_jwks(&public_key, "test-kid", None);
1178
1179            let jwt = create_jwt_token(
1180                &private_key,
1181                "test-kid",
1182                Algorithm::RS256,
1183                "https://test-issuer.com",
1184                Some("urn:risingwave:cluster:test-cluster-id"),
1185                get_future_timestamp(),
1186                HashMap::new(),
1187            );
1188
1189            let result = validate_jwt_with_jwks(
1190                &jwt,
1191                &jwks,
1192                "https://test-issuer.com",
1193                "urn:risingwave:cluster:test-cluster-id",
1194                &HashMap::new(),
1195            );
1196
1197            assert!(result.unwrap());
1198        }
1199
1200        #[test]
1201        fn test_jwt_with_mixed_jwks_ignores_non_matching_ec_key() {
1202            let (private_key, public_key) = create_test_rsa_keys();
1203            let mut jwks = create_test_jwks(&public_key, "rsa-kid", Some("RS256"));
1204            jwks.keys.push(Jwk {
1205                kty: Some("EC".to_owned()),
1206                kid: Some("ec-kid".to_owned()),
1207                alg: Some("ES256".to_owned()),
1208                n: None,
1209                e: None,
1210            });
1211
1212            let jwt = create_jwt_token(
1213                &private_key,
1214                "rsa-kid",
1215                Algorithm::RS256,
1216                "https://test-issuer.com",
1217                Some("urn:risingwave:cluster:test-cluster-id"),
1218                get_future_timestamp(),
1219                HashMap::new(),
1220            );
1221
1222            let result = validate_jwt_with_jwks(
1223                &jwt,
1224                &jwks,
1225                "https://test-issuer.com",
1226                "urn:risingwave:cluster:test-cluster-id",
1227                &HashMap::new(),
1228            );
1229
1230            assert!(result.unwrap());
1231        }
1232
1233        #[test]
1234        fn test_jwks_deserializes_ec_key_without_kid() {
1235            let jwks: Jwks = serde_json::from_value(json!({
1236                "keys": [
1237                    {
1238                        "alg": "ES256",
1239                        "kty": "EC",
1240                        "crv": "P-256",
1241                        "x": "x-coordinate",
1242                        "y": "y-coordinate"
1243                    },
1244                    {
1245                        "kid": "rsa-kid",
1246                        "alg": "RS256",
1247                        "kty": "RSA",
1248                        "n": "modulus",
1249                        "e": "AQAB"
1250                    }
1251                ]
1252            }))
1253            .unwrap();
1254
1255            assert_eq!(jwks.keys.len(), 2);
1256            assert!(jwks.keys[0].kid.is_none());
1257            assert_eq!(jwks.keys[1].kid.as_deref(), Some("rsa-kid"));
1258            assert_eq!(jwks.keys[1].n.as_deref(), Some("modulus"));
1259        }
1260
1261        #[test]
1262        fn test_jwt_uses_compatible_rsa_key_when_ec_key_has_same_kid() {
1263            let (private_key, public_key) = create_test_rsa_keys();
1264            let mut jwks = create_test_jwks(&public_key, "shared-kid", Some("RS256"));
1265            jwks.keys.insert(
1266                0,
1267                Jwk {
1268                    kty: Some("EC".to_owned()),
1269                    kid: Some("shared-kid".to_owned()),
1270                    alg: Some("ES256".to_owned()),
1271                    n: None,
1272                    e: None,
1273                },
1274            );
1275
1276            let jwt = create_jwt_token(
1277                &private_key,
1278                "shared-kid",
1279                Algorithm::RS256,
1280                "https://test-issuer.com",
1281                Some("urn:risingwave:cluster:test-cluster-id"),
1282                get_future_timestamp(),
1283                HashMap::new(),
1284            );
1285
1286            let result = validate_jwt_with_jwks(
1287                &jwt,
1288                &jwks,
1289                "https://test-issuer.com",
1290                "urn:risingwave:cluster:test-cluster-id",
1291                &HashMap::new(),
1292            );
1293
1294            assert!(result.unwrap());
1295        }
1296
1297        #[test]
1298        fn test_jwt_with_jwk_missing_alg_rejects_disallowed_header_alg() {
1299            let (_, public_key) = create_test_rsa_keys();
1300            let jwks = create_test_jwks(&public_key, "test-kid", None);
1301
1302            // Craft a token whose header claims HS256. The allow-list check
1303            // must reject it before the signature is ever verified — this is
1304            // the defence against the classic "alg=HS256 forged with the RSA
1305            // public key as the HMAC secret" confusion attack.
1306            let mut header = Header::new(Algorithm::HS256);
1307            header.kid = Some("test-kid".to_owned());
1308            let claims = json!({
1309                "iss": "https://test-issuer.com",
1310                "aud": "urn:risingwave:cluster:test-cluster-id",
1311                "exp": get_future_timestamp(),
1312            });
1313            let jwt = jsonwebtoken::encode(
1314                &header,
1315                &claims,
1316                &EncodingKey::from_secret(b"attacker-chosen"),
1317            )
1318            .unwrap();
1319
1320            let result = validate_jwt_with_jwks(
1321                &jwt,
1322                &jwks,
1323                "https://test-issuer.com",
1324                "urn:risingwave:cluster:test-cluster-id",
1325                &HashMap::new(),
1326            );
1327
1328            let error = result.unwrap_err();
1329            assert!(
1330                error.to_string().contains("is not allowed"),
1331                "unexpected error: {}",
1332                error
1333            );
1334        }
1335
1336        #[test]
1337        fn test_jwt_alg_mismatch_between_header_and_jwk() {
1338            let (private_key, public_key) = create_test_rsa_keys();
1339            // JWK pins RS384; token header declares RS256 — must be rejected.
1340            let jwks = create_test_jwks(&public_key, "test-kid", Some("RS384"));
1341
1342            let jwt = create_jwt_token(
1343                &private_key,
1344                "test-kid",
1345                Algorithm::RS256,
1346                "https://test-issuer.com",
1347                Some("urn:risingwave:cluster:test-cluster-id"),
1348                get_future_timestamp(),
1349                HashMap::new(),
1350            );
1351
1352            let result = validate_jwt_with_jwks(
1353                &jwt,
1354                &jwks,
1355                "https://test-issuer.com",
1356                "urn:risingwave:cluster:test-cluster-id",
1357                &HashMap::new(),
1358            );
1359
1360            let error = result.unwrap_err();
1361            assert_eq!(
1362                error.to_string(),
1363                "No compatible RSA key found in JWKS for kid: 'test-kid'"
1364            );
1365        }
1366    }
1367}
pgwire/pg_server.rs

pgwire/
pg_server.rs
