pgwire/api/auth/
sasl.rs

1use std::fmt::Debug;
2use std::sync::Arc;
3
4use async_trait::async_trait;
5use futures::{Sink, SinkExt};
6use tokio::sync::Mutex;
7
8use crate::api::auth::sasl::scram::ScramServerAuthWaitingForClientFinal;
9use crate::api::{
10    ClientInfo, ConnectionManager, PgWireConnectionState, PidSecretKeyGenerator,
11    RandomPidSecretKeyGenerator,
12};
13use crate::error::{PgWireError, PgWireResult};
14use crate::messages::startup::{Authentication, PasswordMessageFamily};
15use crate::messages::{PgWireBackendMessage, PgWireFrontendMessage};
16
17use super::{ServerParameterProvider, StartupHandler};
18
19pub mod oauth;
20pub mod scram;
21
22/// SASL mechanism name for SCRAM-SHA-256.
23pub const SCRAM_SHA_256_METHOD: &str = "SCRAM-SHA-256";
24/// SASL mechanism name for SCRAM-SHA-256-PLUS with channel binding.
25pub const SCRAM_SHA_256_PLUS_METHOD: &str = "SCRAM-SHA-256-PLUS";
26/// SASL mechanism name for OAUTHBEARER.
27pub const OAUTHBEARER_METHOD: &str = "OAUTHBEARER";
28
29/// States of the SASL authentication state machine.
30#[derive(Debug)]
31pub enum SASLState {
32    Initial,
33    // scram authentication method selected
34    ScramClientFirstReceived,
35    // cached password, channel_binding and partial auth-message
36    ScramServerFirstSent(Box<ScramServerAuthWaitingForClientFinal>),
37    // oauth authentication method selected
38    OauthStateInit,
39    // failure during authentication
40    OauthStateError,
41    // finished
42    Finished,
43}
44
45impl SASLState {
46    fn is_scram(&self) -> bool {
47        matches!(
48            self,
49            SASLState::ScramClientFirstReceived | SASLState::ScramServerFirstSent(_)
50        )
51    }
52
53    fn is_oauth(&self) -> bool {
54        matches!(self, SASLState::OauthStateInit | SASLState::OauthStateError)
55    }
56}
57
58/// Startup handler for SASL-based authentication (SCRAM and OAUTHBEARER).
59pub struct SASLAuthStartupHandler<P> {
60    parameter_provider: Arc<P>,
61    pid_secret_key_generator: Arc<dyn PidSecretKeyGenerator>,
62    connection_manager: Option<Arc<ConnectionManager>>,
63    /// state of the SASL auth
64    state: Mutex<SASLState>,
65    /// scram configuration
66    scram: Option<scram::ScramAuth>,
67    /// oauth configuration
68    oauth: Option<oauth::Oauth>,
69}
70
71impl<P> SASLAuthStartupHandler<P> {
72    /// Creates a new SASL auth handler.
73    pub fn new(parameter_provider: Arc<P>) -> Self {
74        SASLAuthStartupHandler {
75            parameter_provider,
76            pid_secret_key_generator: Arc::new(RandomPidSecretKeyGenerator::default()),
77            connection_manager: None,
78            state: Mutex::new(SASLState::Initial),
79            scram: None,
80            oauth: None,
81        }
82    }
83
84    /// Configures SCRAM authentication.
85    pub fn with_scram(mut self, scram_auth: scram::ScramAuth) -> Self {
86        self.scram = Some(scram_auth);
87        self
88    }
89
90    /// Configures OAuth authentication.
91    pub fn with_oauth(mut self, oauth: oauth::Oauth) -> Self {
92        self.oauth = Some(oauth);
93        self
94    }
95
96    /// Sets a custom PID/secret key generator.
97    pub fn with_pid_secret_key_generator(
98        mut self,
99        generator: Arc<dyn PidSecretKeyGenerator>,
100    ) -> Self {
101        self.pid_secret_key_generator = generator;
102        self
103    }
104
105    /// Sets a connection manager.
106    pub fn with_connection_manager(mut self, manager: Arc<ConnectionManager>) -> Self {
107        self.connection_manager = Some(manager);
108        self
109    }
110
111    fn supported_mechanisms(&self) -> Vec<String> {
112        let mut mechanisms = vec![];
113
114        if let Some(scram) = &self.scram {
115            mechanisms.push(SCRAM_SHA_256_METHOD.to_owned());
116
117            if scram.supports_channel_binding() {
118                mechanisms.push(SCRAM_SHA_256_PLUS_METHOD.to_owned());
119            }
120        }
121
122        if self.oauth.is_some() {
123            mechanisms.push(OAUTHBEARER_METHOD.to_owned());
124        }
125
126        mechanisms
127    }
128}
129
130#[async_trait]
131impl<P: ServerParameterProvider> StartupHandler for SASLAuthStartupHandler<P> {
132    async fn on_startup<C>(
133        &self,
134        client: &mut C,
135        message: PgWireFrontendMessage,
136    ) -> PgWireResult<()>
137    where
138        C: ClientInfo + Sink<PgWireBackendMessage> + Unpin + Send + Sync,
139        C::Error: Debug,
140        PgWireError: From<<C as Sink<PgWireBackendMessage>>::Error>,
141    {
142        match message {
143            PgWireFrontendMessage::Startup(ref startup) => {
144                super::protocol_negotiation(client, startup).await?;
145                super::save_startup_parameters_to_metadata(client, startup);
146                client.set_state(PgWireConnectionState::AuthenticationInProgress);
147                let supported_mechanisms = self.supported_mechanisms();
148                client
149                    .send(PgWireBackendMessage::Authentication(Authentication::SASL(
150                        supported_mechanisms,
151                    )))
152                    .await?;
153            }
154            PgWireFrontendMessage::PasswordMessageFamily(mut msg) => {
155                let mut state = self.state.lock().await;
156
157                msg = if let SASLState::Initial = *state {
158                    let sasl_initial_response = msg.into_sasl_initial_response()?;
159                    let selected_mechanism = sasl_initial_response.auth_method.as_str();
160
161                    *state = if [SCRAM_SHA_256_METHOD, SCRAM_SHA_256_PLUS_METHOD]
162                        .contains(&selected_mechanism)
163                    {
164                        SASLState::ScramClientFirstReceived
165                    } else if OAUTHBEARER_METHOD == selected_mechanism {
166                        SASLState::OauthStateInit
167                    } else {
168                        return Err(PgWireError::UnsupportedSASLAuthMethod(
169                            selected_mechanism.to_string(),
170                        ));
171                    };
172
173                    PasswordMessageFamily::SASLInitialResponse(sasl_initial_response)
174                } else {
175                    let sasl_response = msg.into_sasl_response()?;
176                    PasswordMessageFamily::SASLResponse(sasl_response)
177                };
178
179                if state.is_scram() {
180                    let scram = self.scram.as_ref().ok_or_else(|| {
181                        PgWireError::UnsupportedSASLAuthMethod("SCRAM".to_string())
182                    })?;
183                    let (res, new_state) = scram.process_scram_message(client, msg, &state).await?;
184                    client
185                        .send(PgWireBackendMessage::Authentication(res))
186                        .await?;
187                    *state = new_state;
188                } else if state.is_oauth() {
189                    let oauth = self.oauth.as_ref().ok_or_else(|| {
190                        PgWireError::UnsupportedSASLAuthMethod("OAUTHBEARER".to_string())
191                    })?;
192                    let (res, new_state) = oauth.process_oauth_message(client, msg, &state).await?;
193                    if let Some(res) = res {
194                        client
195                            .send(PgWireBackendMessage::Authentication(res))
196                            .await?;
197                    }
198                    *state = new_state;
199                } else {
200                    return Err(PgWireError::InvalidSASLState);
201                };
202
203                if matches!(*state, SASLState::Finished) {
204                    let (pid, secret_key) = self.pid_secret_key_generator.generate(client);
205                    client.set_pid_and_secret_key(pid, secret_key);
206                    if let Some(manager) = &self.connection_manager {
207                        super::register_connection(client, manager);
208                    }
209                    super::finish_authentication(client, self.parameter_provider.as_ref()).await?;
210                }
211            }
212            _ => {}
213        }
214
215        Ok(())
216    }
217}
pgwire/api/auth/sasl.rs

pgwire/api/auth/
sasl.rs
