ydb/
credentials.rs

1use crate::client::TimeoutSettings;
2use crate::errors::{YdbError, YdbResult};
3use crate::grpc_connection_manager::GrpcConnectionManager;
4use crate::grpc_wrapper::raw_auth_service::client::RawAuthClient;
5use crate::grpc_wrapper::raw_auth_service::login::RawLoginRequest;
6use crate::grpc_wrapper::runtime_interceptors::MultiInterceptor;
7use crate::load_balancer::{SharedLoadBalancer, StaticLoadBalancer};
8use crate::pub_traits::{Credentials, TokenInfo};
9use chrono::DateTime;
10use http::Uri;
11
12use jsonwebtoken::{encode, Algorithm, EncodingKey, Header};
13use secrecy::{ExposeSecret, SecretString};
14use serde::{Deserialize, Serialize};
15use std::env;
16use std::fmt::Debug;
17use std::ops::Add;
18use std::process::Command;
19use std::sync::{Arc, Mutex};
20use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
21use tracing::{debug, trace};
22
23const YDB_ANONYMOUS_CREDENTIALS: &str = "YDB_ANONYMOUS_CREDENTIALS";
24const YDB_SERVICE_ACCOUNT_KEY_FILE_CREDENTIALS: &str = "YDB_SERVICE_ACCOUNT_KEY_FILE_CREDENTIALS";
25const YDB_METADATA_CREDENTIALS: &str = "YDB_METADATA_CREDENTIALS";
26const YDB_ACCESS_TOKEN_CREDENTIALS: &str = "YDB_ACCESS_TOKEN_CREDENTIALS";
27
28const YC_METADATA_URL: &str =
29    "http://169.254.169.254/computeMetadata/v1/instance/service-accounts/default/token";
30const GCE_METADATA_URL: &str =
31    "http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/token";
32
33const EMPTY_TOKEN: &str = "";
34
35#[deprecated(note = "use AccessTokenCredentials instead")]
36pub type StaticToken = AccessTokenCredentials;
37#[deprecated(note = "use CommandLineCredentials instead")]
38pub type CommandLineYcToken = CommandLineCredentials;
39#[deprecated(note = "use StaticCredentials instead")]
40pub type StaticCredentialsAuth = StaticCredentials;
41#[deprecated(note = "use MetadataUrlCredentials instead")]
42pub type YandexMetadata = MetadataUrlCredentials;
43
44pub(crate) type CredentialsRef = Arc<Box<dyn Credentials>>;
45
46pub(crate) fn credencials_ref<T: 'static + Credentials>(cred: T) -> CredentialsRef {
47    Arc::new(Box::new(cred))
48}
49
50/// Get token of service account of instance
51///
52/// Yandex cloud support GCE token compatible. Use it.
53/// Example:
54/// ```
55/// use ydb::MetadataUrlCredentials;
56/// let cred = MetadataUrlCredentials::new();
57/// ```
58pub struct MetadataUrlCredentials {
59    inner: GCEMetadata,
60}
61
62impl MetadataUrlCredentials {
63    pub fn new() -> Self {
64        Self {
65            inner: GCEMetadata::from_url(YC_METADATA_URL).unwrap(),
66        }
67    }
68
69    /// Create GCEMetadata with custom url (may need for debug or spec infrastructure with non standard metadata)
70    ///
71    /// Example:
72    /// ```
73    /// # use ydb::YdbResult;
74    /// # fn main()->YdbResult<()>{
75    /// use ydb::MetadataUrlCredentials;
76    /// let cred = MetadataUrlCredentials::from_url("http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/token")?;
77    /// # return Ok(());
78    /// # }
79    /// ```
80    pub fn from_url<T: Into<String>>(url: T) -> YdbResult<Self> {
81        Ok(Self {
82            inner: GCEMetadata::from_url(url)?,
83        })
84    }
85}
86
87impl Default for MetadataUrlCredentials {
88    fn default() -> Self {
89        Self::new()
90    }
91}
92
93impl Credentials for MetadataUrlCredentials {
94    fn create_token(&self) -> YdbResult<TokenInfo> {
95        self.inner.create_token()
96    }
97}
98
99pub struct AnonymousCredentials {
100    inner: AccessTokenCredentials,
101}
102
103impl AnonymousCredentials {
104    pub fn new() -> Self {
105        Self {
106            inner: AccessTokenCredentials::from(EMPTY_TOKEN),
107        }
108    }
109}
110
111impl Default for AnonymousCredentials {
112    fn default() -> Self {
113        Self::new()
114    }
115}
116
117impl Credentials for AnonymousCredentials {
118    fn create_token(&self) -> YdbResult<TokenInfo> {
119        self.inner.create_token()
120    }
121}
122
123pub struct FromEnvCredentials {
124    inner: Box<dyn Credentials>,
125}
126
127/// Select credentials from environment
128/// reference: https://ydb.tech/docs/en/reference/ydb-sdk/auth
129impl FromEnvCredentials {
130    pub fn new() -> YdbResult<Self> {
131        Ok(Self {
132            inner: get_credentials_from_env()?,
133        })
134    }
135}
136
137impl Credentials for FromEnvCredentials {
138    fn create_token(&self) -> YdbResult<TokenInfo> {
139        self.inner.create_token()
140    }
141}
142
143fn get_credentials_from_env() -> YdbResult<Box<dyn Credentials>> {
144    if let Ok(file_creds) = env::var(YDB_SERVICE_ACCOUNT_KEY_FILE_CREDENTIALS) {
145        return Ok(Box::new(ServiceAccountCredentials::from_file(file_creds)?));
146    }
147
148    if let Ok(v) = env::var(YDB_ANONYMOUS_CREDENTIALS) {
149        if v == "1" {
150            return Ok(Box::new(
151                // anonymous credentials is empty token
152                AnonymousCredentials::new(),
153            ));
154        }
155    }
156
157    if let Ok(v) = env::var(YDB_METADATA_CREDENTIALS) {
158        if v == "1" {
159            return Ok(Box::new(MetadataUrlCredentials::new()));
160        }
161    }
162
163    if let Ok(token) = env::var(YDB_ACCESS_TOKEN_CREDENTIALS) {
164        return Ok(Box::new(AccessTokenCredentials::from(token)));
165    }
166
167    Ok(Box::new(MetadataUrlCredentials::new()))
168}
169
170/// Credentials with static token without renewing
171///
172/// Example:
173/// ```no_run
174/// # use ydb::{ClientBuilder, AccessTokenCredentials, YdbResult};
175/// # fn main()->YdbResult<()>{
176/// let builder = ClientBuilder::new_from_connection_string("grpc://localhost:2136?database=/local")?;
177/// let client = builder.with_credentials(AccessTokenCredentials::from("asd")).client()?;
178/// # return Ok(());
179/// # }
180/// ```
181#[derive(Debug, Clone)]
182pub struct AccessTokenCredentials {
183    pub(crate) token: String,
184}
185
186impl AccessTokenCredentials {
187    /// Create static token from string
188    ///
189    /// Example:
190    /// ```
191    /// # use ydb::AccessTokenCredentials;
192    /// AccessTokenCredentials::from("asd");
193    /// ```
194    pub fn from<T: Into<String>>(token: T) -> Self {
195        AccessTokenCredentials {
196            token: token.into(),
197        }
198    }
199}
200
201impl Credentials for AccessTokenCredentials {
202    fn create_token(&self) -> YdbResult<TokenInfo> {
203        Ok(TokenInfo::token(self.token.clone()))
204    }
205
206    fn debug_string(&self) -> String {
207        let (begin, end) = if self.token.len() > 20 {
208            (
209                &self.token.as_str()[0..3],
210                &self.token.as_str()[(self.token.len() - 3)..self.token.len()],
211            )
212        } else {
213            ("xxx", "xxx")
214        };
215        format!("static token: {begin}...{end}")
216    }
217}
218
219/// Get from stdout of command
220///
221/// Example create token from yandex cloud command line utility:
222/// ```rust
223/// use ydb::CommandLineCredentials;
224///
225/// let cred = CommandLineCredentials::from_cmd("yc iam create-token").unwrap();
226/// ```
227#[derive(Debug)]
228pub struct CommandLineCredentials {
229    command: Arc<Mutex<Command>>,
230}
231
232impl CommandLineCredentials {
233    /// Command line for create token
234    ///
235    /// The command will be called every time when token needed (token cache by default and will call rare).
236    #[allow(dead_code)]
237    pub fn from_cmd<T: Into<String>>(cmd: T) -> YdbResult<Self> {
238        let cmd = cmd.into();
239        let cmd_parts: Vec<&str> = cmd.split_whitespace().collect();
240
241        if cmd_parts.is_empty() {
242            return Err(YdbError::Custom(format!(
243                "can't split get token command: '{cmd}'"
244            )));
245        }
246
247        let mut command = Command::new(cmd_parts[0]);
248        command.args(&cmd_parts.as_slice()[1..]);
249
250        Ok(CommandLineCredentials {
251            command: Arc::new(Mutex::new(command)),
252        })
253    }
254}
255
256impl Credentials for CommandLineCredentials {
257    fn create_token(&self) -> YdbResult<TokenInfo> {
258        let result = self.command.lock()?.output()?;
259        if !result.status.success() {
260            let err = String::from_utf8(result.stderr)?;
261            return Err(YdbError::Custom(format!(
262                "can't execute yc ({}): {}",
263                result.status.code().unwrap(),
264                err
265            )));
266        }
267        let token = String::from_utf8(result.stdout)?.trim().to_string();
268        Ok(TokenInfo::token(token))
269    }
270
271    fn debug_string(&self) -> String {
272        let token_describe: String = match self.create_token() {
273            Ok(token_info) => {
274                let token = token_info.token.expose_secret();
275                let desc: String = if token.len() > 20 {
276                    format!(
277                        "{}..{}",
278                        &token.as_str()[0..3],
279                        &token.as_str()[(token.len() - 3)..token.len()]
280                    )
281                } else {
282                    "short_token".to_string()
283                };
284                desc
285            }
286            Err(err) => {
287                format!("err: {err}")
288            }
289        };
290
291        token_describe
292    }
293}
294
295/// Get service account credentials instance
296/// service account key should be:
297/// - in the local file and YDB_SERVICE_ACCOUNT_KEY_FILE_CREDENTIALS environment variable should point to it
298/// - in the local file and it's path is specified
299/// - in the json format string
300///
301/// Example:
302/// ```
303/// use ydb::ServiceAccountCredentials;
304/// let cred = ServiceAccountCredentials::from_env();
305/// ```
306/// or
307/// ```
308/// use ydb::ServiceAccountCredentials;
309/// let json = "....";
310/// let cred = ServiceAccountCredentials::from_json(json);
311/// ```
312/// or
313/// ```
314/// use ydb::ServiceAccountCredentials;
315/// let cred = ServiceAccountCredentials::new("service_account_id", "key_id", "private_key");
316/// ```
317/// or
318/// ```
319/// use ydb::ServiceAccountCredentials;
320/// let cred = ServiceAccountCredentials::from_file("/path/to/file");
321/// ```
322pub struct ServiceAccountCredentials {
323    audience_url: String,
324    private_key: SecretString,
325    service_account_id: String,
326    key_id: String,
327}
328
329impl ServiceAccountCredentials {
330    pub fn new(
331        service_account_id: impl Into<String>,
332        key_id: impl Into<String>,
333        private_key: impl Into<String>,
334    ) -> Self {
335        Self {
336            audience_url: Self::IAM_TOKEN_DEFAULT.to_string(),
337            private_key: SecretString::new(private_key.into()),
338            service_account_id: service_account_id.into(),
339            key_id: key_id.into(),
340        }
341    }
342
343    pub fn with_url(mut self, url: impl Into<String>) -> Self {
344        self.audience_url = url.into();
345        self
346    }
347
348    pub fn from_env() -> YdbResult<Self> {
349        let path = std::env::var(YDB_SERVICE_ACCOUNT_KEY_FILE_CREDENTIALS)?;
350
351        ServiceAccountCredentials::from_file(path)
352    }
353
354    pub fn from_file(path: impl AsRef<std::path::Path>) -> YdbResult<Self> {
355        let json_key = std::fs::read_to_string(path)?;
356        ServiceAccountCredentials::from_json(&json_key)
357    }
358
359    pub fn from_json(json_key: &str) -> YdbResult<Self> {
360        #[derive(Debug, Serialize, Deserialize)]
361        struct JsonKey {
362            public_key: String,
363            private_key: String,
364            service_account_id: String,
365            id: String,
366        }
367
368        let key: JsonKey = serde_json::from_str(json_key)?;
369
370        Ok(Self {
371            audience_url: Self::IAM_TOKEN_DEFAULT.to_string(),
372            key_id: key.id,
373            service_account_id: key.service_account_id,
374            private_key: SecretString::new(key.private_key),
375        })
376    }
377
378    const IAM_TOKEN_DEFAULT: &'static str = "https://iam.api.cloud.yandex.net/iam/v1/tokens";
379    const JWT_TOKEN_LIFE_TIME: usize = 720; // max 3600
380
381    fn build_jwt(&self) -> YdbResult<String> {
382        let private_key = self.private_key.expose_secret().as_bytes();
383
384        #[derive(Debug, Serialize, Deserialize)]
385        struct Claims {
386            aud: String, // Optional. Audience
387            exp: usize, // Required (validate_exp defaults to true in validation). Expiration time (as UTC timestamp)
388            iat: usize, // Optional. Issued at (as UTC timestamp)
389            iss: String, // Optional. Issuer
390        }
391
392        let iat = SystemTime::now()
393            .duration_since(UNIX_EPOCH)
394            .expect("Time went backwards")
395            .as_secs() as usize;
396
397        let mut header = Header::new(Algorithm::PS256);
398        header.kid = Some(self.key_id.clone());
399        header.alg = Algorithm::PS256;
400        header.typ = Some("JWT".to_string());
401
402        let claims = Claims {
403            exp: iat + Self::JWT_TOKEN_LIFE_TIME,
404            aud: self.audience_url.clone(),
405            iat,
406            iss: self.service_account_id.clone(),
407        };
408        let token = encode(
409            &header,
410            &claims,
411            &EncodingKey::from_rsa_pem(private_key).map_err(|e| YdbError::custom(e.to_string()))?,
412        )
413        .map_err(|e| YdbError::custom(format!("can't build jwt: {e}")))?;
414
415        debug!("Token was built");
416        Ok(token)
417    }
418
419    fn get_renew_time_for_lifetime(time: chrono::DateTime<chrono::Utc>) -> Instant {
420        let duration = time - chrono::Utc::now();
421        let seconds = (0.1 * duration.num_seconds() as f64) as u64;
422        trace!("renew in: {}", seconds);
423
424        Instant::now() + Duration::from_secs(seconds)
425    }
426}
427
428impl Credentials for ServiceAccountCredentials {
429    fn create_token(&self) -> YdbResult<TokenInfo> {
430        use chrono::Utc;
431        #[derive(Deserialize)]
432        struct TokenResponse {
433            #[serde(rename = "iamToken")]
434            iam_token: String,
435            #[serde(rename = "expiresAt")]
436            expires_at: DateTime<Utc>,
437        }
438
439        #[derive(Serialize)]
440        struct TokenRequest {
441            jwt: String,
442        }
443
444        let jwt = self.build_jwt()?;
445
446        let req = TokenRequest { jwt };
447        let client = reqwest::blocking::Client::new();
448        let res: TokenResponse = client
449            .request(reqwest::Method::POST, self.audience_url.clone())
450            .json(&req)
451            .send()?
452            .json()?;
453
454        Ok(TokenInfo::token(format!("Bearer {}", res.iam_token))
455            .with_renew(Self::get_renew_time_for_lifetime(res.expires_at)))
456    }
457}
458
459/// Get instance service account token from GCE instance
460///
461/// Get token from google cloud engine instance metadata.
462/// By default from url: http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/token"
463///
464/// Example:
465/// ```
466/// use ydb::GCEMetadata;
467///
468/// let cred = GCEMetadata::new();
469/// ```
470pub struct GCEMetadata {
471    uri: String,
472}
473
474impl GCEMetadata {
475    /// Create GCEMetadata with default url for receive token
476    pub fn new() -> Self {
477        Self::from_url(GCE_METADATA_URL).unwrap()
478    }
479
480    /// Create GCEMetadata with custom url (may need for debug or spec infrastructure with non standard metadata)
481    ///
482    /// Example:
483    /// ```
484    /// # use ydb::YdbResult;
485    /// # fn main()->YdbResult<()>{
486    /// use ydb::GCEMetadata;
487    /// let cred = GCEMetadata::from_url("http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/token")?;
488    /// # return Ok(());
489    /// # }
490    /// ```
491    pub fn from_url<T: Into<String>>(url: T) -> YdbResult<Self> {
492        Ok(Self {
493            uri: url.into().parse()?,
494        })
495    }
496}
497
498impl Default for GCEMetadata {
499    fn default() -> Self {
500        Self::new()
501    }
502}
503
504impl Credentials for GCEMetadata {
505    fn create_token(&self) -> YdbResult<TokenInfo> {
506        #[derive(Deserialize)]
507        struct TokenResponse {
508            access_token: String,
509            expires_in: u64,
510            token_type: String,
511        }
512
513        let client = reqwest::blocking::Client::new();
514        let res: TokenResponse = client
515            .request(reqwest::Method::GET, self.uri.as_str())
516            .header("Metadata-Flavor", "Google")
517            .send()?
518            .json()?;
519        Ok(
520            TokenInfo::token(format!("{} {}", res.token_type, res.access_token))
521                .with_renew(Instant::now().add(Duration::from_secs(res.expires_in))),
522        )
523    }
524
525    fn debug_string(&self) -> String {
526        format!("GoogleComputeEngineMetadata from {}", self.uri.as_str())
527    }
528}
529
530pub struct StaticCredentials {
531    username: String,
532    password: SecretString,
533    database: String,
534    endpoint: Uri,
535    cert_path: Option<String>,
536}
537
538impl StaticCredentials {
539    pub async fn acquire_token(&self) -> YdbResult<String> {
540        let static_balancer = StaticLoadBalancer::new(self.endpoint.clone());
541        let empty_connection_manager = GrpcConnectionManager::new(
542            SharedLoadBalancer::new_with_balancer(Box::new(static_balancer)),
543            self.database.clone(),
544            MultiInterceptor::new(),
545            self.cert_path.clone(),
546        );
547
548        let mut auth_client = empty_connection_manager
549            .get_auth_service(RawAuthClient::new)
550            .await
551            .unwrap();
552
553        // TODO: add configurable authorization request timeout
554        let raw_request = RawLoginRequest {
555            operation_params: TimeoutSettings::default().operation_params(),
556            user: self.username.clone(),
557            password: self.password.expose_secret().clone(),
558        };
559
560        let raw_response = auth_client.login(raw_request).await?;
561        Ok(raw_response.token)
562    }
563
564    pub fn new(username: String, password: String, endpoint: Uri, database: String) -> Self {
565        Self {
566            username,
567            password: SecretString::new(password),
568            database,
569            endpoint,
570            cert_path: None,
571        }
572    }
573
574    pub fn new_with_ca(
575        username: String,
576        password: String,
577        endpoint: Uri,
578        database: String,
579        cert_path: String,
580    ) -> Self {
581        Self {
582            username,
583            password: SecretString::new(password),
584            database,
585            endpoint,
586            cert_path: Some(cert_path),
587        }
588    }
589}
590
591impl Credentials for StaticCredentials {
592    #[tokio::main(flavor = "current_thread")]
593    async fn create_token(&self) -> YdbResult<TokenInfo> {
594        Ok(TokenInfo::token(self.acquire_token().await?))
595    }
596}
ydb/credentials.rs

ydb/
credentials.rs
