rocketmq_client_rust/common/
session_credentials.rs

1// Copyright 2023 The RocketMQ Rust Authors
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::collections::HashMap;
16use std::fmt;
17use std::fs;
18use std::path::PathBuf;
19
20use cheetah_string::CheetahString;
21
22pub const ACCESS_KEY: &str = "AccessKey";
23pub const SECRET_KEY: &str = "SecretKey";
24pub const SIGNATURE: &str = "Signature";
25pub const SECURITY_TOKEN: &str = "SecurityToken";
26
27fn get_key_file_path() -> PathBuf {
28    if let Ok(key_file) = std::env::var("rocketmq.client.keyFile") {
29        PathBuf::from(key_file)
30    } else if let Ok(home) = std::env::var("HOME").or_else(|_| std::env::var("USERPROFILE")) {
31        PathBuf::from(home).join("key")
32    } else {
33        PathBuf::from("key")
34    }
35}
36
37#[derive(Debug, Clone, Default)]
38pub struct SessionCredentials {
39    access_key: Option<CheetahString>,
40    secret_key: Option<CheetahString>,
41    security_token: Option<CheetahString>,
42    signature: Option<CheetahString>,
43}
44
45impl SessionCredentials {
46    pub fn new() -> Self {
47        let mut credentials = Self::default();
48        credentials.load_from_file();
49        credentials
50    }
51
52    pub fn with_keys(access_key: impl Into<CheetahString>, secret_key: impl Into<CheetahString>) -> Self {
53        Self {
54            access_key: Some(access_key.into()),
55            secret_key: Some(secret_key.into()),
56            security_token: None,
57            signature: None,
58        }
59    }
60
61    pub fn with_token(
62        access_key: impl Into<CheetahString>,
63        secret_key: impl Into<CheetahString>,
64        security_token: impl Into<CheetahString>,
65    ) -> Self {
66        Self {
67            access_key: Some(access_key.into()),
68            secret_key: Some(secret_key.into()),
69            security_token: Some(security_token.into()),
70            signature: None,
71        }
72    }
73
74    fn load_from_file(&mut self) {
75        let key_file = get_key_file_path();
76        if let Ok(content) = fs::read_to_string(key_file) {
77            if let Ok(properties) = self.parse_properties(&content) {
78                self.update_content(&properties);
79            }
80        }
81    }
82
83    fn parse_properties(&self, content: &str) -> Result<HashMap<String, String>, ()> {
84        let mut properties = HashMap::new();
85        for line in content.lines() {
86            let line = line.trim();
87            if line.is_empty() || line.starts_with('#') {
88                continue;
89            }
90            if let Some(pos) = line.find('=') {
91                let key = line[..pos].trim().to_string();
92                let value = line[pos + 1..].trim().to_string();
93                properties.insert(key, value);
94            }
95        }
96        Ok(properties)
97    }
98
99    pub fn update_content(&mut self, properties: &HashMap<String, String>) {
100        if let Some(value) = properties.get(ACCESS_KEY) {
101            self.access_key = Some(CheetahString::from_string(value.trim().to_string()));
102        }
103        if let Some(value) = properties.get(SECRET_KEY) {
104            self.secret_key = Some(CheetahString::from_string(value.trim().to_string()));
105        }
106        if let Some(value) = properties.get(SECURITY_TOKEN) {
107            self.security_token = Some(CheetahString::from_string(value.trim().to_string()));
108        }
109    }
110
111    pub fn access_key(&self) -> Option<&CheetahString> {
112        self.access_key.as_ref()
113    }
114
115    pub fn set_access_key(&mut self, access_key: impl Into<CheetahString>) {
116        self.access_key = Some(access_key.into());
117    }
118
119    pub fn secret_key(&self) -> Option<&CheetahString> {
120        self.secret_key.as_ref()
121    }
122
123    pub fn set_secret_key(&mut self, secret_key: impl Into<CheetahString>) {
124        self.secret_key = Some(secret_key.into());
125    }
126
127    pub fn signature(&self) -> Option<&CheetahString> {
128        self.signature.as_ref()
129    }
130
131    pub fn set_signature(&mut self, signature: impl Into<CheetahString>) {
132        self.signature = Some(signature.into());
133    }
134
135    pub fn security_token(&self) -> Option<&CheetahString> {
136        self.security_token.as_ref()
137    }
138
139    pub fn set_security_token(&mut self, security_token: impl Into<CheetahString>) {
140        self.security_token = Some(security_token.into());
141    }
142}
143
144impl PartialEq for SessionCredentials {
145    fn eq(&self, other: &Self) -> bool {
146        self.access_key == other.access_key && self.secret_key == other.secret_key && self.signature == other.signature
147    }
148}
149
150impl Eq for SessionCredentials {}
151
152impl std::hash::Hash for SessionCredentials {
153    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
154        self.access_key.hash(state);
155        self.secret_key.hash(state);
156        self.signature.hash(state);
157    }
158}
159
160impl fmt::Display for SessionCredentials {
161    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
162        write!(
163            f,
164            "SessionCredentials [accessKey={:?}, secretKey={:?}, signature={:?}, SecurityToken={:?}]",
165            self.access_key, self.secret_key, self.signature, self.security_token
166        )
167    }
168}
169
170#[cfg(test)]
171mod tests {
172    use super::*;
173
174    #[test]
175    fn session_credentials_new_and_default() {
176        let credentials = SessionCredentials::new();
177        assert!(credentials.access_key().is_none());
178        let credentials = SessionCredentials::default();
179        assert!(credentials.secret_key().is_none());
180    }
181
182    #[test]
183    fn session_credentials_with_keys() {
184        let access_key = "ak";
185        let secret_key = "sk";
186        let credentials = SessionCredentials::with_keys(access_key, secret_key);
187        assert_eq!(credentials.access_key().unwrap(), access_key);
188        assert_eq!(credentials.secret_key().unwrap(), secret_key);
189        assert!(credentials.security_token().is_none());
190        assert!(credentials.signature().is_none());
191    }
192
193    #[test]
194    fn session_credentials_with_token() {
195        let access_key = "ak";
196        let secret_key = "sk";
197        let token = "token";
198        let credentials = SessionCredentials::with_token(access_key, secret_key, token);
199        assert_eq!(credentials.access_key().unwrap(), access_key);
200        assert_eq!(credentials.secret_key().unwrap(), secret_key);
201        assert_eq!(credentials.security_token().unwrap(), token);
202        assert!(credentials.signature().is_none());
203    }
204
205    #[test]
206    fn session_credentials_parse_properties() {
207        let content = "AccessKey=ak123\nSecretKey=sk456\n#comment\n \nSecurityToken=tk789\nSignature=sig123";
208        let credentials = SessionCredentials::default();
209        let properties = credentials.parse_properties(content).unwrap();
210        assert_eq!(properties.get("AccessKey").unwrap(), "ak123");
211        assert_eq!(properties.get("SecretKey").unwrap(), "sk456");
212        assert_eq!(properties.get("SecurityToken").unwrap(), "tk789");
213        assert_eq!(properties.get("Signature").unwrap(), "sig123");
214    }
215
216    #[test]
217    fn session_credentials_update_content() {
218        let mut properties = HashMap::new();
219        properties.insert(ACCESS_KEY.to_string(), "ak".to_string());
220        properties.insert(SECRET_KEY.to_string(), "sk".to_string());
221        properties.insert(SECURITY_TOKEN.to_string(), "token".to_string());
222
223        let mut credentials = SessionCredentials::default();
224        credentials.update_content(&properties);
225
226        assert_eq!(credentials.access_key().unwrap(), "ak");
227        assert_eq!(credentials.secret_key().unwrap(), "sk");
228        assert_eq!(credentials.security_token().unwrap(), "token");
229    }
230
231    #[test]
232    fn session_credentials_getters_and_setters() {
233        let mut credentials = SessionCredentials::default();
234        credentials.set_access_key("ak");
235        credentials.set_secret_key("sk");
236        credentials.set_security_token("token");
237        credentials.set_signature("sig");
238
239        assert_eq!(credentials.access_key().unwrap(), "ak");
240        assert_eq!(credentials.secret_key().unwrap(), "sk");
241        assert_eq!(credentials.security_token().unwrap(), "token");
242        assert_eq!(credentials.signature().unwrap(), "sig");
243    }
244
245    #[test]
246    fn session_credentials_eq_and_hash() {
247        let c1 = SessionCredentials::with_keys("ak", "sk");
248        let c2 = SessionCredentials::with_keys("ak", "sk");
249        let c3 = SessionCredentials::with_keys("ak2", "sk");
250
251        assert_eq!(c1, c2);
252        assert_ne!(c1, c3);
253
254        use std::collections::hash_map::DefaultHasher;
255        use std::hash::Hash;
256        use std::hash::Hasher;
257        let mut h1 = DefaultHasher::new();
258        c1.hash(&mut h1);
259        let mut h2 = DefaultHasher::new();
260        c2.hash(&mut h2);
261        assert_eq!(h1.finish(), h2.finish());
262    }
263
264    #[test]
265    fn session_credentials_display() {
266        let credentials = SessionCredentials::with_token("ak", "sk", "tk");
267        let display = format!("{}", credentials);
268        let expected = "SessionCredentials [accessKey=Some(\"ak\"), secretKey=Some(\"sk\"), signature=None, \
269                        SecurityToken=Some(\"tk\")]";
270        assert_eq!(display, expected);
271    }
272}
rocketmq_client_rust/common/session_credentials.rs

rocketmq_client_rust/common/
session_credentials.rs
