pgwire/api/auth/
md5pass.rs

1use std::fmt::Debug;
2use std::sync::Arc;
3
4use async_trait::async_trait;
5use futures::sink::{Sink, SinkExt};
6use tokio::sync::Mutex;
7
8use super::{
9    AuthSource, ClientInfo, LoginInfo, PgWireConnectionState, ServerParameterProvider,
10    StartupHandler,
11};
12use crate::api::{ConnectionManager, PidSecretKeyGenerator, RandomPidSecretKeyGenerator};
13use crate::error::{PgWireError, PgWireResult};
14use crate::messages::startup::Authentication;
15use crate::messages::{PgWireBackendMessage, PgWireFrontendMessage};
16
17/// Startup handler for MD5 password authentication.
18pub struct Md5PasswordAuthStartupHandler<A, P> {
19    auth_source: Arc<A>,
20    parameter_provider: Arc<P>,
21    pid_secret_key_generator: Arc<dyn PidSecretKeyGenerator>,
22    connection_manager: Option<Arc<ConnectionManager>>,
23    cached_password: Mutex<Vec<u8>>,
24}
25
26impl<A, P> Md5PasswordAuthStartupHandler<A, P> {
27    /// Creates a new MD5 password auth handler.
28    pub fn new(auth_source: Arc<A>, parameter_provider: Arc<P>) -> Self {
29        Md5PasswordAuthStartupHandler {
30            auth_source,
31            parameter_provider,
32            pid_secret_key_generator: Arc::new(RandomPidSecretKeyGenerator::default()),
33            connection_manager: None,
34            cached_password: Mutex::new(vec![]),
35        }
36    }
37
38    /// Sets a custom PID/secret key generator.
39    pub fn with_pid_secret_key_generator(
40        mut self,
41        generator: Arc<dyn PidSecretKeyGenerator>,
42    ) -> Self {
43        self.pid_secret_key_generator = generator;
44        self
45    }
46
47    /// Sets a connection manager.
48    pub fn with_connection_manager(mut self, manager: Arc<ConnectionManager>) -> Self {
49        self.connection_manager = Some(manager);
50        self
51    }
52}
53
54#[async_trait]
55impl<A: AuthSource, P: ServerParameterProvider> StartupHandler
56    for Md5PasswordAuthStartupHandler<A, P>
57{
58    async fn on_startup<C>(
59        &self,
60        client: &mut C,
61        message: PgWireFrontendMessage,
62    ) -> PgWireResult<()>
63    where
64        C: ClientInfo + Sink<PgWireBackendMessage> + Unpin + Send,
65        C::Error: Debug,
66        PgWireError: From<<C as Sink<PgWireBackendMessage>>::Error>,
67    {
68        match message {
69            PgWireFrontendMessage::Startup(ref startup) => {
70                super::protocol_negotiation(client, startup).await?;
71                super::save_startup_parameters_to_metadata(client, startup);
72                client.set_state(PgWireConnectionState::AuthenticationInProgress);
73
74                let login_info = LoginInfo::from_client_info(client);
75                let salt_and_pass = self.auth_source.get_password(&login_info).await?;
76
77                let salt = salt_and_pass
78                    .salt
79                    .as_ref()
80                    .expect("Salt is required for Md5Password authentication");
81
82                self.cached_password
83                    .lock()
84                    .await
85                    .clone_from(&salt_and_pass.password);
86
87                client
88                    .send(PgWireBackendMessage::Authentication(
89                        Authentication::MD5Password(salt.clone()),
90                    ))
91                    .await?;
92            }
93            PgWireFrontendMessage::PasswordMessageFamily(pwd) => {
94                let pwd = pwd.into_password()?;
95                let cached_pass = self.cached_password.lock().await;
96
97                if pwd.password.as_bytes() == *cached_pass {
98                    let (pid, secret_key) = self.pid_secret_key_generator.generate(client);
99                    client.set_pid_and_secret_key(pid, secret_key);
100                    if let Some(manager) = &self.connection_manager {
101                        super::register_connection(client, manager);
102                    }
103                    super::finish_authentication(client, self.parameter_provider.as_ref()).await?;
104                } else {
105                    let login_info = LoginInfo::from_client_info(client);
106                    return Err(PgWireError::InvalidPassword(
107                        login_info.user().map(|x| x.to_owned()).unwrap_or_default(),
108                    ));
109                }
110            }
111            _ => {}
112        }
113        Ok(())
114    }
115}
116
117/// This function is to compute postgres standard md5 hashed password
118///
119/// concat('md5', md5(concat(md5(concat(password, username)), random-salt)))
120///
121/// the input parameter `md5hashed_username_password` represents
122/// `md5(concat(password, username))` so that your can store hashed password in
123/// storage.
124pub fn hash_md5_password(username: &str, password: &str, salt: &[u8]) -> String {
125    let hashed_bytes = format!("{:x}", md5::compute(format!("{password}{username}")));
126    let mut bytes = Vec::with_capacity(hashed_bytes.len() + 4);
127    bytes.extend_from_slice(hashed_bytes.as_ref());
128    bytes.extend_from_slice(salt);
129
130    format!("md5{:x}", md5::compute(bytes))
131}
132
133#[cfg(test)]
134mod tests {
135
136    #[test]
137    fn test_hash_md5_passwd() {
138        let salt = vec![20, 247, 107, 249];
139        let username = "zmjiang";
140        let password = "themanwhochangedchina";
141
142        let result = "md521fe459d77d3e3ea9c9fcd5c11030d30";
143
144        assert_eq!(result, super::hash_md5_password(username, password, &salt));
145    }
146}
pgwire/api/auth/md5pass.rs

pgwire/api/auth/
md5pass.rs
