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