rocketmq_client_rust/common/
acl_client_rpc_hook.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::fmt;
16use std::net::SocketAddr;
17
18use crate::ClientError;
19use cheetah_string::CheetahString;
20use rocketmq_auth::cal_signature_segments_with_algorithm;
21use rocketmq_auth::SignatureAlgorithm;
22use rocketmq_error::REDACTED;
23use rocketmq_protocol::protocol::remoting_command::RemotingCommand;
24use rocketmq_transport::api::RPCHook;
25use smallvec::SmallVec;
26
27use crate::common::session_credentials::SessionCredentials;
28use crate::common::session_credentials::ACCESS_KEY;
29use crate::common::session_credentials::SECURITY_TOKEN;
30use crate::common::session_credentials::SIGNATURE;
31
32/// Client-side ACL RPC hook compatible with the Java RocketMQ client.
33///
34/// The hook signs the request exactly like Java's `AclClientRPCHook`: merge the
35/// custom header into `extFields`, sort all extension fields by key, concatenate
36/// every value except `Signature`, append the request body, and write the HMAC
37/// signature back to `extFields`.
38#[derive(Clone)]
39pub struct AclClientRPCHook {
40    session_credentials: SessionCredentials,
41    signature_algorithm: SignatureAlgorithm,
42}
43
44impl fmt::Debug for AclClientRPCHook {
45    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
46        formatter
47            .debug_struct("AclClientRPCHook")
48            .field("session_credentials", &REDACTED)
49            .field("signature_algorithm", &self.signature_algorithm)
50            .finish()
51    }
52}
53
54impl AclClientRPCHook {
55    pub fn new(session_credentials: SessionCredentials) -> Self {
56        Self::with_signature_algorithm(session_credentials, SignatureAlgorithm::default())
57    }
58
59    pub fn with_signature_algorithm(
60        session_credentials: SessionCredentials,
61        signature_algorithm: SignatureAlgorithm,
62    ) -> Self {
63        Self {
64            session_credentials,
65            signature_algorithm,
66        }
67    }
68
69    pub fn session_credentials(&self) -> &SessionCredentials {
70        &self.session_credentials
71    }
72
73    fn sorted_ext_fields_for_signing(request: &RemotingCommand) -> SmallVec<[(&CheetahString, &CheetahString); 16]> {
74        let mut sorted_fields = SmallVec::<[(&CheetahString, &CheetahString); 16]>::new();
75        if let Some(ext_fields) = request.ext_fields() {
76            for (key, value) in ext_fields {
77                sorted_fields.push((key, value));
78            }
79        }
80        sorted_fields.sort_unstable_by(|(left, _), (right, _)| left.as_str().cmp(right.as_str()));
81        sorted_fields
82    }
83}
84
85impl RPCHook for AclClientRPCHook {
86    fn do_before_request(
87        &self,
88        _remote_addr: SocketAddr,
89        request: &mut RemotingCommand,
90    ) -> Result<(), rocketmq_error::SharedError> {
91        let access_key = self
92            .session_credentials
93            .access_key()
94            .filter(|value| !value.is_empty())
95            .ok_or_else(|| ClientError::illegal_argument("ACL AccessKey must not be blank").into_shared_error())?;
96        let secret_key = self
97            .session_credentials
98            .secret_key()
99            .filter(|value| !value.is_empty())
100            .ok_or_else(|| ClientError::illegal_argument("ACL SecretKey must not be blank").into_shared_error())?;
101
102        request.ensure_ext_fields_initialized();
103        request.add_ext_field(ACCESS_KEY, access_key.clone());
104        if let Some(security_token) = self.session_credentials.security_token() {
105            request.add_ext_field(SECURITY_TOKEN, security_token.clone());
106        }
107
108        request.make_custom_header_to_net();
109        let signature = {
110            let sorted_fields = Self::sorted_ext_fields_for_signing(request);
111            let field_segments = sorted_fields
112                .iter()
113                .filter(|(key, _)| key.as_str() != SIGNATURE)
114                .map(|(_, value)| value.as_bytes());
115            let body_segments = request.body().into_iter().map(|body| body.as_ref());
116
117            cal_signature_segments_with_algorithm(
118                field_segments.chain(body_segments),
119                secret_key.as_str(),
120                self.signature_algorithm,
121            )
122            .map_err(|source| ClientError::illegal_argument_source(source).into_shared_error())?
123        };
124        request.add_ext_field(SIGNATURE, signature);
125
126        Ok(())
127    }
128
129    fn do_after_response(
130        &self,
131        _remote_addr: SocketAddr,
132        _request: &RemotingCommand,
133        _response: &mut RemotingCommand,
134    ) -> Result<(), rocketmq_error::SharedError> {
135        Ok(())
136    }
137}
138
139#[cfg(test)]
140mod tests {
141    use bytes::Bytes;
142    use cheetah_string::CheetahString;
143    use rocketmq_protocol::protocol::remoting_command::RemotingCommand;
144
145    use super::*;
146
147    #[test]
148    fn acl_hook_signs_sorted_ext_fields_and_body_like_java_client() {
149        let hook = AclClientRPCHook::new(SessionCredentials::with_token("ak", "secret", "token"));
150        let mut request = RemotingCommand::create_remoting_command(10).set_body(Bytes::from_static(b"body"));
151        request.ensure_ext_fields_initialized();
152        request.add_ext_field("z", "last");
153        request.add_ext_field("a", "first");
154
155        hook.do_before_request("127.0.0.1:9876".parse().unwrap(), &mut request)
156            .unwrap();
157
158        let ext_fields = request.ext_fields().unwrap();
159        assert_eq!(
160            ext_fields.get(&CheetahString::from_static_str(ACCESS_KEY)).unwrap(),
161            "ak"
162        );
163        assert_eq!(
164            ext_fields.get(&CheetahString::from_static_str(SECURITY_TOKEN)).unwrap(),
165            "token"
166        );
167        assert_eq!(
168            ext_fields.get(&CheetahString::from_static_str(SIGNATURE)).unwrap(),
169            "wuJkFDNRMKY6aFXTDwT/vg0UC8M="
170        );
171    }
172
173    #[test]
174    fn acl_hook_excludes_existing_signature_when_resigning_like_java_client() {
175        let hook = AclClientRPCHook::new(SessionCredentials::with_token("ak", "secret", "token"));
176        let mut request = RemotingCommand::create_remoting_command(10).set_body(Bytes::from_static(b"body"));
177        request.ensure_ext_fields_initialized();
178        request.add_ext_field("z", "last");
179        request.add_ext_field("a", "first");
180        request.add_ext_field(SIGNATURE, "stale-signature");
181
182        hook.do_before_request("127.0.0.1:9876".parse().unwrap(), &mut request)
183            .unwrap();
184
185        let ext_fields = request.ext_fields().unwrap();
186        assert_eq!(
187            ext_fields.get(&CheetahString::from_static_str(SIGNATURE)).unwrap(),
188            "wuJkFDNRMKY6aFXTDwT/vg0UC8M="
189        );
190    }
191
192    #[test]
193    fn acl_hook_supports_hmac_sha256_for_secure_profiles() {
194        let hook = AclClientRPCHook::with_signature_algorithm(
195            SessionCredentials::with_token("ak", "secret", "token"),
196            SignatureAlgorithm::HmacSha256,
197        );
198        let mut request = RemotingCommand::create_remoting_command(10).set_body(Bytes::from_static(b"body"));
199        request.ensure_ext_fields_initialized();
200        request.add_ext_field("z", "last");
201        request.add_ext_field("a", "first");
202
203        hook.do_before_request("127.0.0.1:9876".parse().unwrap(), &mut request)
204            .unwrap();
205
206        assert_eq!(
207            request
208                .ext_fields()
209                .unwrap()
210                .get(&CheetahString::from_static_str(SIGNATURE))
211                .unwrap(),
212            "mBwsFNM/fUVM2NmLpnXIS+x+x7EZXmyuboG/4ZGmwa0="
213        );
214    }
215
216    #[test]
217    fn acl_hook_rejects_blank_credentials_before_signing() {
218        let hook = AclClientRPCHook::new(SessionCredentials::default());
219        let mut request = RemotingCommand::create_remoting_command(10);
220
221        assert!(hook
222            .do_before_request("127.0.0.1:9876".parse().unwrap(), &mut request)
223            .is_err());
224    }
225
226    #[test]
227    fn acl_hook_debug_redacts_session_credentials() {
228        let hook = AclClientRPCHook::new(SessionCredentials::with_token("ak", "secret", "token"));
229
230        let debug = format!("{hook:?}");
231
232        assert!(debug.contains(REDACTED));
233        assert!(!debug.contains("ak"));
234        assert!(!debug.contains("secret"));
235        assert!(!debug.contains("token"));
236    }
237}
rocketmq_client_rust/common/acl_client_rpc_hook.rs

rocketmq_client_rust/common/
acl_client_rpc_hook.rs
