1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143

use ipnet::{IpNet, Ipv4Net, Ipv6Net};
use iprange::IpRange;
use std::{borrow::Borrow, net::IpAddr};

const NUM_ALPHABET: usize = 26;
const NUM_DIGIT: usize = 10;
const NUM_SPECIAL: usize = 2; // `.`, `-`
const NUM_CHILDREN: usize = NUM_ALPHABET + NUM_DIGIT + NUM_SPECIAL;
use std::fmt;
use std::usize;

const MATCHED: usize = usize::MAX;
const NOT_MATCHED: usize = usize::MAX - 1;

pub struct Pattern {
  // usize::MAX -> matched
  // usize::MAX-1 -> Not matched
  // other -> index of first child, should be <= self.0.len() - NUM_CHILDREN
  // should not be empty
  host_trie: Vec<usize>,
  ipv4: IpRange<Ipv4Net>,
  ipv6: IpRange<Ipv6Net>,
}

#[derive(Debug)]
pub enum Error {
  UnexpectedCharacter(u8),
}

impl fmt::Display for Error {
  fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
    match self {
      Error::UnexpectedCharacter(ch) => write!(f, "unknown character: '{:?}'", *ch as char),
    }
  }
}

impl std::error::Error for Error {}

pub type Result<T> = std::result::Result<T, Error>;

impl Default for Pattern {
  fn default() -> Self {
    Pattern {
      host_trie: vec![NOT_MATCHED],
      ipv4: Default::default(),
      ipv6: Default::default(),
    }
  }
}
impl Pattern {
  fn codec(ch: u8) -> Result<usize> {
    if b'A' <= ch && ch <= b'Z' {
      Ok((ch - b'A') as usize)
    } else if b'a' <= ch && ch <= b'z' {
      Ok((ch - b'a') as usize)
    } else if b'0' <= ch && ch <= b'9' {
      Ok((ch - b'0') as usize + NUM_ALPHABET)
    } else if ch == b'.' {
      Ok(NUM_ALPHABET + NUM_DIGIT)
    } else if ch == b'-' {
      Ok(NUM_ALPHABET + NUM_DIGIT + 1)
    } else {
      Err(Error::UnexpectedCharacter(ch))
    }
  }

  fn add_host(&mut self, suffix: &[u8]) -> Result<()> {
    let mut current = 0;
    for &b in suffix.iter().rev() {
      let child = match self.host_trie[current] {
        MATCHED => return Ok(()),
        NOT_MATCHED => {
          let child = self.host_trie.len();
          self
            .host_trie
            .extend(std::iter::repeat(NOT_MATCHED).take(NUM_CHILDREN));
          self.host_trie[current] = child;
          child
        }
        child => child,
      };
      current = child + Self::codec(b)?;
    }
    self.host_trie[current] = MATCHED;
    Ok(())
  }

  pub fn add(&mut self, host_or_ip: &str) -> Result<()> {
    if let Ok(ipnet) = host_or_ip.parse::<IpNet>() {
      match ipnet {
        IpNet::V4(ipnet) => {
          self.ipv4.add(ipnet);
        }
        IpNet::V6(ipnet) => {
          self.ipv6.add(ipnet);
        }
      }
      Ok(())
    } else {
      self.add_host(host_or_ip.as_bytes())
    }
  }

  pub fn build(&mut self) {
    self.host_trie.shrink_to_fit();
    self.ipv4.simplify();
    self.ipv6.simplify();
  }

  pub fn contain_host(&self, uri: &[u8]) -> bool {
    let mut current = 0usize;
    for &b in uri.iter().rev() {
      match self.host_trie[current] {
        MATCHED => return true,
        NOT_MATCHED => return false,
        child => match Self::codec(b) {
          Ok(n) => current = child + n,
          Err(_) => return false,
        },
      }
    }
    self.host_trie[current] == MATCHED
  }

  pub fn contain_ip(&self, ip: &IpAddr) -> bool {
    match ip {
      IpAddr::V4(ip) => self.ipv4.contains(ip),
      IpAddr::V6(ip) => self.ipv6.contains(ip),
    }
  }

  pub fn from_strs<'a>(iter: impl Iterator<Item = &'a str>) -> Result<Pattern> {
    let mut pattern = Pattern::default();

    for s in iter {
      pattern.add(s.borrow())?
    }
    pattern.build();

    Ok(pattern)
  }
}