rocketmq-protocol 1.0.0

// Copyright 2023 The RocketMQ Rust Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::collections::HashMap;

use rocketmq_model::message::MessageQueue;
use serde::de::Error as DeError;
use serde::de::MapAccess;
use serde::de::Visitor;
use serde::Deserialize;
use serde::Deserializer;
use serde::Serialize;
use serde_json_any_key::*;

use crate::protocol::admin::offset_wrapper::OffsetWrapper;
use crate::protocol::RemotingDeserializable;

#[derive(Debug, Default, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ConsumeStats {
    #[serde(with = "any_key_map")]
    pub offset_table: HashMap<MessageQueue, OffsetWrapper>,
    pub consume_tps: f64,
}

#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
struct StrictConsumeStats {
    #[serde(deserialize_with = "deserialize_strict_offset_table")]
    offset_table: HashMap<MessageQueue, OffsetWrapper>,
    consume_tps: f64,
}

fn deserialize_strict_offset_table<'de, D>(deserializer: D) -> Result<HashMap<MessageQueue, OffsetWrapper>, D::Error>
where
    D: Deserializer<'de>,
{
    struct StrictOffsetTableVisitor;

    impl<'de> Visitor<'de> for StrictOffsetTableVisitor {
        type Value = HashMap<MessageQueue, OffsetWrapper>;

        fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
            formatter.write_str("an offset table with unique MessageQueue keys")
        }

        fn visit_map<A>(self, mut access: A) -> Result<Self::Value, A::Error>
        where
            A: MapAccess<'de>,
        {
            let mut offsets = HashMap::with_capacity(access.size_hint().unwrap_or_default());
            while let Some((encoded_queue, offset)) = access.next_entry::<String, OffsetWrapper>()? {
                let queue = serde_json::from_str::<MessageQueue>(&encoded_queue).map_err(A::Error::custom)?;
                if offsets.insert(queue, offset).is_some() {
                    return Err(A::Error::custom("duplicate MessageQueue key in offset table"));
                }
            }
            Ok(offsets)
        }
    }

    deserializer.deserialize_map(StrictOffsetTableVisitor)
}

impl ConsumeStats {
    pub fn new() -> Self {
        Self {
            offset_table: HashMap::new(),
            consume_tps: 0.0,
        }
    }

    pub fn compute_total_diff(&self) -> i64 {
        self.offset_table
            .values()
            .map(|value| value.get_broker_offset() - value.get_consumer_offset())
            .sum()
    }

    pub fn compute_inflight_total_diff(&self) -> i64 {
        self.offset_table
            .values()
            .map(|value| value.get_pull_offset() - value.get_consumer_offset())
            .sum()
    }

    pub fn get_offset_table(&self) -> &HashMap<MessageQueue, OffsetWrapper> {
        &self.offset_table
    }

    pub fn get_offset_table_mut(&mut self) -> &mut HashMap<MessageQueue, OffsetWrapper> {
        &mut self.offset_table
    }

    pub fn set_offset_table(&mut self, offset_table: HashMap<MessageQueue, OffsetWrapper>) {
        self.offset_table = offset_table;
    }

    pub fn get_consume_tps(&self) -> f64 {
        self.consume_tps
    }

    pub fn set_consume_tps(&mut self, consume_tps: f64) {
        self.consume_tps = consume_tps;
    }

    pub fn decode(body: &[u8]) -> rocketmq_error::Result<Self> {
        match <Self as RemotingDeserializable>::decode(body) {
            Ok(stats) => Ok(stats),
            Err(error) => {
                let Ok(raw_body) = std::str::from_utf8(body) else {
                    return Err(error);
                };
                let normalized_body = normalize_nonstandard_offset_table_keys(raw_body);
                if normalized_body == raw_body {
                    return Err(error);
                }
                <Self as RemotingDeserializable>::decode_str(&normalized_body)
            }
        }
    }

    /// Decodes a consume-stats body while rejecting duplicate logical queue keys.
    pub fn decode_strict(body: &[u8]) -> rocketmq_error::Result<Self> {
        let raw_body = std::str::from_utf8(body)
            .map_err(|source| crate::error::serialization_source("deserialize", "JSON", source))?;
        let normalized_body = normalize_nonstandard_offset_table_keys(raw_body);
        let strict = serde_json::from_str::<StrictConsumeStats>(&normalized_body)
            .map_err(|error| crate::error::serialization_source("deserialize", "JSON", error))?;
        Ok(Self {
            offset_table: strict.offset_table,
            consume_tps: strict.consume_tps,
        })
    }

    pub fn encode_java_compatible(&self) -> rocketmq_error::Result<Vec<u8>> {
        Ok(self.to_java_compatible_json()?.into_bytes())
    }

