rocketmq_protocol/protocol/remoting_command/
frame.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 bytes::BufMut;
16use bytes::Bytes;
17use bytes::BytesMut;
18
19use super::RemotingCommand;
20use super::SerializeType;
21
22impl RemotingCommand {
23    #[inline]
24    pub(super) fn try_encode_header_with_body_length(&mut self, body_length: usize) -> rocketmq_error::Result<Bytes> {
25        // Encode header data
26        let header_data = self.try_header_encode()?;
27        let header_len = header_data.len();
28        let (total_length, marked_header_length) =
29            Self::checked_frame_lengths(header_len, body_length, self.serialize_type)?;
30
31        // Allocate exact capacity
32        let mut result = BytesMut::with_capacity(8 + header_len);
33
34        // Write total length
35        result.put_i32(total_length);
36
37        // Write serialize type with embedded header length
38        result.put_i32(marked_header_length);
39
40        // Write header data
41        result.put(header_data);
42
43        Ok(result.freeze())
44    }
45
46    pub(super) fn checked_frame_lengths(
47        header_length: usize,
48        body_length: usize,
49        serialize_type: SerializeType,
50    ) -> rocketmq_error::Result<(i32, i32)> {
51        const MAX_HEADER_LENGTH: usize = 0x00ff_ffff;
52        if header_length > MAX_HEADER_LENGTH {
53            return Err(crate::error::serialization_encode_failed(
54                "remoting-command",
55                format!("encoded header is {header_length} bytes, exceeding the 24-bit wire limit"),
56            ));
57        }
58        let payload_length = 4usize
59            .checked_add(header_length)
60            .and_then(|length| length.checked_add(body_length))
61            .ok_or_else(|| {
62                crate::error::serialization_encode_failed("remoting-command", "encoded frame length overflow")
63            })?;
64        let total_length = i32::try_from(payload_length).map_err(|_| {
65            crate::error::serialization_encode_failed(
66                "remoting-command",
67                format!("encoded payload is {payload_length} bytes, exceeding the signed 32-bit wire limit"),
68            )
69        })?;
70        let header_length = i32::try_from(header_length).map_err(|_| {
71            crate::error::serialization_encode_failed("remoting-command", "encoded header length overflow")
72        })?;
73        Ok((
74            total_length,
75            RemotingCommand::mark_serialize_type(header_length, serialize_type),
76        ))
77    }
78
79    #[inline]
80    pub fn mark_serialize_type(header_length: i32, protocol_type: SerializeType) -> i32 {
81        ((protocol_type.get_code() as i32) << 24) | (header_length & 0x00FFFFFF)
82    }
83}
rocketmq_protocol/protocol/remoting_command/frame.rs

rocketmq_protocol/protocol/remoting_command/
frame.rs
