ydb_unofficial/auth/
sa.rs

1use std::sync::{Arc, RwLock};
2use std::time::{UNIX_EPOCH, SystemTime, Duration};
3
4use jwt_simple::prelude::{Claims, RSAKeyPairLike};
5pub use jwt_simple::prelude::PS256KeyPair;
6use serde::Deserialize;
7use tonic::transport::Uri;
8use yandex_cloud::yandex::cloud::iam::v1::CreateIamTokenResponse;
9
10use crate::AsciiValue;
11use super::{Credentials, UpdatableToken};
12
13#[derive(Debug, Clone, Deserialize)]
14pub struct ServiceAccountKey {
15    pub id: String,
16    pub service_account_id: String,
17    #[serde(with="ps256_private_key")]
18    pub private_key: PS256KeyPair,
19}
20
21#[derive(Debug, Clone)]
22pub struct UpdateConfig {
23    // Auth endpoint. Default is grpcs://iam.api.cloud.yandex.net:443
24    pub endpoint: Uri,
25    /// One of JWT cliams. Default is https://iam.api.cloud.yandex.net/iam/v1/tokens
26    pub audience: String,
27    /// Default update period. Used if received auth response without `expired_at`. Default is 50 minutes.
28    pub update_period: Duration,
29    /// Time reserve to update token. Default is 1 minute.
30    pub update_time_reserve: Duration,
31    /// JWT claim expiration. Default is 1 minute.
32    pub token_request_claim_time: Duration,
33}
34
35impl Default for UpdateConfig {
36    fn default() -> Self {
37        Self { 
38            endpoint: "grpcs://iam.api.cloud.yandex.net:443".parse().unwrap(), 
39            audience: "https://iam.api.cloud.yandex.net/iam/v1/tokens".into(), 
40            update_period: Duration::from_secs(50*60), 
41            update_time_reserve: Duration::from_secs(60),
42            token_request_claim_time: Duration::from_secs(60),
43        }
44    }
45}
46
47#[derive(Clone)]
48pub struct ServiceAccountCredentials {
49    token: Arc<RwLock<AsciiValue>>,
50}
51
52impl Credentials for ServiceAccountCredentials {
53    fn token(&self) -> crate::AsciiValue {
54        self.token.read().unwrap().clone()
55    }
56}
57impl Into<UpdatableToken> for ServiceAccountCredentials {
58    fn into(self) -> UpdatableToken {
59        let Self{token} = self;
60        UpdatableToken { token }
61    }
62}
63
64impl ServiceAccountCredentials {
65    pub async fn create(key: ServiceAccountKey) -> Result<Self, tonic::Status> {
66        Self::create_with_config(Default::default(), key).await
67    }
68    pub async fn create_with_config(conf: UpdateConfig, key: ServiceAccountKey) -> Result<Self, tonic::Status> {
69        //TODO: переделать на Stream
70        let response = conf.request_iam_token(&key).await?;
71        let mut sleep_duration = conf.invoke_sleep_duration(&response);
72        let token = Arc::new(RwLock::new(response.iam_token.clone().try_into().unwrap()));
73        let update_me = Arc::downgrade(&token);
74        tokio::spawn(async move {
75            loop {
76                tokio::time::sleep(sleep_duration).await;
77                if let Some(token) = update_me.upgrade() {
78                    match conf.request_iam_token(&key).await {
79                        Ok(response) => {
80                            sleep_duration = conf.invoke_sleep_duration(&response);
81                            *token.write().unwrap() = response.iam_token.clone().try_into().unwrap();
82                            log::info!("Iam token updated");
83                        }
84                        Err(e) => {
85                            log::error!("Cannot update iam token: {:?}", e);
86                            sleep_duration = Duration::from_secs(5);
87                        }
88                    }
89                } else {
90                    log::info!("ServiceAccountCredentials removed");
91                    break;
92                }
93            }
94        });
95        Ok(Self {token})
96    }
97}
98
99
100impl UpdateConfig {
101    pub async fn request_iam_token(&self, key: &ServiceAccountKey) -> Result<CreateIamTokenResponse, tonic::Status> {
102        let jwt = self.make_jwt(key);
103        let endpoint = crate::client::create_endpoint(self.endpoint.clone());
104        let mut client = yandex_cloud::yandex::cloud::iam::v1::iam_token_service_client::IamTokenServiceClient::new(endpoint.connect_lazy());
105        let request = yandex_cloud::yandex::cloud::iam::v1::CreateIamTokenRequest {
106            identity: Some(yandex_cloud::yandex::cloud::iam::v1::create_iam_token_request::Identity::Jwt(jwt))
107        };
108        let resp = client.create(request).await?;
109        Ok(resp.into_inner())
110    }
111
112    pub fn make_jwt(&self, key: &ServiceAccountKey) -> String {
113        let claims = Claims::create(self.token_request_claim_time.into())
114        .with_issuer(&key.service_account_id)
115        .with_audience(&self.audience);
116        let pair = key.private_key.clone().with_key_id(&key.id);
117        let token = pair.sign(claims).unwrap();
118        token
119    }
120    pub fn invoke_sleep_duration(&self, response: &CreateIamTokenResponse) -> tokio::time::Duration {
121        let CreateIamTokenResponse {iam_token: _, expires_at} = response;
122        let expires = if let Some(ts) = expires_at {
123            (UNIX_EPOCH + Duration::from_secs(ts.seconds as u64)) //convert ts to SystemTime
124            .duration_since(SystemTime::now() + self.update_time_reserve) //расчитываем время сна (с запасом)
125            .unwrap_or_default()
126        } else  {
127            self.update_period
128        };
129        expires.into()
130    }
131}
132
133mod ps256_private_key {
134    use jwt_simple::prelude::PS256KeyPair;
135    use serde::{self, Deserialize, Deserializer};
136    pub fn deserialize<'de, D>(
137        deserializer: D,
138    ) -> Result<PS256KeyPair, D::Error>
139    where
140    D: Deserializer<'de>,
141    {
142        let s = String::deserialize(deserializer)?;
143        PS256KeyPair::from_pem(&s).map_err(serde::de::Error::custom)
144    }
145}
ydb_unofficial/auth/sa.rs

ydb_unofficial/auth/
sa.rs