    pub fn to_java_compatible_json(&self) -> rocketmq_error::Result<String> {
        let mut body = String::new();
        body.push_str("{\"offsetTable\":{");

        for (index, (queue, offset)) in self.offset_table.iter().enumerate() {
            if index > 0 {
                body.push(',');
            }
            append_message_queue_object_key(&mut body, queue)?;
            body.push(':');
            body.push_str(
                &serde_json::to_string(offset)
                    .map_err(|error| crate::error::serialization_source("serialize", "JSON", error))?,
            );
        }

        body.push_str("},\"consumeTps\":");
        body.push_str(
            &serde_json::to_string(&self.consume_tps)
                .map_err(|error| crate::error::serialization_source("serialize", "JSON", error))?,
        );
        body.push('}');
        Ok(body)
    }
}

pub(crate) fn normalize_nonstandard_offset_table_keys(input: &str) -> String {
    let chars = input.chars().collect::<Vec<_>>();
    let mut output = String::with_capacity(input.len());
    let mut index = 0;

    while index < chars.len() {
        if chars[index] != '"' {
            output.push(chars[index]);
            index += 1;
            continue;
        }

        let Some(string_end) = consume_string_end(&chars, index) else {
            return input.to_owned();
        };
        output.extend(chars[index..=string_end].iter());
        let key = chars[index + 1..string_end].iter().collect::<String>();
        index = string_end + 1;

        if key != "offsetTable" {
            continue;
        }

        while index < chars.len() && chars[index].is_whitespace() {
            output.push(chars[index]);
            index += 1;
        }
        if index >= chars.len() || chars[index] != ':' {
            continue;
        }
        output.push(':');
        index += 1;

        while index < chars.len() && chars[index].is_whitespace() {
            output.push(chars[index]);
            index += 1;
        }
        if index >= chars.len() || chars[index] != '{' {
            continue;
        }

        output.push('{');
        index += 1;
        let Some((normalized_map, next_index)) = normalize_object_key_map_body(&chars, index) else {
            return input.to_owned();
        };
        output.push_str(&normalized_map);
        index = next_index;
    }
    output
}

pub(crate) fn append_message_queue_object_key(output: &mut String, queue: &MessageQueue) -> rocketmq_error::Result<()> {
    output.push('{');
    output.push_str("\"topic\":");
    output.push_str(
        &serde_json::to_string(queue.topic_str())
            .map_err(|error| crate::error::serialization_source("serialize", "JSON", error))?,
    );
    output.push_str(",\"brokerName\":");
    output.push_str(
        &serde_json::to_string(queue.broker_name().as_str())
            .map_err(|error| crate::error::serialization_source("serialize", "JSON", error))?,
    );
    output.push_str(",\"queueId\":");
    output.push_str(&queue.queue_id().to_string());
    output.push('}');
    Ok(())
}

fn normalize_object_key_map_body(chars: &[char], mut index: usize) -> Option<(String, usize)> {
    let mut output = String::new();
    let mut expecting_key = true;
    let mut nested_value_depth = 0usize;
    let length = chars.len();

    while index < length {
        let current = chars[index];
        if expecting_key {
            match current {
                '}' => {
                    output.push('}');
                    return Some((output, index + 1));
                }
                '{' => {
                    let object_end = consume_balanced_object_end(chars, index)?;
                    let raw_key = chars[index..=object_end].iter().collect::<String>();
                    append_json_string(&mut output, &raw_key);
                    index = object_end + 1;
                    expecting_key = false;
                }
                '"' => {
                    let string_end = consume_string_end(chars, index)?;
                    output.extend(chars[index..=string_end].iter());
                    index = string_end + 1;
                    expecting_key = false;
                }
                _ => {
                    output.push(current);
                    index += 1;
                }
            }
            continue;
        }

        match current {
            '"' => {
                let string_end = consume_string_end(chars, index)?;
                output.extend(chars[index..=string_end].iter());
                index = string_end + 1;
            }
            '{' | '[' => {
                nested_value_depth += 1;
                output.push(current);
                index += 1;
            }
            '}' => {
                if nested_value_depth == 0 {
                    output.push('}');
                    return Some((output, index + 1));
                }
                nested_value_depth -= 1;
                output.push('}');
                index += 1;
            }
            ']' => {
                nested_value_depth = nested_value_depth.saturating_sub(1);
                output.push(']');
                index += 1;
            }
            ',' if nested_value_depth == 0 => {
                output.push(',');
                index += 1;
                expecting_key = true;
            }
            _ => {
                output.push(current);
                index += 1;
            }
        }
    }

    Some((output, index))
}

