iggy_binary_protocol/primitives/
partitioning.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use crate::WireError;
19use crate::codec::{WireDecode, WireEncode, read_bytes, read_u8, read_u32_le};
20use bytes::{BufMut, BytesMut};
21use std::borrow::Cow;
22
23const KIND_BALANCED: u8 = 1;
24const KIND_PARTITION_ID: u8 = 2;
25const KIND_MESSAGES_KEY: u8 = 3;
26
27/// Maximum key length for `MessagesKey` partitioning (u8 length prefix on the wire).
28pub const MAX_MESSAGES_KEY_LENGTH: usize = 255;
29
30/// Partitioning strategy for message routing.
31///
32/// Wire format: `[kind:1][length:1][value:0..255]`
33/// - `Balanced`:    kind=1, length=0, no value bytes
34/// - `PartitionId`: kind=2, length=4, value=u32 LE
35/// - `MessagesKey`: kind=3, length=1..255, value=raw bytes
36///
37/// Use [`WirePartitioning::messages_key`] to construct a `MessagesKey` variant
38/// with validation. Direct enum construction is possible but skips the length check.
39#[derive(Debug, Clone, PartialEq, Eq)]
40pub enum WirePartitioning {
41    Balanced,
42    PartitionId(u32),
43    MessagesKey(Vec<u8>),
44}
45
46impl WirePartitioning {
47    /// Create a `MessagesKey` partitioning with key length validation.
48    ///
49    /// # Errors
50    /// Returns `WireError::Validation` if `key` is empty or exceeds 255 bytes.
51    pub fn messages_key(key: Vec<u8>) -> Result<Self, WireError> {
52        if key.is_empty() {
53            return Err(WireError::Validation(Cow::Borrowed(
54                "messages_key partitioning cannot have empty key",
55            )));
56        }
57        if key.len() > MAX_MESSAGES_KEY_LENGTH {
58            return Err(WireError::Validation(Cow::Owned(format!(
59                "messages_key length {} exceeds maximum {MAX_MESSAGES_KEY_LENGTH}",
60                key.len()
61            ))));
62        }
63        Ok(Self::MessagesKey(key))
64    }
65}
66
67impl WireEncode for WirePartitioning {
68    fn encoded_size(&self) -> usize {
69        match self {
70            Self::Balanced => 2,
71            Self::PartitionId(_) => 2 + 4,
72            Self::MessagesKey(key) => 2 + key.len(),
73        }
74    }
75
76    fn encode(&self, buf: &mut BytesMut) {
77        match self {
78            Self::Balanced => {
79                buf.put_u8(KIND_BALANCED);
80                buf.put_u8(0);
81            }
82            Self::PartitionId(id) => {
83                buf.put_u8(KIND_PARTITION_ID);
84                buf.put_u8(4);
85                buf.put_u32_le(*id);
86            }
87            Self::MessagesKey(key) => {
88                debug_assert!(
89                    !key.is_empty() && key.len() <= MAX_MESSAGES_KEY_LENGTH,
90                    "MessagesKey length {} out of valid range 1..={MAX_MESSAGES_KEY_LENGTH}; \
91                     use WirePartitioning::messages_key() for validated construction",
92                    key.len()
93                );
94                buf.put_u8(KIND_MESSAGES_KEY);
95                #[allow(clippy::cast_possible_truncation)]
96                buf.put_u8(key.len() as u8);
97                buf.put_slice(key);
98            }
99        }
100    }
101}
102
103impl WireDecode for WirePartitioning {
104    fn decode(buf: &[u8]) -> Result<(Self, usize), WireError> {
105        let kind = read_u8(buf, 0)?;
106        let length = read_u8(buf, 1)? as usize;
107
108        match kind {
109            KIND_BALANCED => {
110                if length != 0 {
111                    return Err(WireError::Validation(Cow::Owned(format!(
112                        "balanced partitioning must have length 0, got {length}"
113                    ))));
114                }
115                Ok((Self::Balanced, 2))
116            }
117            KIND_PARTITION_ID => {
118                if length != 4 {
119                    return Err(WireError::Validation(Cow::Owned(format!(
120                        "partition_id partitioning must have length 4, got {length}"
121                    ))));
122                }
123                let id = read_u32_le(buf, 2)?;
124                Ok((Self::PartitionId(id), 6))
125            }
126            KIND_MESSAGES_KEY => {
127                if length == 0 {
128                    return Err(WireError::Validation(Cow::Borrowed(
129                        "messages_key partitioning cannot have empty key",
130                    )));
131                }
132                let key = read_bytes(buf, 2, length)?;
133                Ok((Self::MessagesKey(key.to_vec()), 2 + length))
134            }
135            _ => Err(WireError::UnknownDiscriminant {
136                type_name: "WirePartitioning",
137                value: kind,
138                offset: 0,
139            }),
140        }
141    }
142}
143
144#[cfg(test)]
145mod tests {
146    use super::*;
147
148    #[test]
149    fn roundtrip_balanced() {
150        let p = WirePartitioning::Balanced;
151        let bytes = p.to_bytes();
152        assert_eq!(bytes.len(), 2);
153        let (decoded, consumed) = WirePartitioning::decode(&bytes).unwrap();
154        assert_eq!(consumed, 2);
155        assert_eq!(decoded, p);
156    }
157
158    #[test]
159    fn roundtrip_partition_id() {
160        let p = WirePartitioning::PartitionId(42);
161        let bytes = p.to_bytes();
162        assert_eq!(bytes.len(), 6);
163        let (decoded, consumed) = WirePartitioning::decode(&bytes).unwrap();
164        assert_eq!(consumed, 6);
165        assert_eq!(decoded, p);
166    }
167
168    #[test]
169    fn roundtrip_messages_key() {
170        let p = WirePartitioning::MessagesKey(b"user-123".to_vec());
171        let bytes = p.to_bytes();
172        assert_eq!(bytes.len(), 2 + 8);
173        let (decoded, consumed) = WirePartitioning::decode(&bytes).unwrap();
174        assert_eq!(consumed, bytes.len());
175        assert_eq!(decoded, p);
176    }
177
178    #[test]
179    fn messages_key_max_length() {
180        let key = vec![0xAB; 255];
181        let p = WirePartitioning::MessagesKey(key.clone());
182        let bytes = p.to_bytes();
183        let (decoded, _) = WirePartitioning::decode(&bytes).unwrap();
184        assert_eq!(decoded, WirePartitioning::MessagesKey(key));
185    }
186
187    #[test]
188    fn messages_key_single_byte() {
189        let p = WirePartitioning::MessagesKey(vec![0x01]);
190        let bytes = p.to_bytes();
191        let (decoded, consumed) = WirePartitioning::decode(&bytes).unwrap();
192        assert_eq!(consumed, 3);
193        assert_eq!(decoded, p);
194    }
195
196    #[test]
197    fn empty_messages_key_rejected() {
198        let buf = [KIND_MESSAGES_KEY, 0x00];
199        assert!(WirePartitioning::decode(&buf).is_err());
200    }
201
202    #[test]
203    fn unknown_kind_rejected() {
204        let buf = [0xFF, 0x00];
205        assert!(WirePartitioning::decode(&buf).is_err());
206    }
207
208    #[test]
209    fn balanced_with_nonzero_length_rejected() {
210        let buf = [KIND_BALANCED, 0x01, 0x00];
211        assert!(WirePartitioning::decode(&buf).is_err());
212    }
213
214    #[test]
215    fn partition_id_with_wrong_length_rejected() {
216        let buf = [KIND_PARTITION_ID, 0x03, 0x00, 0x00, 0x00];
217        assert!(WirePartitioning::decode(&buf).is_err());
218    }
219
220    #[test]
221    fn truncated_buffer() {
222        let p = WirePartitioning::PartitionId(1);
223        let bytes = p.to_bytes();
224        for i in 0..bytes.len() {
225            assert!(
226                WirePartitioning::decode(&bytes[..i]).is_err(),
227                "expected error for truncation at byte {i}"
228            );
229        }
230    }
231
232    #[test]
233    fn wire_compat_balanced_layout() {
234        let bytes = WirePartitioning::Balanced.to_bytes();
235        assert_eq!(&bytes[..], &[KIND_BALANCED, 0x00]);
236    }
237
238    #[test]
239    fn wire_compat_partition_id_layout() {
240        let bytes = WirePartitioning::PartitionId(1).to_bytes();
241        assert_eq!(
242            &bytes[..],
243            &[KIND_PARTITION_ID, 0x04, 0x01, 0x00, 0x00, 0x00]
244        );
245    }
246
247    #[test]
248    fn messages_key_constructor_validates_empty() {
249        assert!(WirePartitioning::messages_key(vec![]).is_err());
250    }
251
252    #[test]
253    fn messages_key_constructor_validates_too_long() {
254        let key = vec![0xAB; 256];
255        assert!(WirePartitioning::messages_key(key).is_err());
256    }
257
258    #[test]
259    fn messages_key_constructor_accepts_max() {
260        let key = vec![0xAB; 255];
261        let p = WirePartitioning::messages_key(key.clone()).unwrap();
262        assert_eq!(p, WirePartitioning::MessagesKey(key));
263    }
264
265    #[test]
266    fn messages_key_constructor_accepts_single_byte() {
267        let p = WirePartitioning::messages_key(vec![0x01]).unwrap();
268        let bytes = p.to_bytes();
269        let (decoded, _) = WirePartitioning::decode(&bytes).unwrap();
270        assert_eq!(decoded, p);
271    }
272}
iggy_binary_protocol/primitives/partitioning.rs

iggy_binary_protocol/primitives/
partitioning.rs
