405 lines
14 KiB
Rust
405 lines
14 KiB
Rust
//! MaxMind DB (`.mmdb`) country lookup, with the `G:<cc>` geoban extban, the
|
||
//! `GEOIP` command and a WHOIS country line. The binary format is parsed by hand:
|
||
//! the metadata section, the record-size-aware search tree, and the typed data
|
||
//! decoder.
|
||
//!
|
||
//! Config: `geoip_database = /path/to/GeoLite2-Country.mmdb` (loaded once at boot).
|
||
|
||
use std::net::IpAddr;
|
||
use std::sync::Arc;
|
||
|
||
use crate::command::{CmdResult, Command};
|
||
use crate::numeric::ERR_NOPRIVILEGES;
|
||
use crate::server::Server;
|
||
use crate::Uid;
|
||
|
||
const MARKER: &[u8] = b"\xab\xcd\xefMaxMind.com";
|
||
const SEPARATOR: usize = 16;
|
||
|
||
/// A loaded MaxMind DB, cached in `Server.ext`.
|
||
pub struct GeoDb(pub Arc<Mmdb>);
|
||
|
||
/// A resolved country: the ISO 3166-1 alpha-2 code (for the `G:` geoban) and the
|
||
/// full English name (for display).
|
||
pub struct Country {
|
||
pub iso: String,
|
||
pub name: String,
|
||
}
|
||
|
||
/// A parsed `.mmdb` file: the raw bytes plus the tree geometry from its metadata.
|
||
pub struct Mmdb {
|
||
data: Vec<u8>,
|
||
node_count: usize,
|
||
node_byte_size: usize, // record_size * 2 / 8
|
||
record_size: u32,
|
||
data_start: usize, // node_count*node_byte_size + 16 (first byte after the separator)
|
||
ip_version: u16,
|
||
}
|
||
|
||
/// Big-endian value of up to 8 bytes.
|
||
fn be(bytes: &[u8]) -> u64 {
|
||
bytes.iter().fold(0u64, |acc, &b| (acc << 8) | b as u64)
|
||
}
|
||
|
||
/// A cursor over the data section for pointer-aware decoding. `base` is the offset
|
||
/// pointers are relative to (the start of the section being decoded).
|
||
struct Decoder<'a> {
|
||
data: &'a [u8],
|
||
base: usize,
|
||
}
|
||
|
||
impl<'a> Decoder<'a> {
|
||
/// Parse the control byte at `off`: returns `(type, size, payload_offset)`.
|
||
/// Not used for pointers (they encode size differently — see `resolve`).
|
||
fn control(&self, off: usize) -> Option<(u8, usize, usize)> {
|
||
let b = *self.data.get(off)?;
|
||
let mut typ = b >> 5;
|
||
let mut o = off + 1;
|
||
if typ == 0 {
|
||
typ = 7 + *self.data.get(o)?;
|
||
o += 1;
|
||
}
|
||
let mut size = (b & 0x1f) as usize;
|
||
if size >= 29 {
|
||
match size {
|
||
29 => {
|
||
size = 29 + *self.data.get(o)? as usize;
|
||
o += 1;
|
||
}
|
||
30 => {
|
||
size = 285 + be(self.data.get(o..o + 2)?) as usize;
|
||
o += 2;
|
||
}
|
||
_ => {
|
||
size = 65821 + be(self.data.get(o..o + 3)?) as usize;
|
||
o += 3;
|
||
}
|
||
}
|
||
}
|
||
Some((typ, size, o))
|
||
}
|
||
|
||
/// Follow pointer(s) at `off` to the concrete value offset (non-pointers pass
|
||
/// through). Depth-limited against malformed data.
|
||
fn resolve(&self, off: usize, depth: u8) -> Option<usize> {
|
||
if depth > 8 {
|
||
return None;
|
||
}
|
||
let b = *self.data.get(off)?;
|
||
if b >> 5 != 1 {
|
||
return Some(off);
|
||
}
|
||
let ps = ((b >> 3) & 0x3) as usize;
|
||
let v0 = (b & 0x7) as usize;
|
||
let ptr = match ps {
|
||
0 => (v0 << 8) | *self.data.get(off + 1)? as usize,
|
||
1 => ((v0 << 16) | be(self.data.get(off + 1..off + 3)?) as usize) + 2048,
|
||
2 => ((v0 << 24) | be(self.data.get(off + 1..off + 4)?) as usize) + 526336,
|
||
_ => be(self.data.get(off + 1..off + 5)?) as usize,
|
||
};
|
||
self.resolve(self.base + ptr, depth + 1)
|
||
}
|
||
|
||
/// The number of bytes the value at `off` occupies (pointers are 2–5 bytes; not
|
||
/// followed). Recurses into maps/arrays. Returns 0 on malformed input.
|
||
fn value_len(&self, off: usize) -> usize {
|
||
let b = match self.data.get(off) {
|
||
Some(&b) => b,
|
||
None => return 0,
|
||
};
|
||
if b >> 5 == 1 {
|
||
return 1 + ((b >> 3) & 0x3) as usize + 1; // 2..=5 bytes
|
||
}
|
||
let Some((typ, size, payload)) = self.control(off) else {
|
||
return 0;
|
||
};
|
||
let header = payload - off;
|
||
match typ {
|
||
7 => {
|
||
// map: `size` key/value pairs
|
||
let mut cur = payload;
|
||
for _ in 0..size {
|
||
cur += self.value_len(cur); // key
|
||
cur += self.value_len(cur); // value
|
||
}
|
||
cur - off
|
||
}
|
||
11 => {
|
||
// array: `size` elements
|
||
let mut cur = payload;
|
||
for _ in 0..size {
|
||
cur += self.value_len(cur);
|
||
}
|
||
cur - off
|
||
}
|
||
14 => header, // bool: value is in the size field, no payload
|
||
_ => header + size, // string/bytes/ints/float/double
|
||
}
|
||
}
|
||
|
||
/// Read a UTF-8 string at `off` (following pointers).
|
||
fn string(&self, off: usize) -> Option<String> {
|
||
let off = self.resolve(off, 0)?;
|
||
let (typ, size, payload) = self.control(off)?;
|
||
if typ != 2 {
|
||
return None;
|
||
}
|
||
std::str::from_utf8(self.data.get(payload..payload + size)?)
|
||
.ok()
|
||
.map(str::to_string)
|
||
}
|
||
|
||
/// Read an unsigned integer at `off` (following pointers).
|
||
fn uint(&self, off: usize) -> Option<u64> {
|
||
let off = self.resolve(off, 0)?;
|
||
let (typ, size, payload) = self.control(off)?;
|
||
if !matches!(typ, 5 | 6 | 9 | 10) {
|
||
return None;
|
||
}
|
||
Some(be(self.data.get(payload..payload + size)?))
|
||
}
|
||
|
||
/// The value offset for `key` in the map at `off` (following pointers).
|
||
fn map_get(&self, off: usize, key: &str) -> Option<usize> {
|
||
let off = self.resolve(off, 0)?;
|
||
let (typ, size, payload) = self.control(off)?;
|
||
if typ != 7 {
|
||
return None;
|
||
}
|
||
let mut cur = payload;
|
||
for _ in 0..size {
|
||
let k = self.string(cur)?;
|
||
cur += self.value_len(cur);
|
||
if k == key {
|
||
return Some(cur);
|
||
}
|
||
cur += self.value_len(cur);
|
||
}
|
||
None
|
||
}
|
||
}
|
||
|
||
impl Mmdb {
|
||
/// Load and parse an `.mmdb` file.
|
||
pub fn open(path: &str) -> Option<Mmdb> {
|
||
let data = std::fs::read(path).ok()?;
|
||
let marker = data.windows(MARKER.len()).rposition(|w| w == MARKER)?;
|
||
let meta = marker + MARKER.len();
|
||
let d = Decoder { data: &data, base: meta };
|
||
let node_count = d.uint(d.map_get(meta, "node_count")?)? as usize;
|
||
let record_size = d.uint(d.map_get(meta, "record_size")?)? as u32;
|
||
let ip_version = d.uint(d.map_get(meta, "ip_version")?)? as u16;
|
||
if !matches!(record_size, 24 | 28 | 32) || node_count == 0 {
|
||
return None;
|
||
}
|
||
let node_byte_size = record_size as usize * 2 / 8;
|
||
let data_start = node_count * node_byte_size + SEPARATOR;
|
||
Some(Mmdb {
|
||
data,
|
||
node_count,
|
||
node_byte_size,
|
||
record_size,
|
||
data_start,
|
||
ip_version,
|
||
})
|
||
}
|
||
|
||
/// The left (bit=false) or right (bit=true) record of tree `node`.
|
||
fn record(&self, node: usize, bit: bool) -> Option<usize> {
|
||
let base = node * self.node_byte_size;
|
||
match self.record_size {
|
||
24 => {
|
||
let o = base + if bit { 3 } else { 0 };
|
||
Some(be(self.data.get(o..o + 3)?) as usize)
|
||
}
|
||
28 => {
|
||
let mid = *self.data.get(base + 3)?;
|
||
if bit {
|
||
Some(((mid as usize & 0x0f) << 24) | be(self.data.get(base + 4..base + 7)?) as usize)
|
||
} else {
|
||
Some(((mid as usize >> 4) << 24) | be(self.data.get(base..base + 3)?) as usize)
|
||
}
|
||
}
|
||
_ => {
|
||
let o = base + if bit { 4 } else { 0 };
|
||
Some(be(self.data.get(o..o + 4)?) as usize)
|
||
}
|
||
}
|
||
}
|
||
|
||
/// The country for `ip` — ISO code plus English name — if the database has one.
|
||
pub fn country(&self, ip: IpAddr) -> Option<Country> {
|
||
// build the bit path; IPv4 in an IPv6 db is prefixed with 96 zero bits
|
||
let mut bits: Vec<bool> = Vec::with_capacity(128);
|
||
match ip {
|
||
IpAddr::V4(a) => {
|
||
if self.ip_version == 6 {
|
||
bits.extend(std::iter::repeat(false).take(96));
|
||
}
|
||
for byte in a.octets() {
|
||
for i in (0..8).rev() {
|
||
bits.push((byte >> i) & 1 == 1);
|
||
}
|
||
}
|
||
}
|
||
IpAddr::V6(a) => {
|
||
for byte in a.octets() {
|
||
for i in (0..8).rev() {
|
||
bits.push((byte >> i) & 1 == 1);
|
||
}
|
||
}
|
||
}
|
||
}
|
||
let mut node = 0usize;
|
||
for bit in bits {
|
||
if node >= self.node_count {
|
||
return None;
|
||
}
|
||
let rec = self.record(node, bit)?;
|
||
if rec == self.node_count {
|
||
return None; // no data
|
||
}
|
||
if rec > self.node_count {
|
||
// data pointer: abs = tree_size + (rec - node_count)
|
||
let abs = (self.data_start - SEPARATOR) + rec - self.node_count;
|
||
let d = Decoder {
|
||
data: &self.data,
|
||
base: self.data_start,
|
||
};
|
||
let country = d.map_get(abs, "country")?;
|
||
let iso = d.string(d.map_get(country, "iso_code")?)?;
|
||
// country.names.en — the `names` submap is usually a shared pointer;
|
||
// map_get/string follow it. Fall back to the code if it's absent.
|
||
let name = d
|
||
.map_get(country, "names")
|
||
.and_then(|names| d.map_get(names, "en"))
|
||
.and_then(|en| d.string(en))
|
||
.unwrap_or_else(|| iso.clone());
|
||
return Some(Country { iso, name });
|
||
}
|
||
node = rec;
|
||
}
|
||
None
|
||
}
|
||
}
|
||
|
||
/// Load the configured database into `Server.ext` at boot. Called from `Ircd::new`.
|
||
pub fn init(s: &mut Server) {
|
||
let Some(path) = s.conf("geoip_database").map(str::to_string) else {
|
||
return;
|
||
};
|
||
match Mmdb::open(&path) {
|
||
Some(db) => {
|
||
s.ext.set(GeoDb(Arc::new(db)));
|
||
eprintln!("echoircd: loaded GeoIP database {path}");
|
||
}
|
||
None => eprintln!("echoircd: could not read GeoIP database {path}"),
|
||
}
|
||
}
|
||
|
||
/// The country of `ip` per the loaded database (ISO code uppercased).
|
||
pub fn lookup(s: &Server, ip: IpAddr) -> Option<Country> {
|
||
s.ext.get::<GeoDb>().and_then(|db| db.0.country(ip)).map(|c| Country {
|
||
iso: c.iso.to_ascii_uppercase(),
|
||
name: c.name,
|
||
})
|
||
}
|
||
|
||
/// The `G:<cc>` geoban match: does `uid`'s country code equal (case-insensitively)
|
||
/// one of the codes in the extban? Dispatched from `Server::ban_list_hit`.
|
||
pub fn geoban_match(s: &Server, uid: Uid, spec: &str) -> bool {
|
||
let Some(ip) = s.users.get(&uid).map(|u| u.addr.ip()) else {
|
||
return false;
|
||
};
|
||
match lookup(s, ip) {
|
||
Some(c) => spec
|
||
.split(',')
|
||
.any(|want| want.trim().eq_ignore_ascii_case(&c.iso)),
|
||
None => false,
|
||
}
|
||
}
|
||
|
||
/// A WHOIS line (opers only) naming the target's country.
|
||
pub fn whois_line(s: &Server, tuid: Uid) -> Option<String> {
|
||
let ip = s.users.get(&tuid).map(|u| u.addr.ip())?;
|
||
lookup(s, ip).map(|c| format!("is connecting from country {}", c.name))
|
||
}
|
||
|
||
pub fn commands() -> Vec<Box<dyn Command>> {
|
||
vec![Box::new(GeoIpCmd)]
|
||
}
|
||
|
||
/// GEOIP `<nick|ip>` — oper command reporting the country of a user or raw IP.
|
||
struct GeoIpCmd;
|
||
impl Command for GeoIpCmd {
|
||
fn name(&self) -> &'static str {
|
||
"GEOIP"
|
||
}
|
||
fn min_params(&self) -> usize {
|
||
1
|
||
}
|
||
fn handle(&self, s: &mut Server, uid: Uid, params: &[String]) -> CmdResult {
|
||
if !s.is_oper(uid) {
|
||
s.numeric(
|
||
uid,
|
||
ERR_NOPRIVILEGES,
|
||
":Permission Denied- You're not an IRC operator",
|
||
);
|
||
return CmdResult::Fail;
|
||
}
|
||
let target = ¶ms[0];
|
||
let ip = s
|
||
.find_nick(target)
|
||
.and_then(|t| s.users.get(&t))
|
||
.map(|u| u.addr.ip())
|
||
.or_else(|| target.parse::<IpAddr>().ok());
|
||
let nick = s.users.get(&uid).map(|u| u.nick.clone()).unwrap_or_default();
|
||
let msg = match ip {
|
||
None => format!("GEOIP: no such nick, and {target} is not an IP"),
|
||
Some(ip) => match lookup(s, ip) {
|
||
Some(c) => format!("GEOIP: {target} ({ip}) is in {} ({})", c.name, c.iso),
|
||
None => format!("GEOIP: no country found for {target} ({ip})"),
|
||
},
|
||
};
|
||
s.send(uid, format!(":{} NOTICE {nick} :*** {msg}", s.name));
|
||
CmdResult::Ok
|
||
}
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
use std::net::Ipv4Addr;
|
||
|
||
// Use a database from the env var if set, else common locations; the test
|
||
// no-ops when none is present so CI stays green.
|
||
const DB_CANDIDATES: &[&str] = &[
|
||
"/usr/share/GeoIP/GeoLite2-Country.mmdb",
|
||
"/etc/echoircd/GeoLite2-Country.mmdb",
|
||
];
|
||
|
||
fn load() -> Option<Mmdb> {
|
||
std::env::var("ECHOIRCD_TEST_MMDB")
|
||
.ok()
|
||
.and_then(|p| Mmdb::open(&p))
|
||
.or_else(|| DB_CANDIDATES.iter().find_map(|p| Mmdb::open(p)))
|
||
}
|
||
|
||
#[test]
|
||
fn known_ips_resolve() {
|
||
let Some(db) = load() else {
|
||
return;
|
||
};
|
||
// 8.8.8.8 (Google DNS) is US, "United States", in every GeoLite2 vintage.
|
||
let us = db.country(IpAddr::V4(Ipv4Addr::new(8, 8, 8, 8))).unwrap();
|
||
assert_eq!(us.iso, "US");
|
||
assert_eq!(us.name, "United States");
|
||
// A private address has no country record.
|
||
assert!(db.country(IpAddr::V4(Ipv4Addr::new(10, 0, 0, 1))).is_none());
|
||
// IPv6 traversal (Google public DNS) also resolves to US.
|
||
assert_eq!(
|
||
db.country("2001:4860:4860::8888".parse().unwrap()).unwrap().iso,
|
||
"US"
|
||
);
|
||
}
|
||
}
|