Skip to main content

rocketmq_model/common/message/
routing_context.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::net::SocketAddr;
16
17use bytes::BufMut;
18use bytes::Bytes;
19use bytes::BytesMut;
20
21use crate::common::sys_flag::message_sys_flag::MessageSysFlag;
22
23/// Message network transmission metadata
24///
25/// Contains metadata generated during network transmission, such as client address, timestamp, and
26/// system flags.
27#[derive(Clone, Debug, PartialEq, Eq)]
28pub struct RoutingContext {
29    born_host: SocketAddr,
30    born_timestamp: i64,
31    sys_flag: i32,
32}
33
34impl RoutingContext {
35    /// Creates a new routing context
36    pub fn new(born_host: SocketAddr, born_timestamp: i64, sys_flag: i32) -> Self {
37        Self {
38            born_host,
39            born_timestamp,
40            sys_flag,
41        }
42    }
43
44    /// Gets the message sender's client address
45    #[inline]
46    pub fn born_host(&self) -> SocketAddr {
47        self.born_host
48    }
49
50    /// Gets the message creation timestamp
51    #[inline]
52    pub fn born_timestamp(&self) -> i64 {
53        self.born_timestamp
54    }
55
56    /// Gets the system flag bits
57    #[inline]
58    pub fn sys_flag(&self) -> i32 {
59        self.sys_flag
60    }
61
62    /// Sets the IPv6 flag for born_host
63    #[inline]
64    pub fn with_born_host_v6_flag(&mut self) {
65        self.sys_flag |= MessageSysFlag::BORNHOST_V6_FLAG;
66    }
67
68    /// Serializes born_host to bytes
69    pub fn born_host_bytes(&self) -> Bytes {
70        Self::socket_addr_to_bytes(&self.born_host)
71    }
72
73    fn socket_addr_to_bytes(addr: &SocketAddr) -> Bytes {
74        match addr {
75            SocketAddr::V4(v4) => {
76                let mut buf = BytesMut::with_capacity(8);
77                buf.put_slice(&v4.ip().octets());
78                buf.put_i32(v4.port() as i32);
79                buf.freeze()
80            }
81            SocketAddr::V6(v6) => {
82                let mut buf = BytesMut::with_capacity(20);
83                buf.put_slice(&v6.ip().octets());
84                buf.put_i32(v6.port() as i32);
85                buf.freeze()
86            }
87        }
88    }
89}
90
91impl Default for RoutingContext {
92    fn default() -> Self {
93        Self {
94            born_host: SocketAddr::from(([127, 0, 0, 1], 0)),
95            born_timestamp: 0,
96            sys_flag: 0,
97        }
98    }
99}
100
101#[cfg(test)]
102mod tests {
103    use super::*;
104
105    #[test]
106    fn test_routing_context_creation() {
107        let addr: SocketAddr = "192.168.1.100:10911".parse().unwrap();
108        let ctx = RoutingContext::new(addr, 1234567890, 0);
109
110        assert_eq!(ctx.born_host(), addr);
111        assert_eq!(ctx.born_timestamp(), 1234567890);
112        assert_eq!(ctx.sys_flag(), 0);
113    }
114
115    #[test]
116    fn test_with_ipv6_flag() {
117        let addr: SocketAddr = "127.0.0.1:10911".parse().unwrap();
118        let mut ctx = RoutingContext::new(addr, 0, 0);
119
120        ctx.with_born_host_v6_flag();
121        assert_ne!(ctx.sys_flag(), 0);
122    }
123
124    #[test]
125    fn test_born_host_bytes_v4() {
126        let addr: SocketAddr = "192.168.1.100:10911".parse().unwrap();
127        let ctx = RoutingContext::new(addr, 0, 0);
128
129        let bytes = ctx.born_host_bytes();
130        assert_eq!(bytes.len(), 8); // 4 bytes IP + 4 bytes port
131    }
132
133    #[test]
134    fn test_born_host_bytes_v6() {
135        let addr: SocketAddr = "[::1]:10911".parse().unwrap();
136        let ctx = RoutingContext::new(addr, 0, 0);
137
138        let bytes = ctx.born_host_bytes();
139        assert_eq!(bytes.len(), 20); // 16 bytes IP + 4 bytes port
140    }
141
142    #[test]
143    fn test_default() {
144        let ctx = RoutingContext::default();
145        assert_eq!(ctx.born_timestamp(), 0);
146        assert_eq!(ctx.sys_flag(), 0);
147    }
148}