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(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_slice(value.trim()));
102        }
103        if let Some(value) = properties.get(SECRET_KEY) {
104            self.secret_key = Some(CheetahString::from_slice(value.trim()));
105        }
106        if let Some(value) = properties.get(SECURITY_TOKEN) {
107            self.security_token = Some(CheetahString::from_slice(value.trim()));
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
152fn redacted_if_present<T>(value: &Option<T>) -> &'static str {
153    if value.is_some() {
154        "<redacted>"
155    } else {
156        "<unset>"
157    }
158}
159
160impl fmt::Debug for SessionCredentials {
161    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
162        f.debug_struct("SessionCredentials")
163            .field("access_key", &redacted_if_present(&self.access_key))
164            .field("secret_key", &redacted_if_present(&self.secret_key))
165            .field("signature", &redacted_if_present(&self.signature))
166            .field("security_token", &redacted_if_present(&self.security_token))
167            .finish()
168    }
169}
170
171impl std::hash::Hash for SessionCredentials {
172    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
173        self.access_key.hash(state);
174        self.secret_key.hash(state);
175        self.signature.hash(state);
176    }
177}
178
179impl fmt::Display for SessionCredentials {
180    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
181        write!(
182            f,
183            "SessionCredentials [accessKey={}, secretKey={}, signature={}, securityToken={}]",
184            redacted_if_present(&self.access_key),
185            redacted_if_present(&self.secret_key),
186            redacted_if_present(&self.signature),
187            redacted_if_present(&self.security_token)
188        )
189    }
190}
191
192#[cfg(test)]
193mod tests {
194    use super::*;
195
196    #[test]
197    fn session_credentials_new_and_default() {
198        let credentials = SessionCredentials::new();
199        assert!(credentials.access_key().is_none());
200        let credentials = SessionCredentials::default();
201        assert!(credentials.secret_key().is_none());
202    }
203
204    #[test]
205    fn session_credentials_with_keys() {
206        let access_key = "ak";
207        let secret_key = "sk";
208        let credentials = SessionCredentials::with_keys(access_key, secret_key);
209        assert_eq!(credentials.access_key().unwrap(), access_key);
210        assert_eq!(credentials.secret_key().unwrap(), secret_key);
211        assert!(credentials.security_token().is_none());
212        assert!(credentials.signature().is_none());
213    }
214
215    #[test]
216    fn session_credentials_with_token() {
217        let access_key = "ak";
218        let secret_key = "sk";
219        let token = "token";
220        let credentials = SessionCredentials::with_token(access_key, secret_key, token);
221        assert_eq!(credentials.access_key().unwrap(), access_key);
222        assert_eq!(credentials.secret_key().unwrap(), secret_key);
223        assert_eq!(credentials.security_token().unwrap(), token);
224        assert!(credentials.signature().is_none());
225    }
226
227    #[test]
228    fn session_credentials_parse_properties() {
229        let content = "AccessKey=ak123\nSecretKey=sk456\n#comment\n \nSecurityToken=tk789\nSignature=sig123";
230        let credentials = SessionCredentials::default();
231        let properties = credentials.parse_properties(content).unwrap();
232        assert_eq!(properties.get("AccessKey").unwrap(), "ak123");
233        assert_eq!(properties.get("SecretKey").unwrap(), "sk456");
234        assert_eq!(properties.get("SecurityToken").unwrap(), "tk789");
235        assert_eq!(properties.get("Signature").unwrap(), "sig123");
236    }
237
238    #[test]
239    fn session_credentials_update_content() {
240        let mut properties = HashMap::new();
241        properties.insert(ACCESS_KEY.to_string(), "ak".to_string());
242        properties.insert(SECRET_KEY.to_string(), "sk".to_string());
243        properties.insert(SECURITY_TOKEN.to_string(), "token".to_string());
244
245        let mut credentials = SessionCredentials::default();
246        credentials.update_content(&properties);
247
248        assert_eq!(credentials.access_key().unwrap(), "ak");
249        assert_eq!(credentials.secret_key().unwrap(), "sk");
250        assert_eq!(credentials.security_token().unwrap(), "token");
251    }
252
253    #[test]
254    fn session_credentials_getters_and_setters() {
255        let mut credentials = SessionCredentials::default();
256        credentials.set_access_key("ak");
257        credentials.set_secret_key("sk");
258        credentials.set_security_token("token");
259        credentials.set_signature("sig");
260
261        assert_eq!(credentials.access_key().unwrap(), "ak");
262        assert_eq!(credentials.secret_key().unwrap(), "sk");
263        assert_eq!(credentials.security_token().unwrap(), "token");
264        assert_eq!(credentials.signature().unwrap(), "sig");
265    }
266
267    #[test]
268    fn session_credentials_eq_and_hash() {
269        let c1 = SessionCredentials::with_keys("ak", "sk");
270        let c2 = SessionCredentials::with_keys("ak", "sk");
271        let c3 = SessionCredentials::with_keys("ak2", "sk");
272
273        assert_eq!(c1, c2);
274        assert_ne!(c1, c3);
275
276        use std::collections::hash_map::DefaultHasher;
277        use std::hash::Hash;
278        use std::hash::Hasher;
279        let mut h1 = DefaultHasher::new();
280        c1.hash(&mut h1);
281        let mut h2 = DefaultHasher::new();
282        c2.hash(&mut h2);
283        assert_eq!(h1.finish(), h2.finish());
284    }
285
286    #[test]
287    fn session_credentials_display() {
288        let mut credentials = SessionCredentials::with_token("ak", "sk", "tk");
289        credentials.set_signature("sig");
290
291        let display = format!("{}", credentials);
292        let expected = "SessionCredentials [accessKey=<redacted>, secretKey=<redacted>, signature=<redacted>, \
293                        securityToken=<redacted>]";
294
295        assert_eq!(display, expected);
296        for secret in ["ak", "sk", "tk", "sig"] {
297            assert!(!display.contains(&format!("\"{secret}\"")));
298        }
299    }
300
301    #[test]
302    fn session_credentials_debug_redacts_sensitive_fields() {
303        let mut credentials = SessionCredentials::with_token("ak", "sk", "tk");
304        credentials.set_signature("sig");
305
306        let debug = format!("{credentials:?}");
307
308        assert!(debug.contains("access_key: \"<redacted>\""));
309        assert!(debug.contains("secret_key: \"<redacted>\""));
310        assert!(debug.contains("signature: \"<redacted>\""));
311        assert!(debug.contains("security_token: \"<redacted>\""));
312        for secret in ["ak", "sk", "tk", "sig"] {
313            assert!(!debug.contains(&format!("\"{secret}\"")));
314        }
315    }
316}
rocketmq_client_rust/common/session_credentials.rs

rocketmq_client_rust/common/
session_credentials.rs
