1use std::future::Future;
22use std::net::IpAddr;
23use std::net::SocketAddr;
24
25use cheetah_string::CheetahString;
26use rand::seq::IndexedRandom;
27use rocketmq_model::common::config::TopicConfig;
28use rocketmq_model::common::message::message_ext::MessageExt;
29use rocketmq_model::common::mix_all;
30use rocketmq_protocol::common::message::message_decoder as MessageDecoder;
31use rocketmq_protocol::protocol::admin::consume_stats::ConsumeStats;
32use rocketmq_protocol::protocol::body::broker_body::cluster_info::ClusterInfo;
33use rocketmq_protocol::protocol::body::consumer_connection::ConsumerConnection;
34use rocketmq_protocol::protocol::body::group_list::GroupList;
35use rocketmq_protocol::protocol::body::kv_table::KVTable;
36use rocketmq_protocol::protocol::body::producer_connection::ProducerConnection;
37use rocketmq_protocol::protocol::body::producer_table_info::ProducerTableInfo;
38use rocketmq_protocol::protocol::body::proxy_drain::ProxyDrainStateResponseBody;
39use rocketmq_protocol::protocol::body::topic::topic_list::TopicList;
40use rocketmq_protocol::protocol::header::get_consume_stats_request_header::GetConsumeStatsRequestHeader;
41use rocketmq_protocol::protocol::header::query_topic_consume_by_who_request_header::QueryTopicConsumeByWhoRequestHeader;
42use rocketmq_protocol::protocol::header::view_message_request_header::ViewMessageRequestHeader;
43use rocketmq_protocol::protocol::route::route_data_view::BrokerData;
44use rocketmq_protocol::protocol::route::topic_route_data::TopicRouteData;
45use rocketmq_protocol::protocol::route_facade::BrokerDataExt;
46use rocketmq_protocol::protocol::subscription::subscription_group_config::SubscriptionGroupConfig;
47
48use super::default_mq_admin_ext::DefaultMQAdminExt;
49
50#[derive(Clone, Debug, Default, Eq, PartialEq)]
53pub struct BrokerConfigAllowlisted {
54 pub generation: u64,
55 pub send_message_thread_pool_nums: Option<u32>,
56 pub pull_message_thread_pool_nums: Option<u32>,
57 pub flush_delay_offset_interval_ms: Option<u64>,
58 pub max_client_event_count: Option<i32>,
59}
60
61#[derive(Clone, Debug, PartialEq)]
63pub struct TopicConfigVersioned {
64 pub version: u64,
65 pub config: TopicConfig,
66}
67
68#[derive(Clone, Debug)]
71pub struct SubscriptionGroupConfigVersioned {
72 pub version: u64,
73 pub config: SubscriptionGroupConfig,
74}
75
76#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd)]
78pub enum ReadFailureCode {
79 SourceUnavailable,
80 Timeout,
81 PermissionDenied,
82 NotFound,
83 RateLimited,
84 InvalidResponse,
85}
86
87#[derive(Clone, Debug, Eq, PartialEq, Ord, PartialOrd)]
89pub struct BrokerReadFailure {
90 broker_name: String,
91 code: ReadFailureCode,
92 retryable: bool,
93}
94
95impl BrokerReadFailure {
96 pub fn new(broker_name: impl AsRef<str>, code: ReadFailureCode, retryable: bool) -> Self {
99 Self {
100 broker_name: sanitize_broker_logical_target(broker_name.as_ref()),
101 code,
102 retryable,
103 }
104 }
105
106 pub fn broker_name(&self) -> &str {
107 &self.broker_name
108 }
109
110 pub const fn code(&self) -> ReadFailureCode {
111 self.code
112 }
113
114 pub const fn retryable(&self) -> bool {
115 self.retryable
116 }
117}
118
119#[derive(Debug)]
121pub struct ConsumeStatsReadResult {
122 pub stats: ConsumeStats,
123 pub attempted_brokers: usize,
124 pub successful_brokers: usize,
125 pub failures: Vec<BrokerReadFailure>,
126}
127
128#[derive(Clone, Debug, Eq, PartialEq)]
134pub struct MessageMetadataRead {
135 pub topic: String,
136 pub message_id: String,
137 pub unique_message_id: Option<String>,
138 pub born_timestamp: i64,
139 pub store_timestamp: i64,
140 pub queue_id: i32,
141 pub queue_offset: i64,
142 pub store_size: i32,
143 pub reconsume_times: i32,
144 pub sys_flag: i32,
145 pub flag: i32,
146 pub prepared_transaction_offset: i64,
147}
148
149#[allow(async_fn_in_trait)]
150pub trait MQAdminReadExt: Send {
151 async fn start(&mut self) -> crate::ClientResult<()>;
152
153 async fn shutdown(&mut self);
154
155 async fn fetch_all_topic_list(&self) -> crate::ClientResult<TopicList>;
156
157 async fn fetch_broker_runtime_stats(&self, broker_addr: CheetahString) -> crate::ClientResult<KVTable>;
158
159 async fn get_broker_config_allowlisted(
162 &self,
163 broker_addr: CheetahString,
164 ) -> crate::ClientResult<BrokerConfigAllowlisted>;
165
166 async fn proxy_drain_state(&self, proxy_addr: CheetahString) -> crate::ClientResult<ProxyDrainStateResponseBody>;
168
169 async fn examine_consume_stats(
170 &self,
171 consumer_group: CheetahString,
172 topic: Option<CheetahString>,
173 cluster_name: Option<CheetahString>,
174 broker_addr: Option<CheetahString>,
175 timeout_millis: Option<u64>,
176 ) -> crate::ClientResult<ConsumeStats>;
177
178 async fn examine_consume_stats_with_evidence(
182 &self,
183 consumer_group: CheetahString,
184 topic: Option<CheetahString>,
185 broker_addr: Option<CheetahString>,
186 timeout_millis: Option<u64>,
187 ) -> crate::ClientResult<ConsumeStatsReadResult> {
188 let stats = self
189 .examine_consume_stats(consumer_group, topic, None, broker_addr, timeout_millis)
190 .await?;
191 Ok(ConsumeStatsReadResult {
192 stats,
193 attempted_brokers: 1,
194 successful_brokers: 1,
195 failures: Vec::new(),
196 })
197 }
198
199 async fn examine_broker_cluster_info(&self) -> crate::ClientResult<ClusterInfo>;
200
201 async fn examine_topic_route_info(&self, topic: CheetahString) -> crate::ClientResult<Option<TopicRouteData>>;
202
203 async fn topic_config_with_version(
205 &self,
206 broker_addr: CheetahString,
207 topic: CheetahString,
208 ) -> crate::ClientResult<TopicConfigVersioned>;
209
210 async fn subscription_group_config_with_version(
213 &self,
214 broker_addr: CheetahString,
215 group: CheetahString,
216 ) -> crate::ClientResult<SubscriptionGroupConfigVersioned>;
217
218 async fn examine_consumer_connection_info(
219 &self,
220 consumer_group: CheetahString,
221 broker_addr: Option<CheetahString>,
222 ) -> crate::ClientResult<ConsumerConnection>;
223
224 async fn observe_consumer_connection_at(
227 &self,
228 consumer_group: CheetahString,
229 broker_addr: CheetahString,
230 ) -> crate::ClientResult<ConsumerConnection>;
231
232 async fn examine_producer_connection_info(
233 &self,
234 producer_group: CheetahString,
235 topic: CheetahString,
236 ) -> crate::ClientResult<ProducerConnection>;
237
238 async fn observe_producer_connection_info(
241 &self,
242 producer_group: CheetahString,
243 topic: CheetahString,
244 ) -> crate::ClientResult<ProducerConnection>;
245
246 async fn observe_producer_connection_at(
249 &self,
250 producer_group: CheetahString,
251 broker_addr: CheetahString,
252 ) -> crate::ClientResult<ProducerConnection>;
253
254 async fn get_all_producer_info(&self, broker_addr: CheetahString) -> crate::ClientResult<ProducerTableInfo>;
255
256 async fn query_topic_consume_by_who(&self, topic: CheetahString) -> crate::ClientResult<GroupList>;
257}
258
259#[allow(async_fn_in_trait)]
264pub trait MQAdminTopicInventoryReadExt: Send {
265 async fn fetch_topic_inventory(&self, cluster: Option<CheetahString>) -> crate::ClientResult<TopicList>;
268}
269
270#[allow(async_fn_in_trait)]
272pub trait MQAdminMessageReadExt: Send {
273 async fn query_message_metadata(
276 &self,
277 topic: CheetahString,
278 message_id: CheetahString,
279 ) -> crate::ClientResult<MessageMetadataRead>;
280
281 async fn query_message_metadata_by_id(
284 &self,
285 _message_id: CheetahString,
286 ) -> crate::ClientResult<MessageMetadataRead> {
287 Err(crate::ClientError::response_process_failed(
288 "query_message_metadata_by_id",
289 "Topic-free message metadata lookup is not implemented by this adapter",
290 ))
291 }
292}
293
294impl MQAdminReadExt for DefaultMQAdminExt {
295 async fn start(&mut self) -> crate::ClientResult<()> {
296 self.inner_mut().start_admin().await
297 }
298
299 async fn shutdown(&mut self) {
300 self.inner_mut().shutdown_admin().await;
301 }
302
303 async fn fetch_all_topic_list(&self) -> crate::ClientResult<TopicList> {
304 self.inner()
305 .mq_client_api()?
306 .get_all_topic_list_from_name_server(self.inner().remoting_timeout_millis()?)
307 .await
308 }
309
310 async fn fetch_broker_runtime_stats(&self, broker_addr: CheetahString) -> crate::ClientResult<KVTable> {
311 self.inner()
312 .mq_client_api()?
313 .get_broker_runtime_info(&broker_addr, self.inner().remoting_timeout_millis()?)
314 .await
315 }
316
317 async fn get_broker_config_allowlisted(
318 &self,
319 broker_addr: CheetahString,
320 ) -> crate::ClientResult<BrokerConfigAllowlisted> {
321 let snapshot = self
322 .inner()
323 .mq_client_api()?
324 .get_broker_config_snapshot(&broker_addr, self.inner().remoting_timeout_millis()?)
325 .await?;
326 let generation = snapshot.generation.filter(|value| *value > 0).ok_or_else(|| {
327 crate::ClientError::response_process_failed(
328 "get_broker_config_allowlisted",
329 "Broker config response does not include a positive config generation",
330 )
331 })?;
332 let properties = snapshot.properties;
333 Ok(BrokerConfigAllowlisted {
334 generation,
335 send_message_thread_pool_nums: parse_allowlisted_value(&properties, "sendMessageThreadPoolNums")?,
336 pull_message_thread_pool_nums: parse_allowlisted_value(&properties, "pullMessageThreadPoolNums")?,
337 flush_delay_offset_interval_ms: parse_allowlisted_value(&properties, "flushDelayOffsetInterval")?,
338 max_client_event_count: parse_allowlisted_value(&properties, "maxClientEventCount")?,
339 })
340 }
341
342 async fn proxy_drain_state(&self, proxy_addr: CheetahString) -> crate::ClientResult<ProxyDrainStateResponseBody> {
343 self.inner()
344 .mq_client_api()?
345 .get_proxy_drain_state(&proxy_addr, self.inner().remoting_timeout_millis()?)
346 .await
347 }
348
349 async fn examine_consume_stats(
350 &self,
351 consumer_group: CheetahString,
352 topic: Option<CheetahString>,
353 _cluster_name: Option<CheetahString>,
354 broker_addr: Option<CheetahString>,
355 timeout_millis: Option<u64>,
356 ) -> crate::ClientResult<ConsumeStats> {
357 let timeout = timeout_millis.unwrap_or(self.inner().remoting_timeout_millis()?);
358 let topic = topic.unwrap_or_default();
359 if let Some(broker_addr) = broker_addr {
360 return self
361 .inner()
362 .mq_client_api()?
363 .get_consume_stats(
364 &broker_addr,
365 GetConsumeStatsRequestHeader {
366 consumer_group,
367 topic,
368 topic_list: None,
369 topic_request_header: None,
370 },
371 timeout,
372 )
373 .await;
374 }
375
376 let retry_topic: CheetahString = mix_all::get_retry_topic(&consumer_group).into();
377 let route = self
378 .inner()
379 .mq_client_api()?
380 .get_topic_route_info_from_name_server(&retry_topic, timeout)
381 .await?;
382 let mut result = ConsumeStats::new();
383 if let Some(route) = route {
384 for broker in &route.broker_datas {
385 if let Some(master_addr) = broker.broker_addrs().get(&mix_all::MASTER_ID) {
386 let stats = self
387 .inner()
388 .mq_client_api()?
389 .get_consume_stats(
390 master_addr,
391 GetConsumeStatsRequestHeader {
392 consumer_group: consumer_group.clone(),
393 topic: topic.clone(),
394 topic_list: None,
395 topic_request_header: None,
396 },
397 timeout,
398 )
399 .await?;
400 result.get_offset_table_mut().extend(stats.offset_table);
401 result.set_consume_tps(result.get_consume_tps() + stats.consume_tps);
402 }
403 }
404 }
405 Ok(result)
406 }
407
408 async fn examine_consume_stats_with_evidence(
409 &self,
410 consumer_group: CheetahString,
411 topic: Option<CheetahString>,
412 broker_addr: Option<CheetahString>,
413 timeout_millis: Option<u64>,
414 ) -> crate::ClientResult<ConsumeStatsReadResult> {
415 let timeout = timeout_millis.unwrap_or(self.inner().remoting_timeout_millis()?);
416 let topic = topic.unwrap_or_default();
417 let targets = match broker_addr {
418 Some(broker_addr) => vec![ConsumeStatsReadTarget::Ready {
419 broker_name: "selected-broker".to_string(),
420 address: broker_addr,
421 }],
422 None => {
423 let retry_topic: CheetahString = mix_all::get_retry_topic(&consumer_group).into();
424 let route = self
425 .inner()
426 .mq_client_api()?
427 .get_topic_route_info_from_name_server(&retry_topic, timeout)
428 .await?;
429 consume_stats_read_targets(route.map(|route| route.broker_datas).unwrap_or_default())
430 }
431 };
432 let mut result = ConsumeStatsReadResult {
433 stats: ConsumeStats::new(),
434 attempted_brokers: targets.len(),
435 successful_brokers: 0,
436 failures: Vec::new(),
437 };
438 for target in targets {
439 let (broker_name, address) = match target {
440 ConsumeStatsReadTarget::Ready { broker_name, address } => (broker_name, address),
441 ConsumeStatsReadTarget::Failed(failure) => {
442 result.failures.push(failure);
443 continue;
444 }
445 };
446 match self
447 .inner()
448 .mq_client_api()?
449 .get_consume_stats(
450 &address,
451 GetConsumeStatsRequestHeader {
452 consumer_group: consumer_group.clone(),
453 topic: topic.clone(),
454 topic_list: None,
455 topic_request_header: None,
456 },
457 timeout,
458 )
459 .await
460 {
461 Ok(stats) => {
462 result.successful_brokers += 1;
463 result.stats.get_offset_table_mut().extend(stats.offset_table);
464 result
465 .stats
466 .set_consume_tps(result.stats.get_consume_tps() + stats.consume_tps);
467 }
468 Err(error) => result.failures.push(broker_read_failure(broker_name, &error)),
469 }
470 }
471 result.failures.sort();
472 Ok(result)
473 }
474
475 async fn examine_broker_cluster_info(&self) -> crate::ClientResult<ClusterInfo> {
476 self.inner()
477 .mq_client_api()?
478 .get_broker_cluster_info(self.inner().remoting_timeout_millis()?)
479 .await
480 }
481
482 async fn examine_topic_route_info(&self, topic: CheetahString) -> crate::ClientResult<Option<TopicRouteData>> {
483 self.inner()
484 .mq_client_api()?
485 .get_topic_route_info_from_name_server(&topic, self.inner().remoting_timeout_millis()?)
486 .await
487 }
488
489 async fn topic_config_with_version(
490 &self,
491 broker_addr: CheetahString,
492 topic: CheetahString,
493 ) -> crate::ClientResult<TopicConfigVersioned> {
494 self.inner()
495 .mq_client_api()?
496 .get_topic_config_with_version(&broker_addr, topic, self.inner().remoting_timeout_millis()?)
497 .await
498 }
499
500 async fn subscription_group_config_with_version(
501 &self,
502 broker_addr: CheetahString,
503 group: CheetahString,
504 ) -> crate::ClientResult<SubscriptionGroupConfigVersioned> {
505 self.inner()
506 .mq_client_api()?
507 .get_subscription_group_config_with_version(&broker_addr, group, self.inner().remoting_timeout_millis()?)
508 .await
509 }
510
511 async fn examine_consumer_connection_info(
512 &self,
513 consumer_group: CheetahString,
514 broker_addr: Option<CheetahString>,
515 ) -> crate::ClientResult<ConsumerConnection> {
516 let timeout = self.inner().remoting_timeout_millis()?;
517 let selected_addr = match broker_addr {
518 Some(broker_addr) => Some(broker_addr),
519 None => {
520 let retry_topic = CheetahString::from_string(mix_all::get_retry_topic(consumer_group.as_str()));
521 self.inner()
522 .mq_client_api()?
523 .get_topic_route_info_from_name_server(&retry_topic, timeout)
524 .await?
525 .and_then(|route| {
526 route
527 .broker_datas
528 .choose(&mut rand::rng())
529 .and_then(BrokerDataExt::select_broker_addr)
530 })
531 }
532 };
533 let mut result = ConsumerConnection::new();
534 if let Some(broker_addr) = selected_addr {
535 result = self
536 .inner()
537 .mq_client_api()?
538 .get_consumer_connection_list(broker_addr.as_str(), consumer_group, timeout)
539 .await?;
540 }
541 if result.get_connection_set().is_empty() {
542 return Err(crate::mq_client_err!(
543 rocketmq_protocol::code::response_code::ResponseCode::ConsumerNotOnline,
544 "Not found the consumer group connection"
545 ));
546 }
547 Ok(result)
548 }
549
550 async fn observe_consumer_connection_at(
551 &self,
552 consumer_group: CheetahString,
553 broker_addr: CheetahString,
554 ) -> crate::ClientResult<ConsumerConnection> {
555 self.inner()
556 .mq_client_api()?
557 .get_consumer_connection_list(
558 broker_addr.as_str(),
559 consumer_group,
560 self.inner().remoting_timeout_millis()?,
561 )
562 .await
563 }
564
565 async fn examine_producer_connection_info(
566 &self,
567 producer_group: CheetahString,
568 topic: CheetahString,
569 ) -> crate::ClientResult<ProducerConnection> {
570 let timeout = self.inner().remoting_timeout_millis()?;
571 let route = self
572 .inner()
573 .mq_client_api()?
574 .get_topic_route_info_from_name_server(&topic, timeout)
575 .await?;
576 let mut result = ProducerConnection::new();
577 if let Some(broker_addr) = route.and_then(|route| {
578 route
579 .broker_datas
580 .choose(&mut rand::rng())
581 .and_then(BrokerDataExt::select_broker_addr)
582 }) {
583 result = self
584 .inner()
585 .mq_client_api()?
586 .get_producer_connection_list(broker_addr.as_str(), producer_group, timeout)
587 .await?;
588 }
589 if result.connection_set().is_empty() {
590 return Err(crate::mq_client_err!("Not found the producer group connection"));
591 }
592 Ok(result)
593 }
594
595 async fn observe_producer_connection_info(
596 &self,
597 producer_group: CheetahString,
598 topic: CheetahString,
599 ) -> crate::ClientResult<ProducerConnection> {
600 let timeout = self.inner().remoting_timeout_millis()?;
601 let route = self
602 .inner()
603 .mq_client_api()?
604 .get_topic_route_info_from_name_server(&topic, timeout)
605 .await?;
606 let Some(broker_addr) = route.and_then(|route| {
607 route
608 .broker_datas
609 .choose(&mut rand::rng())
610 .and_then(BrokerDataExt::select_broker_addr)
611 }) else {
612 return Ok(ProducerConnection::new());
613 };
614 self.inner()
615 .mq_client_api()?
616 .get_producer_connection_list(broker_addr.as_str(), producer_group, timeout)
617 .await
618 }
619
620 async fn observe_producer_connection_at(
621 &self,
622 producer_group: CheetahString,
623 broker_addr: CheetahString,
624 ) -> crate::ClientResult<ProducerConnection> {
625 self.inner()
626 .mq_client_api()?
627 .get_producer_connection_list(
628 broker_addr.as_str(),
629 producer_group,
630 self.inner().remoting_timeout_millis()?,
631 )
632 .await
633 }
634
635 async fn get_all_producer_info(&self, broker_addr: CheetahString) -> crate::ClientResult<ProducerTableInfo> {
636 self.inner()
637 .mq_client_api()?
638 .get_all_producer_info(broker_addr.as_str(), self.inner().remoting_timeout_millis()?)
639 .await
640 }
641
642 async fn query_topic_consume_by_who(&self, topic: CheetahString) -> crate::ClientResult<GroupList> {
643 let timeout = self.inner().remoting_timeout_millis()?;
644 let route = self
645 .inner()
646 .mq_client_api()?
647 .get_topic_route_info_from_name_server(&topic, timeout)
648 .await?;
649 if let Some(route) = route {
650 for broker in &route.broker_datas {
651 if let Some(master_addr) = broker.broker_addrs().get(&mix_all::MASTER_ID) {
652 return self
653 .inner()
654 .mq_client_api()?
655 .query_topic_consume_by_who(
656 master_addr,
657 QueryTopicConsumeByWhoRequestHeader {
658 topic: topic.clone(),
659 topic_request_header: None,
660 },
661 timeout,
662 )
663 .await;
664 }
665 }
666 }
667 Ok(GroupList::default())
668 }
669}
670
671impl MQAdminTopicInventoryReadExt for DefaultMQAdminExt {
672 async fn fetch_topic_inventory(&self, cluster: Option<CheetahString>) -> crate::ClientResult<TopicList> {
673 let client_api = self.inner().mq_client_api()?;
674 let timeout_millis = self.inner().remoting_timeout_millis()?;
675 fetch_topic_inventory_from(
676 cluster,
677 |cluster| client_api.get_topics_by_cluster(cluster, timeout_millis),
678 || client_api.get_all_topic_list_from_name_server(timeout_millis),
679 )
680 .await
681 }
682}
683
684impl MQAdminMessageReadExt for DefaultMQAdminExt {
685 async fn query_message_metadata(
686 &self,
687 topic: CheetahString,
688 message_id: CheetahString,
689 ) -> crate::ClientResult<MessageMetadataRead> {
690 let (broker_addr, header) = message_lookup_target(message_id.as_str(), Some(topic.clone()))?;
691 let message = self
692 .inner()
693 .mq_client_api()?
694 .view_message(&broker_addr, header, self.inner().remoting_timeout_millis()?)
695 .await?;
696 Ok(MessageMetadataRead {
697 topic: topic.to_string(),
698 message_id: message.msg_id().to_string(),
699 unique_message_id: message
700 .properties()
701 .get(rocketmq_model::common::message::MessageConst::PROPERTY_UNIQ_CLIENT_MESSAGE_ID_KEYIDX)
702 .map(ToString::to_string)
703 .filter(|value| !value.is_empty()),
704 born_timestamp: message.born_timestamp(),
705 store_timestamp: message.store_timestamp(),
706 queue_id: message.queue_id(),
707 queue_offset: message.queue_offset(),
708 store_size: message.store_size(),
709 reconsume_times: message.reconsume_times(),
710 sys_flag: message.sys_flag(),
711 flag: message.flag(),
712 prepared_transaction_offset: message.prepared_transaction_offset(),
713 })
714 }
715
716 async fn query_message_metadata_by_id(
717 &self,
718 message_id: CheetahString,
719 ) -> crate::ClientResult<MessageMetadataRead> {
720 let (broker_addr, header) = message_lookup_target(message_id.as_str(), None)?;
721 let message = self
722 .inner()
723 .mq_client_api()?
724 .view_message(&broker_addr, header, self.inner().remoting_timeout_millis()?)
725 .await?;
726 let topic = message.topic().to_string();
727 Ok(project_message_metadata(message, topic))
728 }
729}
730
731fn message_lookup_target(
732 message_id: &str,
733 topic: Option<CheetahString>,
734) -> crate::ClientResult<(CheetahString, ViewMessageRequestHeader)> {
735 MessageDecoder::validate_message_id(message_id)
736 .map_err(|error| crate::ClientError::illegal_argument(format!("Invalid message ID: {error}")))?;
737 let decoded = MessageDecoder::decode_message_id(message_id).map_err(crate::ClientError::illegal_argument)?;
738 Ok((
739 CheetahString::from_string(decoded.address.to_string()),
740 ViewMessageRequestHeader {
741 topic,
742 offset: decoded.offset,
743 },
744 ))
745}
746
747fn project_message_metadata(message: MessageExt, topic: String) -> MessageMetadataRead {
748 MessageMetadataRead {
749 topic,
750 message_id: message.msg_id().to_string(),
751 unique_message_id: None,
752 born_timestamp: message.born_timestamp(),
753 store_timestamp: message.store_timestamp(),
754 queue_id: message.queue_id(),
755 queue_offset: message.queue_offset(),
756 store_size: message.store_size(),
757 reconsume_times: message.reconsume_times(),
758 sys_flag: message.sys_flag(),
759 flag: message.flag(),
760 prepared_transaction_offset: message.prepared_transaction_offset(),
761 }
762}
763
764async fn fetch_topic_inventory_from<ClusterFetch, ClusterFuture, GlobalFetch, GlobalFuture>(
765 cluster: Option<CheetahString>,
766 fetch_cluster: ClusterFetch,
767 fetch_global: GlobalFetch,
768) -> crate::ClientResult<TopicList>
769where
770 ClusterFetch: FnOnce(CheetahString) -> ClusterFuture,
771 ClusterFuture: Future<Output = crate::ClientResult<TopicList>>,
772 GlobalFetch: FnOnce() -> GlobalFuture,
773 GlobalFuture: Future<Output = crate::ClientResult<TopicList>>,
774{
775 match cluster {
776 Some(cluster) => fetch_cluster(cluster).await,
777 None => fetch_global().await,
778 }
779}
780
781fn parse_allowlisted_value<T>(
782 properties: &std::collections::HashMap<CheetahString, CheetahString>,
783 key: &str,
784) -> crate::ClientResult<Option<T>>
785where
786 T: std::str::FromStr,
787{
788 properties
789 .iter()
790 .find_map(|(name, value)| (name.as_str() == key).then_some(value.as_str()))
791 .map(|value| {
792 value.parse().map_err(|_| {
793 crate::ClientError::illegal_argument(format!("Broker allowlisted configuration `{key}` is malformed"))
794 })
795 })
796 .transpose()
797}
798
799#[derive(Debug, Eq, PartialEq)]
800enum ConsumeStatsReadTarget {
801 Ready {
802 broker_name: String,
803 address: CheetahString,
804 },
805 Failed(BrokerReadFailure),
806}
807
808fn consume_stats_read_targets(brokers: Vec<BrokerData>) -> Vec<ConsumeStatsReadTarget> {
809 brokers
810 .into_iter()
811 .map(|broker| {
812 let broker_name = broker.broker_name().to_string();
813 match broker.broker_addrs().get(&mix_all::MASTER_ID).cloned() {
814 Some(address) if !address.as_str().trim().is_empty() => {
815 ConsumeStatsReadTarget::Ready { broker_name, address }
816 }
817 Some(_) | None => ConsumeStatsReadTarget::Failed(BrokerReadFailure::new(
818 broker_name,
819 ReadFailureCode::InvalidResponse,
820 false,
821 )),
822 }
823 })
824 .collect()
825}
826
827fn broker_read_failure(broker_name: String, error: &crate::ClientError) -> BrokerReadFailure {
828 let descriptor = error.descriptor();
829 let status = descriptor.projection().http().status.as_u16();
830 let code = match status {
831 401 | 403 => ReadFailureCode::PermissionDenied,
832 404 => ReadFailureCode::NotFound,
833 408 | 504 => ReadFailureCode::Timeout,
834 429 => ReadFailureCode::RateLimited,
835 400 | 413 | 422 => ReadFailureCode::InvalidResponse,
836 _ => ReadFailureCode::SourceUnavailable,
837 };
838 let retryable = matches!(
839 descriptor.recovery_hint(),
840 rocketmq_error::RecoveryHint::Backoff
841 | rocketmq_error::RecoveryHint::RefreshRoute
842 | rocketmq_error::RecoveryHint::RefreshLeader
843 | rocketmq_error::RecoveryHint::SwitchBroker
844 );
845 BrokerReadFailure::new(broker_name, code, retryable)
846}
847
848fn sanitize_broker_logical_target(target: &str) -> String {
849 const UNKNOWN_TARGET: &str = "unknown";
850 const MAX_TARGET_BYTES: usize = 128;
851 const SENSITIVE_OR_ERROR_MARKERS: &[&str] = &[
852 "access_key",
853 "accesskey",
854 "authorization",
855 "bearer",
856 "credential",
857 "denied",
858 "error",
859 "exception",
860 "failed",
861 "failure",
862 "password",
863 "passwd",
864 "refused",
865 "secret",
866 "signature",
867 "source_unavailable",
868 "timed_out",
869 "timeout",
870 "token",
871 ];
872
873 let raw_target = target;
874 let target = raw_target.trim();
875 let lowercase = target.to_ascii_lowercase();
876 let invalid = target.is_empty()
877 || raw_target.len() > MAX_TARGET_BYTES
878 || !target.is_ascii()
879 || target.parse::<IpAddr>().is_ok()
880 || target.parse::<SocketAddr>().is_ok()
881 || target.contains([':', '/', '\\', '@', '=', '&', '?', '#'])
882 || raw_target.chars().any(char::is_control)
883 || !target
884 .chars()
885 .all(|character| character.is_ascii_alphanumeric() || matches!(character, '-' | '_' | '.'))
886 || lowercase.starts_with("akia")
887 || lowercase.starts_with("sk-")
888 || SENSITIVE_OR_ERROR_MARKERS
889 .iter()
890 .any(|marker| lowercase.contains(marker));
891 if invalid {
892 UNKNOWN_TARGET.to_string()
893 } else {
894 target.to_string()
895 }
896}
897
898#[cfg(test)]
899mod tests {
900 use crate::test_support::error_assertions::assert_invalid_argument;
901 use std::collections::HashMap;
902 use std::sync::atomic::AtomicUsize;
903 use std::sync::atomic::Ordering;
904 use std::sync::Arc;
905
906 use bytes::Bytes;
907 use cheetah_string::CheetahString;
908 use rocketmq_model::common::message::message_ext::MessageExt;
909 use rocketmq_model::common::message::MessageConst;
910 use rocketmq_model::common::message::MessageTrait;
911 use rocketmq_model::common::mix_all;
912 use rocketmq_protocol::protocol::body::topic::topic_list::TopicList;
913 use rocketmq_protocol::protocol::route::route_data_view::BrokerData;
914
915 use super::consume_stats_read_targets;
916 use super::fetch_topic_inventory_from;
917 use super::message_lookup_target;
918 use super::parse_allowlisted_value;
919 use super::project_message_metadata;
920 use super::BrokerReadFailure;
921 use super::ConsumeStatsReadTarget;
922 use super::ReadFailureCode;
923
924 fn route_broker(name: &str, master: Option<&str>) -> BrokerData {
925 let broker_addrs = master
926 .map(|address| HashMap::from([(mix_all::MASTER_ID, CheetahString::from_string(address.to_string()))]))
927 .unwrap_or_default();
928 BrokerData::new(
929 CheetahString::from_static_str("cluster-a"),
930 CheetahString::from_string(name.to_string()),
931 broker_addrs,
932 None,
933 )
934 }
935
936 #[test]
937 fn topic_free_message_lookup_decodes_exact_endpoint_and_offset() {
938 let (address, header) = message_lookup_target("7F00000100002A9F000000000000002A", None).unwrap();
939
940 assert_eq!(address, "127.0.0.1:10911");
941 assert_eq!(header.offset, 42);
942 assert!(header.topic.is_none());
943 assert!(message_lookup_target("not-a-message-id", None).is_err());
944 }
945
946 #[test]
947 fn topic_free_metadata_projection_ignores_body_properties_and_hosts() {
948 let mut message = MessageExt::default();
949 message.set_topic("actual-topic".into());
950 message.set_msg_id("raw-offset-id".into());
951 message.set_body(Bytes::from_static(b"must-not-be-read"));
952 message.put_property(
953 CheetahString::from_static_str(MessageConst::PROPERTY_UNIQ_CLIENT_MESSAGE_ID_KEYIDX),
954 "must-not-be-read".into(),
955 );
956 message.set_born_host("192.0.2.10:1234".parse().unwrap());
957 message.set_store_host("192.0.2.11:10911".parse().unwrap());
958
959 let metadata = project_message_metadata(message, "actual-topic".to_owned());
960
961 assert_eq!(metadata.topic, "actual-topic");
962 assert_eq!(metadata.message_id, "raw-offset-id");
963 assert_eq!(metadata.unique_message_id, None);
964 }
965
966 #[tokio::test]
967 async fn cluster_topic_inventory_selects_exactly_one_cluster_source_call() {
968 let cluster_calls = Arc::new(AtomicUsize::new(0));
969 let global_calls = Arc::new(AtomicUsize::new(0));
970 let result = fetch_topic_inventory_from(
971 Some(CheetahString::from("cluster-a")),
972 {
973 let cluster_calls = Arc::clone(&cluster_calls);
974 move |cluster| async move {
975 cluster_calls.fetch_add(1, Ordering::SeqCst);
976 assert_eq!(cluster, CheetahString::from("cluster-a"));
977 Ok(TopicList::default())
978 }
979 },
980 {
981 let global_calls = Arc::clone(&global_calls);
982 move || async move {
983 global_calls.fetch_add(1, Ordering::SeqCst);
984 Ok(TopicList::default())
985 }
986 },
987 )
988 .await
989 .unwrap();
990
991 assert!(result.topic_list.is_empty());
992 assert_eq!(cluster_calls.load(Ordering::SeqCst), 1);
993 assert_eq!(global_calls.load(Ordering::SeqCst), 0);
994 }
995
996 #[tokio::test]
997 async fn global_topic_inventory_selects_exactly_one_global_source_call() {
998 let cluster_calls = Arc::new(AtomicUsize::new(0));
999 let global_calls = Arc::new(AtomicUsize::new(0));
1000 let result = fetch_topic_inventory_from(
1001 None,
1002 {
1003 let cluster_calls = Arc::clone(&cluster_calls);
1004 move |_| async move {
1005 cluster_calls.fetch_add(1, Ordering::SeqCst);
1006 Ok(TopicList::default())
1007 }
1008 },
1009 {
1010 let global_calls = Arc::clone(&global_calls);
1011 move || async move {
1012 global_calls.fetch_add(1, Ordering::SeqCst);
1013 Ok(TopicList::default())
1014 }
1015 },
1016 )
1017 .await
1018 .unwrap();
1019
1020 assert!(result.topic_list.is_empty());
1021 assert_eq!(cluster_calls.load(Ordering::SeqCst), 0);
1022 assert_eq!(global_calls.load(Ordering::SeqCst), 1);
1023 }
1024
1025 #[test]
1026 fn parses_only_the_requested_allowlisted_value() {
1027 let properties = HashMap::from([
1028 (
1029 CheetahString::from_static_str("sendMessageThreadPoolNums"),
1030 CheetahString::from_static_str("32"),
1031 ),
1032 (
1033 CheetahString::from_static_str("maxClientEventCount"),
1034 CheetahString::from_static_str("100"),
1035 ),
1036 (
1037 CheetahString::from_static_str("accessKey"),
1038 CheetahString::from_static_str("must-not-cross-boundary"),
1039 ),
1040 ]);
1041
1042 let value = parse_allowlisted_value::<u32>(&properties, "sendMessageThreadPoolNums");
1043 assert_eq!(value.unwrap(), Some(32));
1044 assert_eq!(
1045 parse_allowlisted_value::<i32>(&properties, "maxClientEventCount").unwrap(),
1046 Some(100)
1047 );
1048 assert_eq!(
1049 parse_allowlisted_value::<u32>(&properties, "pullMessageThreadPoolNums").unwrap(),
1050 None
1051 );
1052 }
1053
1054 #[test]
1055 fn rejects_malformed_allowlisted_values() {
1056 let properties = HashMap::from([(
1057 CheetahString::from_static_str("flushDelayOffsetInterval"),
1058 CheetahString::from_static_str("not-a-number"),
1059 )]);
1060
1061 let error = parse_allowlisted_value::<u64>(&properties, "flushDelayOffsetInterval").unwrap_err();
1062 assert_invalid_argument(&error);
1063 }
1064
1065 #[test]
1066 fn consume_stats_targets_preserve_mixed_valid_and_unusable_route_entries() {
1067 let targets = consume_stats_read_targets(vec![
1068 route_broker("broker-a", Some("127.0.0.1:10911")),
1069 route_broker("broker-b", None),
1070 route_broker("broker-c", Some(" ")),
1071 ]);
1072
1073 assert_eq!(targets.len(), 3);
1074 assert!(matches!(
1075 &targets[0],
1076 ConsumeStatsReadTarget::Ready { broker_name, .. } if broker_name == "broker-a"
1077 ));
1078 for target in &targets[1..] {
1079 let ConsumeStatsReadTarget::Failed(failure) = target else {
1080 panic!("unusable master must remain a failed routed attempt");
1081 };
1082 assert_eq!(failure.code(), ReadFailureCode::InvalidResponse);
1083 assert!(!failure.retryable());
1084 }
1085 }
1086
1087 #[test]
1088 fn consume_stats_targets_preserve_all_missing_masters_as_failures() {
1089 let targets = consume_stats_read_targets(vec![route_broker("broker-a", None), route_broker("broker-b", None)]);
1090
1091 assert_eq!(targets.len(), 2);
1092 assert!(targets.iter().all(|target| matches!(
1093 target,
1094 ConsumeStatsReadTarget::Failed(failure)
1095 if failure.code() == ReadFailureCode::InvalidResponse && !failure.retryable()
1096 )));
1097 }
1098
1099 #[test]
1100 fn consume_stats_targets_keep_zero_broker_route_authoritatively_empty() {
1101 assert!(consume_stats_read_targets(Vec::new()).is_empty());
1102 }
1103
1104 #[test]
1105 fn broker_read_failure_preserves_valid_logical_targets() {
1106 let failure = BrokerReadFailure::new(" broker-a_1.prod ", ReadFailureCode::Timeout, true);
1107
1108 assert_eq!(failure.broker_name(), "broker-a_1.prod");
1109 assert_eq!(failure.code(), ReadFailureCode::Timeout);
1110 assert!(failure.retryable());
1111 }
1112
1113 #[test]
1114 fn broker_read_failure_replaces_unsafe_targets_with_one_safe_token() {
1115 let overlong = "a".repeat(129);
1116 let unsafe_targets = [
1117 "127.0.0.1",
1118 "::1",
1119 "127.0.0.1:10911",
1120 "[::1]:10911",
1121 "https://broker.example.invalid/runtime",
1122 "broker\nname",
1123 "broker-a\r",
1124 overlong.as_str(),
1125 "password-secret-value",
1126 "AKIAIOSFODNN7EXAMPLE",
1127 "connection_error_from_backend",
1128 ];
1129
1130 for target in unsafe_targets {
1131 let failure = BrokerReadFailure::new(target, ReadFailureCode::SourceUnavailable, true);
1132 assert_eq!(
1133 failure.broker_name(),
1134 "unknown",
1135 "target should be rejected: {target:?}"
1136 );
1137 }
1138 }
1139}