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