runlight/
mmdb.rs

1//! A reader for MaxMind DB files (the MMDB format that MaxMind's GeoLite2
2//! and DB-IP's free databases use), so location needs no other crate. It
3//! answers what the TypeScript server's mmdb-lib answers: the record for an
4//! address as a JSON value, or `None` when the address is not in the
5//! database.
6//!
7//! A database opened from a file is read a page at a time as lookups need
8//! it, so a 130 MB city database costs each lookup a few hundred kilobytes
9//! of reads.
10//!
11//! Format: <https://maxmind.github.io/MaxMind-DB/>
12
13use std::collections::HashMap;
14use std::fs::File;
15use std::io::{Read, Seek, SeekFrom};
16use std::net::IpAddr;
17use std::sync::Mutex;
18
19use crate::geo::{Found, GeoLookup};
20use crate::js::{self, Object, Value};
21use crate::{BoxError, BoxFuture};
22
23const METADATA_MARKER: &[u8] = b"\xAB\xCD\xEFMaxMind.com";
24/// The metadata sits in the file's last 128 KiB.
25const METADATA_MAX: u64 = 131_072;
26const PAGE: u64 = 4096;
27/// Pages kept from a file at once; a lookup reads a few dozen.
28const PAGES_KEPT: usize = 256;
29
30enum Source {
31    Bytes(Vec<u8>),
32    File(Mutex<(File, HashMap<u64, Vec<u8>>)>),
33}
34
35/// An MMDB database.
36pub struct Mmdb {
37    source: Source,
38    size: u64,
39    /// The database's metadata.
40    pub metadata: Value,
41    node_count: u64,
42    record_size: u64,
43    node_bytes: u64,
44    data_start: u64,
45    ipv4_start: Mutex<Option<u64>>,
46}
47
48/// A database that could not be read.
49#[derive(Debug)]
50pub struct MmdbError(pub String);
51
52impl std::fmt::Display for MmdbError {
53    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
54        f.write_str(&self.0)
55    }
56}
57
58impl std::error::Error for MmdbError {}
59
60fn bad(what: &str) -> MmdbError {
61    MmdbError(what.to_string())
62}
63
64impl Mmdb {
65    /// A database held in memory.
66    pub fn from_bytes(bytes: Vec<u8>) -> Result<Mmdb, MmdbError> {
67        let size = bytes.len() as u64;
68        Mmdb::with(Source::Bytes(bytes), size)
69    }
70
71    /// A database read from its file as lookups need it.
72    pub fn open(path: impl AsRef<std::path::Path>) -> Result<Mmdb, MmdbError> {
73        let path = path.as_ref();
74        let file = File::open(path).map_err(|_| MmdbError(format!("Could not read {}", path.display())))?;
75        let size = file.metadata().map_err(|e| MmdbError(e.to_string()))?.len();
76        Mmdb::with(Source::File(Mutex::new((file, HashMap::new()))), size)
77    }
78
79    fn with(source: Source, size: u64) -> Result<Mmdb, MmdbError> {
80        let mut db = Mmdb {
81            source,
82            size,
83            metadata: Value::Null,
84            node_count: 0,
85            record_size: 0,
86            node_bytes: 0,
87            data_start: 0,
88            ipv4_start: Mutex::new(None),
89        };
90        let tail_start = size.saturating_sub(METADATA_MAX);
91        let tail = db.read(tail_start, size - tail_start)?;
92        let at = tail
93            .windows(METADATA_MARKER.len())
94            .rposition(|w| w == METADATA_MARKER)
95            .ok_or_else(|| bad("Not a MaxMind DB file: no metadata"))?;
96        let start = tail_start + at as u64 + METADATA_MARKER.len() as u64;
97        let (metadata, _) = db.decode(start, start, 0)?;
98        let num = |k: &str| metadata.get(k).and_then(Value::as_f64);
99        let (Some(nodes), Some(record), Some(_)) = (num("node_count"), num("record_size"), num("ip_version")) else {
100            return Err(bad("Not a MaxMind DB file: bad metadata"));
101        };
102        if ![24.0, 28.0, 32.0].contains(&record) {
103            return Err(MmdbError(format!("Unsupported record size {record}")));
104        }
105        db.node_count = nodes as u64;
106        db.record_size = record as u64;
107        db.node_bytes = db.record_size / 4;
108        db.data_start = db.node_count * db.node_bytes + 16;
109        db.metadata = metadata;
110        Ok(db)
111    }
112
113    fn ip_version(&self) -> u64 {
114        self.metadata.get("ip_version").and_then(Value::as_f64).unwrap_or(0.0) as u64
115    }
116
117    /// `length` bytes from `at`, fewer at the end of the database.
118    fn read(&self, at: u64, length: u64) -> Result<Vec<u8>, MmdbError> {
119        let end = (at + length).min(self.size);
120        match &self.source {
121            Source::Bytes(b) => Ok(b.get(at as usize..end as usize).unwrap_or(&[]).to_vec()),
122            Source::File(file) => {
123                let mut guard = file.lock().unwrap_or_else(|e| e.into_inner());
124                let (handle, pages) = &mut *guard;
125                let mut out = Vec::new();
126                let mut at = at;
127                while at < end {
128                    let number = at / PAGE;
129                    if !pages.contains_key(&number) {
130                        if pages.len() >= PAGES_KEPT {
131                            pages.clear();
132                        }
133                        let mut page = Vec::with_capacity(PAGE as usize);
134                        handle.seek(SeekFrom::Start(number * PAGE)).map_err(|e| MmdbError(e.to_string()))?;
135                        handle.by_ref().take(PAGE).read_to_end(&mut page).map_err(|e| MmdbError(e.to_string()))?;
136                        pages.insert(number, page);
137                    }
138                    let page = &pages[&number];
139                    let offset = (at - number * PAGE) as usize;
140                    let take = ((end - at) as usize).min(page.len().saturating_sub(offset));
141                    if take == 0 {
142                        break;
143                    }
144                    out.extend_from_slice(&page[offset..offset + take]);
145                    at += take as u64;
146                }
147                Ok(out)
148            }
149        }
150    }
151
152    fn byte(&self, at: u64) -> Result<u64, MmdbError> {
153        self.read(at, 1)?.first().map(|b| u64::from(*b)).ok_or_else(|| bad("Invalid MaxMind DB: read past the end"))
154    }
155
156    /// The record for an address, or `None`. An error for text that is not
157    /// an IP address, or an IPv6 address in an IPv4-only database.
158    pub fn get(&self, ip: &str) -> Result<Option<Value>, MmdbError> {
159        let address: IpAddr = ip.parse().map_err(|_| MmdbError(format!("Not an IP address: {ip}")))?;
160        let packed: Vec<u8> = match address {
161            IpAddr::V4(a) => a.octets().to_vec(),
162            IpAddr::V6(a) => a.octets().to_vec(),
163        };
164        let v6 = packed.len() == 16;
165        if v6 && self.ip_version() == 4 {
166            return Err(MmdbError(format!("An IPv6 address cannot be looked up in an IPv4-only database: {ip}")));
167        }
168        let mut node = if v6 || self.ip_version() == 4 { 0 } else { self.ipv4_start()? };
169        let bits = packed.len() * 8;
170        let mut i = 0;
171        while i < bits && node < self.node_count {
172            let bit = (packed[i >> 3] >> (7 - (i & 7))) & 1;
173            node = self.record(node, u64::from(bit))?;
174            i += 1;
175        }
176        // The node count itself means no record, and so does a tree that ends before the address does.
177        if node <= self.node_count {
178            return Ok(None);
179        }
180        let (value, _) = self.decode(self.data_start + node - self.node_count - 16, self.data_start, 0)?;
181        Ok(Some(value))
182    }
183
184    /// IPv4 addresses live under ::/96 in an IPv6 tree: the node 96 left
185    /// turns down.
186    fn ipv4_start(&self) -> Result<u64, MmdbError> {
187        let mut cached = self.ipv4_start.lock().unwrap_or_else(|e| e.into_inner());
188        if let Some(n) = *cached {
189            return Ok(n);
190        }
191        let mut node = 0;
192        let mut i = 0;
193        while i < 96 && node < self.node_count {
194            node = self.record(node, 0)?;
195            i += 1;
196        }
197        *cached = Some(node);
198        Ok(node)
199    }
200
201    fn record(&self, node: u64, right: u64) -> Result<u64, MmdbError> {
202        let b = self.read(node * self.node_bytes, self.node_bytes)?;
203        if (b.len() as u64) < self.node_bytes {
204            return Err(bad("Invalid MaxMind DB: read past the end"));
205        }
206        let b: Vec<u64> = b.into_iter().map(u64::from).collect();
207        Ok(match self.record_size {
208            24 => {
209                let at = (right * 3) as usize;
210                (b[at] << 16) | (b[at + 1] << 8) | b[at + 2]
211            }
212            28 => {
213                if right == 0 {
214                    ((b[3] & 0xF0) << 20) | (b[0] << 16) | (b[1] << 8) | b[2]
215                } else {
216                    ((b[3] & 0x0F) << 24) | (b[4] << 16) | (b[5] << 8) | b[6]
217                }
218            }
219            _ => {
220                let at = (right * 4) as usize;
221                (b[at] << 24) | (b[at + 1] << 16) | (b[at + 2] << 8) | b[at + 3]
222            }
223        })
224    }
225
226    /// Decodes the value at `at`; pointers are offsets from `base`. Gives
227    /// the value and the offset just past it.
228    fn decode(&self, at: u64, base: u64, depth: usize) -> Result<(Value, u64), MmdbError> {
229        if depth > 64 {
230            return Err(bad("Invalid MaxMind DB: nested too deeply"));
231        }
232        let mut at = at;
233        let control = self.byte(at)?;
234        at += 1;
235        let mut kind = control >> 5;
236        if kind == 1 {
237            // A pointer: up to four more bytes of offset, then the value found there.
238            let ss = (control >> 3) & 3;
239            let vvv = control & 7;
240            let pointer = match ss {
241                0 => (vvv << 8) | self.byte(at)?,
242                1 => ((vvv << 16) | (self.byte(at)? << 8) | self.byte(at + 1)?) + 2048,
243                2 => ((vvv << 24) | (self.byte(at)? << 16) | (self.byte(at + 1)? << 8) | self.byte(at + 2)?) + 526_336,
244                _ => unsigned(&self.read(at, 4)?),
245            };
246            let (value, _) = self.decode(base + pointer, base, depth + 1)?;
247            return Ok((value, at + ss + 1));
248        }
249        if kind == 0 {
250            kind = 7 + self.byte(at)?;
251            at += 1;
252        }
253        let mut size = control & 0x1F;
254        if size >= 29 {
255            let extra = size - 28;
256            let mut n = 0;
257            for i in 0..extra {
258                n = (n << 8) | self.byte(at + i)?;
259            }
260            size = match size {
261                29 => 29,
262                30 => 285,
263                _ => 65_821,
264            } + n;
265            at += extra;
266        }
267        Ok(match kind {
268            2 => (Value::String(String::from_utf8_lossy(&self.read(at, size)?).into_owned()), at + size),
269            3 => {
270                let b = self.read(at, 8)?;
271                let bytes: [u8; 8] = b.try_into().map_err(|_| bad("Invalid MaxMind DB: read past the end"))?;
272                (Value::Number(f64::from_be_bytes(bytes)), at + 8)
273            }
274            4 => (Value::String(String::from_utf8_lossy(&self.read(at, size)?).into_owned()), at + size),
275            5 | 6 => (Value::Number(unsigned(&self.read(at, size)?) as f64), at + size),
276            7 => {
277                let mut map = Object::new();
278                for _ in 0..size {
279                    let (key, next) = self.decode(at, base, depth + 1)?;
280                    let (value, next) = self.decode(next, base, depth + 1)?;
281                    map.set(js::js_string(&key), value);
282                    at = next;
283                }
284                (Value::Object(map), at)
285            }
286            8 => {
287                let mut n = unsigned(&self.read(at, size)?) as i64;
288                if size == 4 && n >= 0x8000_0000 {
289                    n -= 0x1_0000_0000;
290                }
291                (Value::Number(n as f64), at + size)
292            }
293            9 | 10 => {
294                let bytes = self.read(at, size)?;
295                let n = bytes.iter().fold(0u128, |n, b| (n << 8) | u128::from(*b));
296                // A number past 2^53 is written as its decimal text.
297                let v = if n <= 1u128 << 53 { Value::Number(n as f64) } else { Value::String(n.to_string()) };
298                (v, at + size)
299            }
300            11 => {
301                let mut list = Vec::new();
302                for _ in 0..size {
303                    let (value, next) = self.decode(at, base, depth + 1)?;
304                    list.push(value);
305                    at = next;
306                }
307                (Value::Array(list), at)
308            }
309            14 => (Value::Bool(size != 0), at),
310            15 => {
311                let b = self.read(at, 4)?;
312                let bytes: [u8; 4] = b.try_into().map_err(|_| bad("Invalid MaxMind DB: read past the end"))?;
313                (Value::Number(f64::from(f32::from_be_bytes(bytes))), at + 4)
314            }
315            _ => return Err(MmdbError(format!("Invalid MaxMind DB: unknown data type {kind}"))),
316        })
317    }
318}
319
320fn unsigned(bytes: &[u8]) -> u64 {
321    bytes.iter().fold(0u64, |n, b| (n << 8) | u64::from(*b))
322}
323
324/// A city as people say it, without a trailing bracketed district.
325pub fn city_name(name: &str) -> String {
326    let re = crate::re::uni_re!(
327        r"[\t\n\x{0B}\f\r \x{A0}\x{1680}\x{2000}-\x{200A}\x{2028}\x{2029}\x{202F}\x{205F}\x{3000}\x{FEFF}]*\([^)]*\)[\t\n\x{0B}\f\r \x{A0}\x{1680}\x{2000}-\x{200A}\x{2028}\x{2029}\x{202F}\x{205F}\x{3000}\x{FEFF}]*$"
328    );
329    js::trim(&re.replace(name, "")).to_string()
330}
331
332/// What a lookup makes of a database record. DB-IP's records follow
333/// MaxMind's city layout, with names but no subdivision codes; a city loses
334/// the district DB-IP adds in brackets, as in "Toronto (Old Toronto)". This
335/// is the TypeScript server's `lookupFrom`.
336pub fn found_in(record: &Value) -> Option<Found> {
337    let country = record.get("country").and_then(|c| c.get("iso_code"))?;
338    if !js::truthy(country) {
339        return None;
340    }
341    let sub = record.get("subdivisions").and_then(|s| s.as_array()).and_then(|a| a.first());
342    let region = sub
343        .and_then(|s| s.get("iso_code").filter(|v| !v.is_null()))
344        .or_else(|| sub.and_then(|s| s.get("names")).and_then(|n| n.get("en")).filter(|v| !v.is_null()));
345    let city = record.get("city").and_then(|c| c.get("names")).and_then(|n| n.get("en")).filter(|v| !v.is_null());
346    Some(Found {
347        country: Some(js::js_string(country)),
348        region: Some(region.map(js::js_string).unwrap_or_default()),
349        city: Some(match city {
350            Some(Value::String(s)) => city_name(s),
351            Some(other) => js::js_string(other),
352            None => String::new(),
353        }),
354    })
355}
356
357impl GeoLookup for Mmdb {
358    fn lookup<'a>(&'a self, ip: &'a str) -> BoxFuture<'a, Result<Option<Found>, BoxError>> {
359        Box::pin(async move {
360            Ok(match self.get(ip) {
361                Ok(Some(record)) => found_in(&record),
362                _ => None,
363            })
364        })
365    }
366}
runlight/mmdb.rs

runlight/
mmdb.rs
