auth/user_provider/
watch_file_user_provider.rs

1// Copyright 2023 Greptime Team
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::sync::{Arc, Mutex};
16
17use async_trait::async_trait;
18use common_config::file_watcher::{FileWatcherBuilder, FileWatcherConfig};
19use common_telemetry::{info, warn};
20use snafu::ResultExt;
21
22use crate::error::{FileWatchSnafu, Result};
23use crate::user_provider::{
24    PgAuthInfo, UserInfoMap, authenticate_with_credential, load_credential_from_file,
25    postgres_auth_info_with_credential,
26};
27use crate::{Identity, Password, UserInfoRef, UserProvider};
28
29pub(crate) const WATCH_FILE_USER_PROVIDER: &str = "watch_file_user_provider";
30
31type WatchedCredentialRef = Arc<Mutex<UserInfoMap>>;
32
33/// A user provider that reads user credential from a file and watches the file for changes.
34///
35/// Both empty file and non-existent file are invalid and will cause initialization to fail.
36#[derive(Debug)]
37pub(crate) struct WatchFileUserProvider {
38    users: WatchedCredentialRef,
39}
40
41impl WatchFileUserProvider {
42    pub fn new(filepath: &str) -> Result<Self> {
43        let credential = load_credential_from_file(filepath)?;
44        let users = Arc::new(Mutex::new(credential));
45
46        let users_clone = users.clone();
47        let filepath_owned = filepath.to_string();
48
49        FileWatcherBuilder::new()
50            .watch_path(filepath)
51            .context(FileWatchSnafu)?
52            .config(FileWatcherConfig::new())
53            .spawn(move || match load_credential_from_file(&filepath_owned) {
54                Ok(credential) => {
55                    let mut users = users_clone.lock().expect("users credential must be valid");
56                    #[cfg(not(test))]
57                    info!("User provider file {} reloaded", &filepath_owned);
58                    #[cfg(test)]
59                    info!(
60                        "User provider file {} reloaded: {:?}",
61                        &filepath_owned, credential
62                    );
63                    *users = credential;
64                }
65                Err(err) => {
66                    warn!(
67                        ?err,
68                        "Fail to load credential from file {}; keep the old one", &filepath_owned
69                    )
70                }
71            })
72            .context(FileWatchSnafu)?;
73
74        Ok(WatchFileUserProvider { users })
75    }
76}
77
78#[async_trait]
79impl UserProvider for WatchFileUserProvider {
80    fn name(&self) -> &str {
81        WATCH_FILE_USER_PROVIDER
82    }
83
84    async fn authenticate(&self, id: Identity<'_>, password: Password<'_>) -> Result<UserInfoRef> {
85        let users = self.users.lock().expect("users credential must be valid");
86        authenticate_with_credential(&users, id, password)
87    }
88
89    async fn postgres_auth_info(&self, id: Identity<'_>) -> Result<PgAuthInfo> {
90        let users = self.users.lock().expect("users credential must be valid");
91        postgres_auth_info_with_credential(&users, id)
92    }
93
94    async fn authorize(&self, _: &str, _: &str, _: &UserInfoRef) -> Result<()> {
95        // default allow all
96        Ok(())
97    }
98}
99
100#[cfg(test)]
101pub mod test {
102    use std::time::{Duration, Instant};
103
104    use common_test_util::temp_dir::create_temp_dir;
105    use tokio::time::sleep;
106
107    use crate::UserProvider;
108    use crate::user_provider::watch_file_user_provider::WatchFileUserProvider;
109    use crate::user_provider::{Identity, Password};
110
111    async fn test_authenticate(
112        provider: &dyn UserProvider,
113        username: &str,
114        password: &str,
115        ok: bool,
116        timeout: Option<Duration>,
117    ) {
118        if let Some(timeout) = timeout {
119            let deadline = Instant::now().checked_add(timeout).unwrap();
120            loop {
121                let re = provider
122                    .authenticate(
123                        Identity::UserId(username, None),
124                        Password::PlainText(password.to_string().into()),
125                    )
126                    .await;
127                if re.is_ok() == ok {
128                    break;
129                } else if Instant::now() < deadline {
130                    sleep(Duration::from_millis(100)).await;
131                } else {
132                    panic!("timeout (username: {username}, password: {password}, expected: {ok})");
133                }
134            }
135        } else {
136            let re = provider
137                .authenticate(
138                    Identity::UserId(username, None),
139                    Password::PlainText(password.to_string().into()),
140                )
141                .await;
142            assert_eq!(
143                re.is_ok(),
144                ok,
145                "username: {}, password: {}",
146                username,
147                password
148            );
149        }
150    }
151
152    #[tokio::test]
153    async fn test_file_provider_initialization_with_missing_file() {
154        common_telemetry::init_default_ut_logging();
155
156        let dir = create_temp_dir("test_missing_file");
157        let file_path = format!("{}/non_existent_file", dir.path().to_str().unwrap());
158
159        // Try to create provider with non-existent file should fail
160        let result = WatchFileUserProvider::new(file_path.as_str());
161        assert!(result.is_err());
162
163        let error = result.unwrap_err();
164        assert!(error.to_string().contains("UserProvider file must exist"));
165    }
166
167    #[tokio::test]
168    async fn test_file_provider() {
169        common_telemetry::init_default_ut_logging();
170
171        let dir = create_temp_dir("test_file_provider");
172        let file_path = format!("{}/test_file_provider", dir.path().to_str().unwrap());
173
174        // write a tmp file
175        assert!(std::fs::write(&file_path, "root=123456\nadmin=654321\n").is_ok());
176        let provider = WatchFileUserProvider::new(file_path.as_str()).unwrap();
177        let timeout = Duration::from_secs(60);
178
179        test_authenticate(&provider, "root", "123456", true, None).await;
180        test_authenticate(&provider, "admin", "654321", true, None).await;
181        test_authenticate(&provider, "root", "654321", false, None).await;
182
183        // update the tmp file
184        assert!(std::fs::write(&file_path, "root=654321\n").is_ok());
185        test_authenticate(&provider, "root", "123456", false, Some(timeout)).await;
186        test_authenticate(&provider, "root", "654321", true, Some(timeout)).await;
187        test_authenticate(&provider, "admin", "654321", false, Some(timeout)).await;
188
189        // remove the tmp file
190        assert!(std::fs::remove_file(&file_path).is_ok());
191        // When file is deleted during runtime, keep the last known good credentials
192        test_authenticate(&provider, "root", "654321", true, Some(timeout)).await;
193        test_authenticate(&provider, "root", "123456", false, Some(timeout)).await;
194        test_authenticate(&provider, "admin", "654321", false, Some(timeout)).await;
195
196        // recreate the tmp file
197        assert!(std::fs::write(&file_path, "root=123456\n").is_ok());
198        test_authenticate(&provider, "root", "123456", true, Some(timeout)).await;
199        test_authenticate(&provider, "root", "654321", false, Some(timeout)).await;
200        test_authenticate(&provider, "admin", "654321", false, Some(timeout)).await;
201    }
202}
auth/user_provider/watch_file_user_provider.rs

auth/user_provider/
watch_file_user_provider.rs
