ndb_tcp_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 = "tcp-services.csv";
8pub const BIN_NAME: &str = "tcp-services.bin";
9
10/// Represents a single TCP service entry
11#[derive(Debug, Clone, Serialize, Deserialize)]
12pub struct TcpServiceEntry {
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 TCP service database
23pub struct TcpServiceDb {
24    inner: HashMap<u16, TcpServiceEntry>,
25}
26
27impl TcpServiceDb {
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::<TcpServiceEntry>() {
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<TcpServiceEntry>) -> 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 TCP service database from a binary slice
48    pub fn from_slice(slice: &[u8]) -> Result<Self> {
49        let (entries, _): (Vec<TcpServiceEntry>, _) =
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/tcp-services.bin");
58        Self::from_slice(BIN_DATA).expect("Failed to load bundled tcp-services.bin")
59    }
60
61    /// Lookup a TCP 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 TCP service entry by port
67    pub fn get(&self, port: u16) -> Option<&TcpServiceEntry> {
68        self.inner.get(&port)
69    }
70
71    /// Get all TCP service entries as an iterator
72    pub fn all(&self) -> impl Iterator<Item = (&u16, &TcpServiceEntry)> {
73        self.inner.iter()
74    }
75
76    /// Get well-known TCP service entries
77    pub fn wellknown(&self) -> impl Iterator<Item = (&u16, &TcpServiceEntry)> {
78        self.inner.iter().filter(|(_, e)| e.wellknown)
79    }
80
81    /// Get common TCP service entries
82    pub fn common(&self) -> impl Iterator<Item = (&u16, &TcpServiceEntry)> {
83        self.inner.iter().filter(|(_, e)| e.common)
84    }
85
86    /// Get all TCP service entries as a vector
87    pub fn entries(&self) -> Vec<TcpServiceEntry> {
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 = TcpServiceDb::bundled();
100        let entry = db.get(80).expect("Port 80 should exist");
101        assert_eq!(entry.name, "http");
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 = TcpServiceDb::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 = TcpServiceDb::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 = TcpServiceDb::bundled();
124        assert!(db.wellknown().any(|(port, _)| *port == 22));
125    }
126
127    #[test]
128    #[cfg(feature = "bundled")]
129    fn test_iter_common() {
130        let db = TcpServiceDb::bundled();
131        assert!(db.common().any(|(port, _)| *port == 53));
132    }
133}
ndb_tcp_service/lib.rs

ndb_tcp_service/
lib.rs
