Skip to main content

tld/
lib.rs

1//! A native Rust library for Mozilla's Public Suffix List
2
3#![no_std]
4
5extern crate alloc;
6
7mod error;
8mod types;
9
10use alloc::{collections::BTreeMap, string::ToString, vec::Vec};
11use core::str::{FromStr, from_utf8};
12
13pub use error::Error;
14pub use types::{Domain, Info, List as Psl, Suffix, Type};
15
16/// The official URL of the list
17pub const LIST_URL: &str = "https://publicsuffix.org/list/public_suffix_list.dat";
18
19type Children = BTreeMap<Vec<u8>, Node>;
20
21const WILDCARD: &str = "*";
22
23const PUBLIC_SUFFIX_LIST_DATA: &str = include_str!("./public_suffix_list.txt");
24
25pub fn public_suffix_list() -> Result<List, Error> {
26    PUBLIC_SUFFIX_LIST_DATA.parse()
27}
28
29#[derive(Debug, Clone, Default, Eq, PartialEq)]
30struct Node {
31    children: Children,
32    leaf: Option<Leaf>,
33}
34
35#[derive(Debug, Clone, Copy, Eq, PartialEq)]
36struct Leaf {
37    is_exception: bool,
38    typ: Type,
39}
40
41/// A dynamic public suffix list
42#[derive(Debug, Clone, Default, Eq, PartialEq)]
43pub struct List {
44    rules: Node,
45    typ: Option<Type>,
46}
47
48impl List {
49    /// Creates a new list with default wildcard rule support
50    #[inline]
51    #[must_use]
52    pub fn new() -> Self {
53        Self::default()
54    }
55
56    /// Creates a new list from a byte slice
57    ///
58    /// # Errors
59    ///
60    /// Returns an `Err` if the list is not UTF-8 encoded
61    /// or if its format is invalid.
62    #[inline]
63    pub fn from_bytes(bytes: &[u8]) -> Result<Self, Error> {
64        from_utf8(bytes).map_err(|_| Error::ListNotUtf8Encoded)?.parse()
65    }
66
67    /// Checks to see if the list is empty, ignoring the wildcard rule
68    #[inline]
69    #[must_use]
70    pub fn is_empty(&self) -> bool {
71        self.rules.children.is_empty()
72    }
73
74    #[inline]
75    fn append(&mut self, mut rule: &str, typ: Type) -> Result<(), Error> {
76        let mut is_exception = false;
77        if rule.starts_with('!') {
78            if !rule.contains('.') {
79                return Err(Error::ExceptionAtFirstLabel(rule.to_string()));
80            }
81            is_exception = true;
82            rule = &rule[1..];
83        }
84
85        let mut current = &mut self.rules;
86        for label in rule.rsplit('.') {
87            if label.is_empty() {
88                return Err(Error::EmptyLabel(rule.to_string()));
89            }
90
91            #[cfg(not(feature = "anycase"))]
92            let key = label.as_bytes().to_vec();
93            #[cfg(feature = "anycase")]
94            let key = UniCase::new(Cow::from(label.to_string()));
95
96            current = current.children.entry(key).or_insert_with(Default::default);
97        }
98
99        current.leaf = Some(Leaf {
100            is_exception,
101            typ,
102        });
103
104        Ok(())
105    }
106}
107
108#[cfg(feature = "anycase")]
109macro_rules! anycase_key {
110    ($label:ident) => {
111        match from_utf8($label) {
112            Ok(label) => UniCase::new(Cow::from(label)),
113            Err(_) => {
114                return Info {
115                    len: 0,
116                    typ: None,
117                }
118            }
119        }
120    };
121}
122
123impl Psl for List {
124    #[inline]
125    fn find<'a, T>(&self, mut labels: T) -> Info
126    where
127        T: Iterator<Item = &'a [u8]>,
128    {
129        let mut rules = &self.rules;
130
131        // the first label
132        // it's special because we always need it whether or not
133        // it's in our hash map (because of the implicit wildcard)
134        let mut info = match labels.next() {
135            Some(label) => {
136                let mut info = Info {
137                    len: label.len(),
138                    typ: None,
139                };
140                #[cfg(not(feature = "anycase"))]
141                let node_opt = rules.children.get(label);
142                #[cfg(feature = "anycase")]
143                let node_opt = rules.children.get(&anycase_key!(label));
144                match node_opt {
145                    Some(node) => {
146                        info.typ = node.leaf.map(|leaf| leaf.typ);
147                        rules = node;
148                    }
149                    None => return info,
150                }
151                info
152            }
153            None => {
154                return Info {
155                    len: 0,
156                    typ: None,
157                };
158            }
159        };
160
161        // the rest of the labels
162        let mut len_so_far = info.len;
163        for label in labels {
164            #[cfg(not(feature = "anycase"))]
165            let node_opt = rules.children.get(label);
166            #[cfg(feature = "anycase")]
167            let node_opt = rules.children.get(&anycase_key!(label));
168            match node_opt {
169                Some(node) => rules = node,
170                None => {
171                    #[cfg(not(feature = "anycase"))]
172                    let node_opt = rules.children.get(WILDCARD.as_bytes());
173                    #[cfg(feature = "anycase")]
174                    let node_opt = rules.children.get(&UniCase::new(Cow::from(WILDCARD)));
175                    match node_opt {
176                        Some(node) => rules = node,
177                        None => break,
178                    }
179                }
180            }
181            let label_plus_dot = label.len() + 1;
182            if let Some(leaf) = rules.leaf {
183                if self.typ.is_none() || self.typ == Some(leaf.typ) {
184                    info.typ = Some(leaf.typ);
185                    if leaf.is_exception {
186                        info.len = len_so_far;
187                        break;
188                    }
189                    info.len = len_so_far + label_plus_dot;
190                }
191            }
192            len_so_far += label_plus_dot;
193        }
194
195        info
196    }
197}
198
199impl FromStr for List {
200    type Err = Error;
201
202    #[inline]
203    fn from_str(s: &str) -> Result<Self, Self::Err> {
204        let mut typ = None;
205        let mut list = List::new();
206        for line in s.lines() {
207            match line {
208                line if line.contains("BEGIN ICANN DOMAINS") => {
209                    typ = Some(Type::Icann);
210                }
211                line if line.contains("BEGIN PRIVATE DOMAINS") => {
212                    typ = Some(Type::Private);
213                }
214                line if line.starts_with("//") => {
215                    continue;
216                }
217                line => match typ {
218                    Some(typ) => {
219                        let rule = match line.split_whitespace().next() {
220                            Some(rule) => rule,
221                            None => continue,
222                        };
223                        list.append(rule, typ)?;
224                        #[cfg(feature = "punycode")]
225                        {
226                            let ascii = idna::domain_to_ascii(rule).map_err(|_| Error::InvalidRule(rule.to_owned()))?;
227                            list.append(&ascii, typ)?;
228                        }
229                    }
230                    None => {
231                        continue;
232                    }
233                },
234            }
235        }
236        if list.is_empty() {
237            return Err(Error::InvalidList);
238        }
239        Ok(list)
240    }
241}
242
243/// A list of only ICANN suffixes
244#[derive(Debug, Clone, Default, Eq, PartialEq)]
245pub struct IcannList(List);
246
247impl From<List> for IcannList {
248    #[inline]
249    fn from(mut list: List) -> Self {
250        list.typ = Some(Type::Icann);
251        Self(list)
252    }
253}
254
255impl From<IcannList> for List {
256    #[inline]
257    fn from(IcannList(mut list): IcannList) -> Self {
258        list.typ = None;
259        list
260    }
261}
262
263impl IcannList {
264    /// Creates a new list from a byte slice
265    ///
266    /// # Errors
267    ///
268    /// Returns an `Err` if the list is not UTF-8 encoded
269    /// or if its format is invalid.
270    #[inline]
271    pub fn from_bytes(bytes: &[u8]) -> Result<Self, Error> {
272        let list = List::from_bytes(bytes)?;
273        Ok(list.into())
274    }
275
276    /// Checks to see if the list is empty, ignoring the wildcard rule
277    #[inline]
278    #[must_use]
279    pub fn is_empty(&self) -> bool {
280        self.0.is_empty()
281    }
282}
283
284impl FromStr for IcannList {
285    type Err = Error;
286
287    #[inline]
288    fn from_str(s: &str) -> Result<Self, Self::Err> {
289        let list = List::from_str(s)?;
290        Ok(list.into())
291    }
292}
293
294impl Psl for IcannList {
295    #[inline]
296    fn find<'a, T>(&self, labels: T) -> Info
297    where
298        T: Iterator<Item = &'a [u8]>,
299    {
300        self.0.find(labels)
301    }
302}
303
304/// A list of only private suffixes
305#[derive(Debug, Clone, Default, Eq, PartialEq)]
306pub struct PrivateList(List);
307
308impl From<List> for PrivateList {
309    #[inline]
310    fn from(mut list: List) -> Self {
311        list.typ = Some(Type::Private);
312        Self(list)
313    }
314}
315
316impl From<PrivateList> for List {
317    #[inline]
318    fn from(PrivateList(mut list): PrivateList) -> Self {
319        list.typ = None;
320        list
321    }
322}
323
324impl PrivateList {
325    /// Creates a new list from a byte slice
326    ///
327    /// # Errors
328    ///
329    /// Returns an `Err` if the list is not UTF-8 encoded
330    /// or if its format is invalid.
331    #[inline]
332    pub fn from_bytes(bytes: &[u8]) -> Result<Self, Error> {
333        let list = List::from_bytes(bytes)?;
334        Ok(list.into())
335    }
336
337    /// Checks to see if the list is empty, ignoring the wildcard rule
338    #[inline]
339    #[must_use]
340    pub fn is_empty(&self) -> bool {
341        self.0.is_empty()
342    }
343}
344
345impl FromStr for PrivateList {
346    type Err = Error;
347
348    #[inline]
349    fn from_str(s: &str) -> Result<Self, Self::Err> {
350        let list = List::from_str(s)?;
351        Ok(list.into())
352    }
353}
354
355impl Psl for PrivateList {
356    #[inline]
357    fn find<'a, T>(&self, labels: T) -> Info
358    where
359        T: Iterator<Item = &'a [u8]>,
360    {
361        self.0.find(labels)
362    }
363}
364
365#[cfg(test)]
366mod tests {
367    use super::*;
368
369    const LIST: &[u8] = b"
370        // BEGIN ICANN DOMAINS
371        com.uk
372        ";
373
374    #[test]
375    fn list_construction() {
376        let list = List::from_bytes(LIST).unwrap();
377        let expected = List {
378            typ: None,
379            rules: Node {
380                children: {
381                    let mut children = Children::default();
382                    children.insert(
383                        #[cfg(not(feature = "anycase"))]
384                        b"uk".to_vec(),
385                        #[cfg(feature = "anycase")]
386                        UniCase::new(Cow::from("uk")),
387                        Node {
388                            children: {
389                                let mut children = Children::default();
390                                children.insert(
391                                    #[cfg(not(feature = "anycase"))]
392                                    b"com".to_vec(),
393                                    #[cfg(feature = "anycase")]
394                                    UniCase::new(Cow::from("com")),
395                                    Node {
396                                        children: Default::default(),
397                                        leaf: Some(Leaf {
398                                            is_exception: false,
399                                            typ: Type::Icann,
400                                        }),
401                                    },
402                                );
403                                children
404                            },
405                            leaf: None,
406                        },
407                    );
408                    children
409                },
410                leaf: None,
411            },
412        };
413        assert_eq!(list, expected);
414    }
415
416    #[test]
417    fn find_localhost() {
418        let list = List::from_bytes(LIST).unwrap();
419        let labels = b"localhost".rsplit(|x| *x == b'.');
420        assert_eq!(
421            list.find(labels),
422            Info {
423                len: 9,
424                typ: None
425            }
426        );
427    }
428
429    #[test]
430    fn find_uk() {
431        let list = List::from_bytes(LIST).unwrap();
432        let labels = b"uk".rsplit(|x| *x == b'.');
433        assert_eq!(
434            list.find(labels),
435            Info {
436                len: 2,
437                typ: None
438            }
439        );
440    }
441
442    #[test]
443    fn find_com_uk() {
444        let list = List::from_bytes(LIST).unwrap();
445        let labels = b"com.uk".rsplit(|x| *x == b'.');
446        assert_eq!(
447            list.find(labels),
448            Info {
449                len: 6,
450                typ: Some(Type::Icann)
451            }
452        );
453    }
454
455    #[test]
456    fn find_ide_kyoto_jp() {
457        let list = List::from_bytes(b"// BEGIN ICANN DOMAINS\nide.kyoto.jp").unwrap();
458        let labels = b"ide.kyoto.jp".rsplit(|x| *x == b'.');
459        assert_eq!(
460            list.find(labels),
461            Info {
462                len: 12,
463                typ: Some(Type::Icann)
464            }
465        );
466    }
467}