Skip to main content

stackforge_core/layer/generic/
mod.rs

1//! Generic / custom protocol layer for user-defined protocols.
2//!
3//! [`GenericLayer`] allows a protocol to be defined entirely at runtime
4//! (e.g. from Python via PyO3 bindings) as a flat list of named, typed
5//! fields.  Field descriptors ([`GenericFieldDesc`]) are shared between all
6//! packet instances using `Arc`, so the definition overhead is paid only once.
7//!
8//! # Example
9//!
10//! ```rust
11//! use std::sync::Arc;
12//! use stackforge_core::layer::generic::{GenericFieldDesc, GenericLayer};
13//! use stackforge_core::layer::field::{FieldType, FieldValue};
14//! use stackforge_core::layer::LayerIndex;
15//! use stackforge_core::layer::LayerKind;
16//!
17//! let descs = Arc::new(vec![
18//!     GenericFieldDesc {
19//!         name: "opcode".to_string(),
20//!         offset: 0,
21//!         size: 1,
22//!         field_type: FieldType::U8,
23//!         default_value: vec![0x00],
24//!     },
25//!     GenericFieldDesc {
26//!         name: "length".to_string(),
27//!         offset: 1,
28//!         size: 2,
29//!         field_type: FieldType::U16,
30//!         default_value: vec![0x00, 0x00],
31//!     },
32//! ]);
33//!
34//! let buf = vec![0x01u8, 0x00, 0x0A]; // opcode=1, length=10
35//! let index = LayerIndex::new(LayerKind::Generic, 0, 3);
36//! let layer = GenericLayer::new(index, Arc::from("MyProto"), descs);
37//!
38//! assert_eq!(
39//!     layer.get_field(&buf, "opcode").unwrap().unwrap(),
40//!     FieldValue::U8(1)
41//! );
42//! assert_eq!(
43//!     layer.get_field(&buf, "length").unwrap().unwrap(),
44//!     FieldValue::U16(10)
45//! );
46//! ```
47
48pub mod builder;
49pub use builder::GenericLayerBuilder;
50
51use std::sync::Arc;
52
53use crate::layer::field::{FieldError, FieldType, FieldValue};
54use crate::layer::{Layer, LayerIndex, LayerKind};
55
56/// Descriptor for a single field in a dynamically-defined protocol layer.
57#[derive(Debug, Clone)]
58pub struct GenericFieldDesc {
59    /// Field name (must be unique within the protocol).
60    pub name: String,
61    /// Byte offset from the start of the layer.
62    pub offset: usize,
63    /// Field size in bytes.
64    pub size: usize,
65    /// How to interpret the raw bytes.
66    pub field_type: FieldType,
67    /// Default value used by the builder when no explicit value is given.
68    /// Should be `size` bytes long; builder pads/truncates as needed.
69    pub default_value: Vec<u8>,
70}
71
72/// A dynamically-defined protocol layer for custom / user-defined protocols.
73///
74/// The field descriptors are shared (`Arc`) so they can be defined once and
75/// reused across many packet instances efficiently.
76#[derive(Debug, Clone)]
77pub struct GenericLayer {
78    /// Position of this layer within the packet buffer.
79    pub index: LayerIndex,
80    /// Protocol name as defined by the caller.
81    pub name: Arc<str>,
82    /// Shared field descriptors.
83    pub field_descs: Arc<Vec<GenericFieldDesc>>,
84}
85
86impl GenericLayer {
87    /// Create a new `GenericLayer`.
88    pub fn new(index: LayerIndex, name: Arc<str>, field_descs: Arc<Vec<GenericFieldDesc>>) -> Self {
89        Self {
90            index,
91            name,
92            field_descs,
93        }
94    }
95
96    /// Human-readable summary: `"<name> (<N> fields)"`.
97    pub fn summary(&self, _buf: &[u8]) -> String {
98        format!("{} ({} fields)", self.name, self.field_descs.len())
99    }
100
101    /// Header length = sum of all field sizes.
102    pub fn header_len(&self, _buf: &[u8]) -> usize {
103        self.field_descs.iter().map(|f| f.size).sum()
104    }
105
106    /// Dynamic list of field names (not `&'static str` because names come from
107    /// user-defined strings at runtime).
108    pub fn field_names_dynamic(&self) -> Vec<String> {
109        self.field_descs.iter().map(|f| f.name.clone()).collect()
110    }
111
112    /// Read a field from the buffer.
113    ///
114    /// Returns:
115    /// - `Some(Ok(value))` if the field was found and the buffer is long enough.
116    /// - `Some(Err(e))` if the field was found but the buffer is too short.
117    /// - `None` if no field with this name exists in this layer.
118    pub fn get_field(&self, buf: &[u8], name: &str) -> Option<Result<FieldValue, FieldError>> {
119        let desc = self.field_descs.iter().find(|f| f.name == name)?;
120        let layer_slice = self.index.slice(buf);
121
122        let field_end = desc.offset + desc.size;
123        if layer_slice.len() < field_end {
124            return Some(Err(FieldError::BufferTooShort {
125                offset: self.index.start + desc.offset,
126                need: desc.size,
127                have: layer_slice.len().saturating_sub(desc.offset),
128            }));
129        }
130
131        let field_slice = &layer_slice[desc.offset..field_end];
132
133        let value = match desc.field_type {
134            FieldType::U8 => FieldValue::U8(field_slice[0]),
135            FieldType::U16 => {
136                let v = u16::from_be_bytes([field_slice[0], field_slice[1]]);
137                FieldValue::U16(v)
138            }
139            FieldType::U32 => {
140                let v = u32::from_be_bytes([
141                    field_slice[0],
142                    field_slice[1],
143                    field_slice[2],
144                    field_slice[3],
145                ]);
146                FieldValue::U32(v)
147            }
148            FieldType::U64 => {
149                let v = u64::from_be_bytes([
150                    field_slice[0],
151                    field_slice[1],
152                    field_slice[2],
153                    field_slice[3],
154                    field_slice[4],
155                    field_slice[5],
156                    field_slice[6],
157                    field_slice[7],
158                ]);
159                FieldValue::U64(v)
160            }
161            FieldType::LEU16 => {
162                let v = u16::from_le_bytes([field_slice[0], field_slice[1]]);
163                FieldValue::U16(v)
164            }
165            FieldType::LEU32 => {
166                let v = u32::from_le_bytes([
167                    field_slice[0],
168                    field_slice[1],
169                    field_slice[2],
170                    field_slice[3],
171                ]);
172                FieldValue::U32(v)
173            }
174            FieldType::LEU64 => {
175                let v = u64::from_le_bytes([
176                    field_slice[0],
177                    field_slice[1],
178                    field_slice[2],
179                    field_slice[3],
180                    field_slice[4],
181                    field_slice[5],
182                    field_slice[6],
183                    field_slice[7],
184                ]);
185                FieldValue::U64(v)
186            }
187            FieldType::Bool => FieldValue::Bool(field_slice[0] != 0),
188            // For all other types, return raw bytes.
189            _ => FieldValue::Bytes(field_slice.to_vec()),
190        };
191
192        Some(Ok(value))
193    }
194
195    /// Write a field into the (mutable) buffer.
196    ///
197    /// Returns:
198    /// - `Some(Ok(()))` if the field was found and written successfully.
199    /// - `Some(Err(e))` if the field was found but the buffer is too short or
200    ///   the value type is wrong.
201    /// - `None` if no field with this name exists in this layer.
202    pub fn set_field(
203        &self,
204        buf: &mut [u8],
205        name: &str,
206        value: FieldValue,
207    ) -> Option<Result<(), FieldError>> {
208        let desc = self.field_descs.iter().find(|f| f.name == name)?;
209
210        let abs_offset = self.index.start + desc.offset;
211        let abs_end = abs_offset + desc.size;
212
213        if buf.len() < abs_end {
214            return Some(Err(FieldError::BufferTooShort {
215                offset: abs_offset,
216                need: desc.size,
217                have: buf.len().saturating_sub(abs_offset),
218            }));
219        }
220
221        let dest = &mut buf[abs_offset..abs_end];
222
223        let result = match (desc.field_type, &value) {
224            (FieldType::U8, FieldValue::U8(v)) => {
225                dest[0] = *v;
226                Ok(())
227            }
228            (FieldType::U16, FieldValue::U16(v)) => {
229                dest.copy_from_slice(&v.to_be_bytes());
230                Ok(())
231            }
232            (FieldType::U32, FieldValue::U32(v)) => {
233                dest.copy_from_slice(&v.to_be_bytes());
234                Ok(())
235            }
236            (FieldType::U64, FieldValue::U64(v)) => {
237                dest.copy_from_slice(&v.to_be_bytes());
238                Ok(())
239            }
240            (FieldType::LEU16, FieldValue::U16(v)) => {
241                dest.copy_from_slice(&v.to_le_bytes());
242                Ok(())
243            }
244            (FieldType::LEU32, FieldValue::U32(v)) => {
245                dest.copy_from_slice(&v.to_le_bytes());
246                Ok(())
247            }
248            (FieldType::LEU64, FieldValue::U64(v)) => {
249                dest.copy_from_slice(&v.to_le_bytes());
250                Ok(())
251            }
252            (FieldType::Bool, FieldValue::Bool(v)) => {
253                dest[0] = if *v { 1 } else { 0 };
254                Ok(())
255            }
256            (_, FieldValue::Bytes(bytes)) => {
257                let copy_len = bytes.len().min(dest.len());
258                dest[..copy_len].copy_from_slice(&bytes[..copy_len]);
259                Ok(())
260            }
261            _ => Err(FieldError::InvalidValue(format!(
262                "field '{}': value type {:?} does not match declared field type {:?}",
263                name, value, desc.field_type
264            ))),
265        };
266
267        Some(result)
268    }
269}
270
271impl Layer for GenericLayer {
272    fn kind(&self) -> LayerKind {
273        LayerKind::Generic
274    }
275
276    fn summary(&self, data: &[u8]) -> String {
277        self.summary(data)
278    }
279
280    fn header_len(&self, data: &[u8]) -> usize {
281        self.header_len(data)
282    }
283
284    /// The `Layer` trait requires `&'static [&'static str]`, which is not
285    /// possible for dynamically-generated field names.  Returns an empty
286    /// slice; callers that need the full list should call
287    /// [`field_names_dynamic`](Self::field_names_dynamic) directly.
288    fn field_names(&self) -> &'static [&'static str] {
289        &[]
290    }
291
292    fn hashret(&self, _data: &[u8]) -> Vec<u8> {
293        vec![]
294    }
295}
296
297#[cfg(test)]
298mod tests {
299    use super::*;
300    use crate::layer::LayerKind;
301
302    fn make_descs() -> Arc<Vec<GenericFieldDesc>> {
303        Arc::new(vec![
304            GenericFieldDesc {
305                name: "opcode".to_string(),
306                offset: 0,
307                size: 1,
308                field_type: FieldType::U8,
309                default_value: vec![0x00],
310            },
311            GenericFieldDesc {
312                name: "length".to_string(),
313                offset: 1,
314                size: 2,
315                field_type: FieldType::U16,
316                default_value: vec![0x00, 0x00],
317            },
318            GenericFieldDesc {
319                name: "seq".to_string(),
320                offset: 3,
321                size: 4,
322                field_type: FieldType::U32,
323                default_value: vec![0x00, 0x00, 0x00, 0x00],
324            },
325            GenericFieldDesc {
326                name: "payload".to_string(),
327                offset: 7,
328                size: 8,
329                field_type: FieldType::Bytes,
330                default_value: vec![0u8; 8],
331            },
332        ])
333    }
334
335    fn make_layer_and_buf() -> (GenericLayer, Vec<u8>) {
336        let descs = make_descs();
337        let buf = vec![
338            0x05u8, // opcode = 5
339            0x00, 0x10, // length = 16
340            0x00, 0x00, 0x00, 0x01, // seq = 1
341            0xDE, 0xAD, 0xBE, 0xEF, 0x00, 0x00, 0x00, 0x00, // payload
342        ];
343        let index = LayerIndex::new(LayerKind::Generic, 0, 15);
344        let layer = GenericLayer::new(index, Arc::from("TestProto"), descs);
345        (layer, buf)
346    }
347
348    #[test]
349    fn test_summary() {
350        let (layer, buf) = make_layer_and_buf();
351        let s = layer.summary(&buf);
352        assert!(s.contains("TestProto"), "summary: {}", s);
353        assert!(s.contains("4 fields"), "summary: {}", s);
354    }
355
356    #[test]
357    fn test_header_len() {
358        let (layer, buf) = make_layer_and_buf();
359        // 1 + 2 + 4 + 8 = 15
360        assert_eq!(layer.header_len(&buf), 15);
361    }
362
363    #[test]
364    fn test_get_field_u8() {
365        let (layer, buf) = make_layer_and_buf();
366        let v = layer.get_field(&buf, "opcode").unwrap().unwrap();
367        assert_eq!(v, FieldValue::U8(5));
368    }
369
370    #[test]
371    fn test_get_field_u16() {
372        let (layer, buf) = make_layer_and_buf();
373        let v = layer.get_field(&buf, "length").unwrap().unwrap();
374        assert_eq!(v, FieldValue::U16(16));
375    }
376
377    #[test]
378    fn test_get_field_u32() {
379        let (layer, buf) = make_layer_and_buf();
380        let v = layer.get_field(&buf, "seq").unwrap().unwrap();
381        assert_eq!(v, FieldValue::U32(1));
382    }
383
384    #[test]
385    fn test_get_field_bytes() {
386        let (layer, buf) = make_layer_and_buf();
387        let v = layer.get_field(&buf, "payload").unwrap().unwrap();
388        if let FieldValue::Bytes(b) = v {
389            assert_eq!(&b[..4], &[0xDE, 0xAD, 0xBE, 0xEF]);
390        } else {
391            panic!("expected Bytes, got {:?}", v);
392        }
393    }
394
395    #[test]
396    fn test_get_field_unknown() {
397        let (layer, buf) = make_layer_and_buf();
398        assert!(layer.get_field(&buf, "nonexistent").is_none());
399    }
400
401    #[test]
402    fn test_get_field_buffer_too_short() {
403        let descs = make_descs();
404        let buf = vec![0x05u8]; // only 1 byte — "length" won't fit
405        let index = LayerIndex::new(LayerKind::Generic, 0, 1);
406        let layer = GenericLayer::new(index, Arc::from("T"), descs);
407        let result = layer.get_field(&buf, "length").unwrap();
408        assert!(result.is_err());
409    }
410
411    #[test]
412    fn test_set_field_u8() {
413        let (layer, mut buf) = make_layer_and_buf();
414        layer
415            .set_field(&mut buf, "opcode", FieldValue::U8(0xFF))
416            .unwrap()
417            .unwrap();
418        assert_eq!(buf[0], 0xFF);
419    }
420
421    #[test]
422    fn test_set_field_u16() {
423        let (layer, mut buf) = make_layer_and_buf();
424        layer
425            .set_field(&mut buf, "length", FieldValue::U16(0x0100))
426            .unwrap()
427            .unwrap();
428        assert_eq!(buf[1], 0x01);
429        assert_eq!(buf[2], 0x00);
430    }
431
432    #[test]
433    fn test_set_field_u32() {
434        let (layer, mut buf) = make_layer_and_buf();
435        layer
436            .set_field(&mut buf, "seq", FieldValue::U32(0xDEAD_BEEF))
437            .unwrap()
438            .unwrap();
439        assert_eq!(&buf[3..7], &[0xDE, 0xAD, 0xBE, 0xEF]);
440    }
441
442    #[test]
443    fn test_set_field_bytes() {
444        let (layer, mut buf) = make_layer_and_buf();
445        let payload = vec![0x01u8, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08];
446        layer
447            .set_field(&mut buf, "payload", FieldValue::Bytes(payload.clone()))
448            .unwrap()
449            .unwrap();
450        assert_eq!(&buf[7..15], payload.as_slice());
451    }
452
453    #[test]
454    fn test_set_field_unknown() {
455        let (layer, mut buf) = make_layer_and_buf();
456        assert!(
457            layer
458                .set_field(&mut buf, "nonexistent", FieldValue::U8(0))
459                .is_none()
460        );
461    }
462
463    #[test]
464    fn test_field_names_dynamic() {
465        let (layer, _) = make_layer_and_buf();
466        let names = layer.field_names_dynamic();
467        assert_eq!(names, vec!["opcode", "length", "seq", "payload"]);
468    }
469
470    #[test]
471    fn test_layer_trait() {
472        let (layer, buf) = make_layer_and_buf();
473        assert_eq!(layer.kind(), LayerKind::Generic);
474        assert_eq!(Layer::header_len(&layer, &buf), 15);
475        assert!(Layer::field_names(&layer).is_empty());
476    }
477
478    #[test]
479    fn test_layer_with_offset() {
480        // Layer embedded inside a larger buffer at offset 4.
481        let descs = Arc::new(vec![GenericFieldDesc {
482            name: "val".to_string(),
483            offset: 0,
484            size: 2,
485            field_type: FieldType::U16,
486            default_value: vec![0, 0],
487        }]);
488        let buf = vec![0x00u8, 0x00, 0x00, 0x00, 0xAB, 0xCD]; // layer starts at byte 4
489        let index = LayerIndex::new(LayerKind::Generic, 4, 6);
490        let layer = GenericLayer::new(index, Arc::from("Offset"), descs);
491        let v = layer.get_field(&buf, "val").unwrap().unwrap();
492        assert_eq!(v, FieldValue::U16(0xABCD));
493    }
494
495    #[test]
496    fn test_bool_field() {
497        let descs = Arc::new(vec![GenericFieldDesc {
498            name: "flag".to_string(),
499            offset: 0,
500            size: 1,
501            field_type: FieldType::Bool,
502            default_value: vec![0x00],
503        }]);
504        let buf = vec![0x01u8];
505        let index = LayerIndex::new(LayerKind::Generic, 0, 1);
506        let layer = GenericLayer::new(index, Arc::from("Flags"), descs);
507        let v = layer.get_field(&buf, "flag").unwrap().unwrap();
508        assert_eq!(v, FieldValue::Bool(true));
509    }
510}