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#[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
22pub struct UdpServiceDb {
24 inner: HashMap<u16, UdpServiceEntry>,
25}
26
27impl UdpServiceDb {
28 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 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 #[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 pub fn get_name(&self, port: u16) -> Option<&str> {
63 self.inner.get(&port).map(|e| e.name.as_str())
64 }
65
66 pub fn get(&self, port: u16) -> Option<&UdpServiceEntry> {
68 self.inner.get(&port)
69 }
70
71 pub fn all(&self) -> impl Iterator<Item = (&u16, &UdpServiceEntry)> {
73 self.inner.iter()
74 }
75
76 pub fn wellknown(&self) -> impl Iterator<Item = (&u16, &UdpServiceEntry)> {
78 self.inner.iter().filter(|(_, e)| e.wellknown)
79 }
80
81 pub fn common(&self) -> impl Iterator<Item = (&u16, &UdpServiceEntry)> {
83 self.inner.iter().filter(|(_, e)| e.common)
84 }
85
86 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}