fn append_json_string(output: &mut String, value: &str) {
    output.push('"');
    for ch in value.chars() {
        match ch {
            '"' => output.push_str("\\\""),
            '\\' => output.push_str("\\\\"),
            '\u{08}' => output.push_str("\\b"),
            '\u{0C}' => output.push_str("\\f"),
            '\n' => output.push_str("\\n"),
            '\r' => output.push_str("\\r"),
            '\t' => output.push_str("\\t"),
            ch if ch <= '\u{1F}' => {
                output.push_str("\\u00");
                const HEX: &[u8; 16] = b"0123456789abcdef";
                let byte = ch as u8;
                output.push(HEX[(byte >> 4) as usize] as char);
                output.push(HEX[(byte & 0x0f) as usize] as char);
            }
            ch => output.push(ch),
        }
    }
    output.push('"');
}

fn consume_string_end(chars: &[char], start: usize) -> Option<usize> {
    let mut index = start + 1;
    let mut escaped = false;
    while index < chars.len() {
        if escaped {
            escaped = false;
        } else if chars[index] == '\\' {
            escaped = true;
        } else if chars[index] == '"' {
            return Some(index);
        }
        index += 1;
    }
    None
}

fn consume_balanced_object_end(chars: &[char], start: usize) -> Option<usize> {
    let mut depth = 0usize;
    let mut index = start;
    let mut in_string = false;
    let mut escaped = false;

    while index < chars.len() {
        let current = chars[index];
        if in_string {
            if escaped {
                escaped = false;
            } else if current == '\\' {
                escaped = true;
            } else if current == '"' {
                in_string = false;
            }
        } else {
            match current {
                '"' => in_string = true,
                '{' => depth += 1,
                '}' => {
                    depth -= 1;
                    if depth == 0 {
                        return Some(index);
                    }
                }
                _ => {}
            }
        }
        index += 1;
    }

    None
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn decode_rejects_unterminated_strings_without_panicking() {
        for body in [
            "\"",
            "{\"offsetTable",
            "{\"offsetTable\":{{\"topic\":\"unfinished",
            "{\"offsetTable\":{{\"topic\":\"topic\"",
            "{\"offsetTable\":{\"",
            "{\"offsetTable\":{\"key\":\"value",
            "{\"offsetTable\":{\"key\":\"value\\",
        ] {
            assert!(ConsumeStats::decode(body.as_bytes()).is_err(), "{body:?}");
            assert!(ConsumeStats::decode_strict(body.as_bytes()).is_err(), "{body:?}");
            assert_eq!(normalize_nonstandard_offset_table_keys(body), body);
        }
    }

    #[test]
    fn decode_preserves_standard_json_and_rejects_unrelated_malformed_input() {
        let decoded = ConsumeStats::decode(br#"{"offsetTable":{},"consumeTps":1.5}"#).unwrap();
        assert!(decoded.get_offset_table().is_empty());
        assert_eq!(decoded.get_consume_tps(), 1.5);
        for body in [r#"{"offsetTable":{},"consumeTps":}"#, r#"{"offsetTable":{},}"#] {
            assert_eq!(normalize_nonstandard_offset_table_keys(body), body);
            assert!(ConsumeStats::decode(body.as_bytes()).is_err());
        }
    }

    fn create_offset_wrapper(broker: i64, consumer: i64, pull: i64) -> OffsetWrapper {
        let mut wrapper = OffsetWrapper::new();
        wrapper.set_broker_offset(broker);
        wrapper.set_consumer_offset(consumer);
        wrapper.set_pull_offset(pull);
        wrapper
    }

    fn create_mq(topic: &str, queue_id: i32) -> MessageQueue {
        let json = format!(
            r#"{{"topic": "{}", "brokerName": "broker_1", "queueId": {}}}"#,
            topic, queue_id
        );
        serde_json::from_str(&json).expect("Failed to create MessageQueue for test")
    }

    #[test]
    fn methods_compute_consumer_lag() {
        let mut stats = ConsumeStats::new();
        assert_eq!(stats.get_consume_tps(), 0.0);
        assert!(stats.get_offset_table().is_empty());
        assert_eq!(stats.compute_total_diff(), 0);
        assert_eq!(stats.compute_inflight_total_diff(), 0);

        stats.set_consume_tps(5.67);
        let mut map = HashMap::new();
        map.insert(create_mq("topic_a", 0), create_offset_wrapper(500, 400, 450));
        stats.set_offset_table(map);
        stats
            .get_offset_table_mut()
            .insert(create_mq("topic_a", 1), create_offset_wrapper(200, 150, 180));

        assert_eq!(stats.get_consume_tps(), 5.67);
        assert_eq!(stats.get_offset_table().len(), 2);
        assert_eq!(stats.compute_total_diff(), 150);
        assert_eq!(stats.compute_inflight_total_diff(), 80);
    }

    #[test]
    fn test_serialization_json_any_key() {
        let mut stats = ConsumeStats::new();
        let mq = create_mq("order", 2);
        stats.get_offset_table_mut().insert(mq, create_offset_wrapper(10, 5, 8));
        stats.set_consume_tps(1.0);

        let serialized = serde_json::to_string(&stats).expect("Serialization failed");

        assert!(serialized.contains("offsetTable"));
        assert!(serialized.contains("consumeTps"));

        let deserialized: ConsumeStats = serde_json::from_str(&serialized).expect("Deserialization failed");
        assert_eq!(deserialized.get_offset_table().len(), 1);
        assert_eq!(deserialized.compute_total_diff(), 5);
    }

    #[test]
    fn decode_accepts_nonstandard_offset_table_object_keys() {
        let body = r#"{"offsetTable":{{"topic":"TopicTest","brokerName":"broker-a","queueId":1}:{"brokerOffset":120,"consumerOffset":20,"pullOffset":80,"lastTimestamp":1700000000000}},"consumeTps":1.5}"#;

        let decoded = ConsumeStats::decode(body.as_bytes()).expect("consume stats should decode");

        assert_eq!(decoded.get_offset_table().len(), 1);
        assert_eq!(decoded.get_consume_tps(), 1.5);
        let (queue, offset) = decoded.get_offset_table().iter().next().expect("one queue");
        assert_eq!(queue.topic_str(), "TopicTest");
        assert_eq!(queue.broker_name(), "broker-a");
        assert_eq!(queue.queue_id(), 1);
        assert_eq!(offset.get_broker_offset(), 120);
        assert_eq!(offset.get_consumer_offset(), 20);
        assert_eq!(offset.get_pull_offset(), 80);
        assert_eq!(offset.get_last_timestamp(), 1_700_000_000_000);
    }

    #[test]
    fn strict_decode_rejects_duplicate_logical_queue_keys() {
        let queue = serde_json::to_string(&create_mq("TopicTest", 1)).expect("queue key");
        let encoded_key = serde_json::to_string(&queue).expect("encoded key");
        let offset = r#"{"brokerOffset":120,"consumerOffset":20,"pullOffset":80,"lastTimestamp":0}"#;
        let body = format!(r#"{{"offsetTable":{{{encoded_key}:{offset},{encoded_key}:{offset}}},"consumeTps":1.5}}"#);

        assert!(ConsumeStats::decode_strict(body.as_bytes()).is_err());
    }

    #[test]
    fn normalize_nonstandard_offset_table_keys_stringifies_object_keys() {
        let input = r#"{"offsetTable":{{"topic":"TopicTest","brokerName":"broker-a","queueId":1}:{"brokerOffset":120}},"consumeTps":1.5}"#;
        let normalized = normalize_nonstandard_offset_table_keys(input);

        assert_eq!(
            normalized,
            r#"{"offsetTable":{"{\"topic\":\"TopicTest\",\"brokerName\":\"broker-a\",\"queueId\":1}":{"brokerOffset":120}},"consumeTps":1.5}"#
        );
    }

    #[test]
    fn normalize_nonstandard_offset_table_keys_escapes_object_key_strings() {
        let input = r#"{"offsetTable":{{"topic":"Topic\"Quoted","brokerName":"broker\\one","queueId":1}:{"brokerOffset":120}},"consumeTps":1.5}"#;
        let normalized = normalize_nonstandard_offset_table_keys(input);
        let value: serde_json::Value = serde_json::from_str(&normalized).expect("normalized JSON should parse");
        let offset_table = value["offsetTable"].as_object().expect("offsetTable object");
        let key = offset_table.keys().next().expect("one queue key");

        assert_eq!(
            key,
            r#"{"topic":"Topic\"Quoted","brokerName":"broker\\one","queueId":1}"#
        );
    }

    #[test]
    fn java_compatible_encode_uses_object_keys_for_offset_table() {
        let mut stats = ConsumeStats::new();
        stats
            .get_offset_table_mut()
            .insert(create_mq("TopicTest", 1), create_offset_wrapper(120, 20, 80));
        stats.set_consume_tps(1.5);

        let encoded = String::from_utf8(stats.encode_java_compatible().expect("encode consume stats"))
            .expect("utf8 consume stats");

        assert!(encoded.contains(r#""offsetTable":{{"topic":"TopicTest","brokerName":"broker_1","queueId":1}:"#));
        assert!(!encoded.contains(r#""{\"topic\""#));

        let decoded = ConsumeStats::decode(encoded.as_bytes()).expect("decode java-compatible consume stats");
        assert_eq!(decoded.get_offset_table().len(), 1);
        assert_eq!(decoded.compute_total_diff(), 100);
    }
}