ndb_udp_service/
lib.rs

1use anyhow::Result;
2use ndb_core::utils::serde::de_u8_to_bool;
3use serde::{Deserialize, Serialize};
4use std::collections::HashMap;
5use std::io::Read;
6
7pub const CSV_NAME: &str = "udp-services.csv";
8pub const BIN_NAME: &str = "udp-services.bin";
9
10/// Represents a single UDP service entry
11#[derive(Debug, Clone, Serialize, Deserialize)]
12pub struct UdpServiceEntry {
13    pub port: u16,
14    pub name: String,
15    pub description: Option<String>,
16    #[serde(deserialize_with = "de_u8_to_bool")]
17    pub wellknown: bool,
18    #[serde(deserialize_with = "de_u8_to_bool")]
19    pub common: bool,
20}
21
22/// Represents the UDP service database
23pub struct UdpServiceDb {
24    inner: HashMap<u16, UdpServiceEntry>,
25}
26
27impl UdpServiceDb {
28    /// Load database from a CSV reader
29    pub fn from_csv<R: Read>(reader: R) -> Result<Self, csv::Error> {
30        let mut rdr = csv::Reader::from_reader(reader);
31        let mut map = HashMap::new();
32        for result in rdr.deserialize::<UdpServiceEntry>() {
33            let entry = result?;
34            map.insert(entry.port, entry);
35        }
36        Ok(Self { inner: map })
37    }
38
39    pub fn from_entries(entries: Vec<UdpServiceEntry>) -> Self {
40        let inner = entries
41            .into_iter()
42            .map(|entry| (entry.port, entry))
43            .collect();
44        Self { inner }
45    }
46
47    /// Create a new UDP service database from a binary slice
48    pub fn from_slice(slice: &[u8]) -> Result<Self> {
49        let (entries, _): (Vec<UdpServiceEntry>, _) =
50            bincode::serde::decode_from_slice(slice, bincode::config::standard())?;
51        Ok(Self::from_entries(entries))
52    }
53
54    /// Load embedded (bundled) database
55    #[cfg(feature = "bundled")]
56    pub fn bundled() -> Self {
57        static BIN_DATA: &[u8] = include_bytes!("../data/udp-services.bin");
58        Self::from_slice(BIN_DATA).expect("Failed to load bundled udp-services.bin")
59    }
60
61    /// Lookup a UDP service name by port
62    pub fn get_name(&self, port: u16) -> Option<&str> {
63        self.inner.get(&port).map(|e| e.name.as_str())
64    }
65
66    /// Lookup a UDP service entry by port
67    pub fn get(&self, port: u16) -> Option<&UdpServiceEntry> {
68        self.inner.get(&port)
69    }
70
71    /// Get all UDP service entries as an iterator
72    pub fn all(&self) -> impl Iterator<Item = (&u16, &UdpServiceEntry)> {
73        self.inner.iter()
74    }
75
76    /// Get well-known UDP service entries
77    pub fn wellknown(&self) -> impl Iterator<Item = (&u16, &UdpServiceEntry)> {
78        self.inner.iter().filter(|(_, e)| e.wellknown)
79    }
80
81    /// Get common UDP service entries
82    pub fn common(&self) -> impl Iterator<Item = (&u16, &UdpServiceEntry)> {
83        self.inner.iter().filter(|(_, e)| e.common)
84    }
85
86    /// Get all UDP service entries as a vector
87    pub fn entries(&self) -> Vec<UdpServiceEntry> {
88        self.inner.values().cloned().collect()
89    }
90}
91
92#[cfg(test)]
93mod tests {
94    use super::*;
95
96    #[test]
97    #[cfg(feature = "bundled")]
98    fn test_lookup_known_port() {
99        let db = UdpServiceDb::bundled();
100        let entry = db.get(53).expect("Port 53 should exist");
101        assert_eq!(entry.name, "domain");
102        assert!(entry.wellknown);
103        assert!(entry.common);
104    }
105
106    #[test]
107    #[cfg(feature = "bundled")]
108    fn test_lookup_unknown_port() {
109        let db = UdpServiceDb::bundled();
110        assert!(db.get(u16::MAX).is_none());
111    }
112
113    #[test]
114    #[cfg(feature = "bundled")]
115    fn test_get_name() {
116        let db = UdpServiceDb::bundled();
117        assert_eq!(db.get_name(443), Some("https"));
118    }
119
120    #[test]
121    #[cfg(feature = "bundled")]
122    fn test_iter_wellknown() {
123        let db = UdpServiceDb::bundled();
124        assert!(db.wellknown().any(|(port, _)| *port == 69));
125    }
126
127    #[test]
128    #[cfg(feature = "bundled")]
129    fn test_iter_common() {
130        let db = UdpServiceDb::bundled();
131        assert!(db.common().any(|(port, _)| *port == 53));
132    }
133}