sqlserver_mcp_catalog/auth/
auth_manager.rs1use crate::core::config_schema::AuthMethod;
4use crate::core::credential_storage::{load_credential, save_credential};
5
6use super::auth_strategy::{AuthConfig, AuthStrategy, Credentials};
7use super::errors::AuthError;
8use super::strategies::azure_ad::AzureAdAuthStrategy;
9use super::strategies::sql_server::SqlServerAuthStrategy;
10use super::strategies::windows::WindowsAuthStrategy;
11
12const CREDENTIAL_ACCOUNT: &str = "active-credentials";
13const ENV_PREFIX: &str = "SQLSERVER";
14
15fn credentials_from_env(auth_method: AuthMethod) -> Option<AuthConfig> {
22 let mut config = AuthConfig::new();
23 match auth_method {
24 AuthMethod::SqlServer | AuthMethod::Windows => {
25 let username = std::env::var(format!("{ENV_PREFIX}_USERNAME")).ok()?;
26 let password = std::env::var(format!("{ENV_PREFIX}_PASSWORD")).ok()?;
27 config.insert("username".to_string(), username);
28 config.insert("password".to_string(), password);
29 }
30 AuthMethod::AzureAd => {
31 let client_id = std::env::var(format!("{ENV_PREFIX}_CLIENT_ID")).ok()?;
32 let client_secret = std::env::var(format!("{ENV_PREFIX}_CLIENT_SECRET")).ok()?;
33 config.insert("client_id".to_string(), client_id);
34 config.insert("client_secret".to_string(), client_secret);
35 if let Ok(tenant_id) = std::env::var(format!("{ENV_PREFIX}_TENANT_ID")) {
36 config.insert("tenant_id".to_string(), tenant_id);
37 }
38 }
39 }
40 Some(config)
41}
42
43fn strategy_for(auth_method: AuthMethod) -> Box<dyn AuthStrategy> {
44 match auth_method {
45 AuthMethod::SqlServer => Box::new(SqlServerAuthStrategy),
46 AuthMethod::Windows => Box::new(WindowsAuthStrategy),
47 AuthMethod::AzureAd => Box::new(AzureAdAuthStrategy),
48 }
49}
50
51pub struct AuthManager {
62 auth_method: AuthMethod,
63 strategy: Box<dyn AuthStrategy>,
64 cached_credentials: Option<Credentials>,
65}
66
67impl AuthManager {
68 pub fn new(auth_method: AuthMethod) -> Self {
69 Self {
70 strategy: strategy_for(auth_method),
71 auth_method,
72 cached_credentials: None,
73 }
74 }
75
76 pub async fn login(&mut self, config: &AuthConfig) -> anyhow::Result<Credentials> {
77 let credentials = self.strategy.authenticate(config).await?;
78 self.cached_credentials = Some(credentials.clone());
79 save_credential(CREDENTIAL_ACCOUNT, &serde_json::to_string(&credentials)?)?;
80 Ok(credentials)
81 }
82
83 pub fn set_credentials(&mut self, credentials: Credentials) {
88 self.cached_credentials = Some(credentials);
89 }
90
91 pub async fn credentials(&mut self) -> anyhow::Result<Credentials> {
92 if let Some(cached) = &self.cached_credentials
93 && self.strategy.validate_credentials(cached)
94 {
95 return Ok(cached.clone());
96 }
97
98 if let Some(env_config) = credentials_from_env(self.auth_method)
99 && let Ok(from_env) = self.strategy.authenticate(&env_config).await
100 && self.strategy.validate_credentials(&from_env)
101 {
102 self.cached_credentials = Some(from_env.clone());
103 return Ok(from_env);
104 }
105
106 if let Some(stored) = load_credential(CREDENTIAL_ACCOUNT)? {
107 let parsed: Credentials = serde_json::from_str(&stored)?;
108 if self.strategy.validate_credentials(&parsed) {
109 self.cached_credentials = Some(parsed.clone());
110 return Ok(parsed);
111 }
112 if let Ok(refreshed) = self.strategy.refresh_token(&parsed).await {
113 self.cached_credentials = Some(refreshed.clone());
114 save_credential(CREDENTIAL_ACCOUNT, &serde_json::to_string(&refreshed)?)?;
115 return Ok(refreshed);
116 }
117 }
118
119 Err(AuthError::NoActiveCredentials(format!("{:?}", self.auth_method)).into())
120 }
121
122 pub async fn resolve_tds_auth(&mut self) -> anyhow::Result<tiberius::AuthMethod> {
125 let credentials = self.credentials().await?;
126 match self.auth_method {
127 AuthMethod::SqlServer => {
128 let username = credentials
129 .get("username")
130 .ok_or_else(|| AuthError::MissingCredentials("username".to_string()))?;
131 let password = credentials
132 .get("password")
133 .ok_or_else(|| AuthError::MissingCredentials("password".to_string()))?;
134 Ok(tiberius::AuthMethod::sql_server(username, password))
135 }
136 AuthMethod::Windows => {
137 let username = credentials
138 .get("username")
139 .ok_or_else(|| AuthError::MissingCredentials("username".to_string()))?;
140 let password = credentials
141 .get("password")
142 .ok_or_else(|| AuthError::MissingCredentials("password".to_string()))?;
143 windows_auth_method(username, password)
144 }
145 AuthMethod::AzureAd => {
146 let token = credentials
147 .get("access_token")
148 .ok_or_else(|| AuthError::MissingCredentials("access_token".to_string()))?;
149 Ok(tiberius::AuthMethod::AADToken(token.clone()))
150 }
151 }
152 }
153}
154
155#[cfg(windows)]
165fn windows_auth_method(username: &str, password: &str) -> anyhow::Result<tiberius::AuthMethod> {
166 Ok(tiberius::AuthMethod::windows(username, password))
167}
168
169#[cfg(not(windows))]
170fn windows_auth_method(_username: &str, _password: &str) -> anyhow::Result<tiberius::AuthMethod> {
171 anyhow::bail!(
172 "Windows Authentication requires tiberius's native SSPI binding, which is only compiled \
173 on Windows targets; this server is running on a non-Windows target with no \
174 non-Windows NTLM/Kerberos implementation linked in. Use `sql_server` or `azure_ad` \
175 auth instead, or add Kerberos support (tiberius's `integrated-auth-gssapi` feature, \
176 which requires the system `libgssapi` library) if Windows Authentication against a \
177 non-Windows client is required."
178 )
179}
180
181#[cfg(test)]
182mod tests {
183 use super::*;
184
185 #[tokio::test]
186 async fn seeded_credentials_win_over_stored_ones() {
187 let mut manager = AuthManager::new(AuthMethod::SqlServer);
188 let mut credentials = Credentials::new();
189 credentials.insert("username".to_string(), "sa".to_string());
190 credentials.insert("password".to_string(), "s3cr3t".to_string());
191 manager.set_credentials(credentials.clone());
192
193 let resolved = manager.credentials().await.unwrap();
194 assert_eq!(resolved.get("username").map(String::as_str), Some("sa"));
195 }
196
197 #[tokio::test]
198 async fn resolve_tds_auth_builds_a_sql_server_auth_method() {
199 let mut manager = AuthManager::new(AuthMethod::SqlServer);
200 let mut credentials = Credentials::new();
201 credentials.insert("username".to_string(), "sa".to_string());
202 credentials.insert("password".to_string(), "s3cr3t".to_string());
203 manager.set_credentials(credentials);
204
205 let auth = manager.resolve_tds_auth().await.unwrap();
206 assert!(matches!(auth, tiberius::AuthMethod::SqlServer(_)));
207 }
208
209 #[tokio::test]
210 async fn resolve_tds_auth_builds_an_azure_ad_auth_method() {
211 let mut manager = AuthManager::new(AuthMethod::AzureAd);
212 let mut credentials = Credentials::new();
213 credentials.insert("access_token".to_string(), "eyJ...".to_string());
214 manager.set_credentials(credentials);
215
216 let auth = manager.resolve_tds_auth().await.unwrap();
217 assert!(matches!(auth, tiberius::AuthMethod::AADToken(_)));
218 }
219
220 #[tokio::test]
221 async fn resolve_tds_auth_errors_without_any_credentials() {
222 let mut manager = AuthManager::new(AuthMethod::SqlServer);
223 assert!(manager.resolve_tds_auth().await.is_err());
224 }
225}