mod config;
mod parse;
mod query;
use std::collections::HashMap;
use std::fmt;
use std::fs;
use std::io::{self, BufRead, BufReader};
use std::net::IpAddr;
use std::ops::Range;
use std::path::{Path, PathBuf};
use std::slice;
use std::time::Duration;
use ipnet::IpNet;
use prefix_trie::joint::JointPrefixMap;
use tracing::{debug, warn};
use crate::geoip::AutoDownloadConfig;
use crate::geoip::DatabaseFormat;
use crate::geoip::download::fetch::{self, Archive, Credential, Transfer};
use crate::geoip::download::verify::Guard;
use crate::geoip::download::{self, GeoIpDownloadError, SECS_PER_DAY, is_fresh};
use parse::Parsed;
pub use config::{Index, Schema, SourceArchive, TableFormat, TableSource};
pub(crate) use parse::probe;
pub use query::{Condition, Fields, Matched, Matches, Query};
const ADDRESS_COLUMNS: [&str; 5] = ["ip", "ip_address", "ipaddress", "address", "addr"];
const PREFIX_COLUMNS: [&str; 5] = ["prefix", "network", "cidr", "range", "subnet"];
const ADDRESS_SAMPLE: usize = 64;
const NO_ROWS: &[usize] = &[];
type Cell = Option<Box<str>>;
#[derive(Debug, thiserror::Error)]
#[non_exhaustive]
pub enum TableError {
#[error("fetching the table failed: {0}")]
Fetch(#[from] GeoIpDownloadError),
#[error("IO error: {0}")]
Io(#[from] io::Error),
#[error("{url} is not a URL that can be fetched: {detail}")]
BadUrl {
url: String,
detail: String,
},
#[error("file {file} is not a plain file name")]
UnsafeFileName {
file: String,
},
#[error("the table is malformed at line {line}: {detail}")]
Malformed {
line: usize,
detail: String,
},
#[error("the source is not a JSON array of objects: {detail}")]
NotAnArrayOfObjects {
detail: String,
},
#[error("a CSV with no header row needs schema.named to supply the column names")]
NamesRequired,
#[error("schema.named names no columns")]
NoNames,
#[error("the source holds no rows")]
Empty,
#[error("the index names column {column}, which is not one of: {columns}")]
UnknownColumn {
column: String,
columns: String,
},
#[error("no column holds IP addresses; name one with index.column instead. columns: {columns}")]
NoAddressColumn {
columns: String,
},
#[error("no column holds CIDR ranges; name one with index.column instead. columns: {columns}")]
NoPrefixColumn {
columns: String,
},
#[error(
"the table outgrew the {ceiling}-byte resident ceiling; raise \
auto_download.resident_max_bytes or point the source at less data"
)]
OverResidentCeiling {
ceiling: u64,
},
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[non_exhaustive]
pub enum TableBacking {
Resident,
}
impl TableBacking {
#[must_use]
pub const fn label(self) -> &'static str {
match self {
Self::Resident => "resident",
}
}
}
#[derive(Debug)]
enum Keys {
Text(HashMap<Box<str>, Vec<usize>>),
Address {
filed: HashMap<IpAddr, Vec<usize>>,
unindexed: Vec<usize>,
},
Prefix {
filed: Box<JointPrefixMap<IpNet, Vec<usize>>>,
unindexed: Vec<usize>,
},
}
#[derive(Debug)]
pub struct Table {
columns: Vec<String>,
key_column: usize,
rows: Vec<Vec<Cell>>,
keys: Keys,
}
impl Table {
pub async fn ensure(
source: &TableSource,
auto: &AutoDownloadConfig,
) -> Result<Self, TableError> {
source.validate()?;
let dest = auto.data_dir.join(&source.file);
if !auto.enabled {
debug!(file = %dest.display(), "auto-download is off, reading what is on disk");
} else if is_fresh(&dest, u64::from(auto.max_age_days) * SECS_PER_DAY) {
debug!(file = %dest.display(), "table is fresh");
} else {
download::check_settings(auto, source.http_client.is_some())?;
let client = fetch::client(
source.http_client.as_ref(),
Duration::from_secs(auto.connect_timeout_secs),
Duration::from_secs(auto.read_timeout_secs),
)?;
let transfer = transfer(source, dest.clone());
let guard = Guard::for_table(source.format, &source.schema, &source.index, auto);
if let Err(e) = transfer.run_guarded(&client, guard).await {
if !dest.exists() {
return Err(e.into());
}
warn!(
file = %dest.display(),
error = %e,
"table download failed, reading the copy already on disk"
);
}
}
let reader = BufReader::new(fs::File::open(&dest)?);
let parsed = parse::read(
reader,
source.format,
&source.schema,
auto.resident_max_bytes,
)?;
let table = Self::build(parsed, &source.index)?;
telemetry::backing(&source.file, table.backing());
Ok(table)
}
pub fn load(
path: &Path,
format: TableFormat,
schema: &Schema,
index: &Index,
) -> Result<Self, TableError> {
Self::from_reader(BufReader::new(fs::File::open(path)?), format, schema, index)
}
pub fn from_reader(
reader: impl BufRead,
format: TableFormat,
schema: &Schema,
index: &Index,
) -> Result<Self, TableError> {
Self::build(parse::read(reader, format, schema, u64::MAX)?, index)
}
#[must_use]
pub fn columns(&self) -> &[String] {
&self.columns
}
#[must_use]
pub fn key_column(&self) -> &str {
&self.columns[self.key_column]
}
#[must_use]
pub const fn keyed_by_address(&self) -> bool {
matches!(self.keys, Keys::Address { .. } | Keys::Prefix { .. })
}
#[must_use]
pub const fn backing(&self) -> TableBacking {
TableBacking::Resident
}
#[must_use]
pub fn len(&self) -> usize {
self.rows.len()
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.rows.is_empty()
}
#[must_use]
pub fn rows(&self) -> Rows<'_> {
Rows {
table: self,
positions: Positions::All(0..self.rows.len()),
}
}
#[must_use]
pub fn get(&self, key: &str) -> Option<Row<'_>> {
self.positions_of(key).first().map(|&at| self.row(at))
}
#[must_use]
pub fn all(&self, key: &str) -> Rows<'_> {
Rows {
table: self,
positions: Positions::Listed(self.positions_of(key).iter()),
}
}
#[must_use]
pub fn get_by_address(&self, address: IpAddr) -> Option<Row<'_>> {
self.positions_at(address).first().map(|&at| self.row(at))
}
#[must_use]
pub fn all_by_address(&self, address: IpAddr) -> Rows<'_> {
Rows {
table: self,
positions: Positions::Listed(self.positions_at(address).iter()),
}
}
fn build(parsed: Parsed, index: &Index) -> Result<Self, TableError> {
let Parsed { columns, rows } = parsed;
let key_column = match index {
Index::Column(name) => columns
.iter()
.position(|column| column == name)
.ok_or_else(|| TableError::UnknownColumn {
column: name.clone(),
columns: columns.join(", "),
})?,
Index::Ip => keyed_column(&columns, &rows, &ADDRESS_COLUMNS, |text| {
as_address(text).is_some()
})
.ok_or_else(|| TableError::NoAddressColumn {
columns: columns.join(", "),
})?,
Index::Prefix => keyed_column(&columns, &rows, &PREFIX_COLUMNS, |text| {
as_prefix(text).is_some()
})
.ok_or_else(|| TableError::NoPrefixColumn {
columns: columns.join(", "),
})?,
};
let mut reachable = 0;
let keys = match index {
Index::Column(_) => {
let mut filed: HashMap<Box<str>, Vec<usize>> = HashMap::new();
for (at, row) in rows.iter().enumerate() {
if let Some(key) = cell(row, key_column) {
filed.entry(Box::from(key)).or_default().push(at);
reachable += 1;
}
}
Keys::Text(filed)
}
Index::Ip => {
let mut filed: HashMap<IpAddr, Vec<usize>> = HashMap::new();
let mut unindexed = Vec::new();
for (at, row) in rows.iter().enumerate() {
let Some(text) = cell(row, key_column) else {
continue;
};
if let Some(address) = as_address(text) {
filed.entry(address).or_default().push(at);
reachable += 1;
} else {
unindexed.push(at);
}
}
Keys::Address { filed, unindexed }
}
Index::Prefix => {
let mut filed: JointPrefixMap<IpNet, Vec<usize>> = JointPrefixMap::default();
let mut unindexed = Vec::new();
for (at, row) in rows.iter().enumerate() {
let Some(text) = cell(row, key_column) else {
continue;
};
if let Some(prefix) = as_prefix(text) {
filed.entry(prefix).or_default().push(at);
reachable += 1;
} else {
unindexed.push(at);
}
}
Keys::Prefix {
filed: Box::new(filed),
unindexed,
}
}
};
if reachable < rows.len() {
debug!(
column = columns[key_column].as_str(),
rows = rows.len(),
unreachable = rows.len() - reachable,
"some table rows carry no key"
);
}
Ok(Self {
columns,
key_column,
rows,
keys,
})
}
fn positions_of(&self, key: &str) -> &[usize] {
match &self.keys {
Keys::Text(filed) => filed.get(key).map_or(NO_ROWS, Vec::as_slice),
Keys::Address { .. } | Keys::Prefix { .. } => NO_ROWS,
}
}
fn positions_at(&self, address: IpAddr) -> &[usize] {
match &self.keys {
Keys::Address { filed, .. } => filed.get(&address).map_or(NO_ROWS, Vec::as_slice),
Keys::Prefix { filed, .. } => filed
.get_lpm(&IpNet::from(address))
.map_or(NO_ROWS, |(_, rows)| rows.as_slice()),
Keys::Text(_) => NO_ROWS,
}
}
fn row(&self, at: usize) -> Row<'_> {
Row {
table: self,
cells: &self.rows[at],
}
}
}
impl TableSource {
pub fn validate(&self) -> Result<(), TableError> {
reqwest::Url::parse(&self.url).map_err(|e| TableError::BadUrl {
url: fetch::loggable(&self.url),
detail: e.to_string(),
})?;
if Path::new(&self.file).file_name() != Some(self.file.as_ref()) {
return Err(TableError::UnsafeFileName {
file: self.file.clone(),
});
}
match (&self.format, &self.schema) {
(TableFormat::Csv { header: false }, Schema::Auto) => Err(TableError::NamesRequired),
(_, Schema::Named(named)) if named.is_empty() => Err(TableError::NoNames),
_ => Ok(()),
}
}
}
#[derive(Clone, Copy)]
pub struct Row<'a> {
table: &'a Table,
cells: &'a [Cell],
}
impl<'a> Row<'a> {
#[must_use]
pub fn get(&self, column: &str) -> Option<&'a str> {
self.at(self.table.position_of(column)?)
}
#[must_use]
pub fn at(&self, column: usize) -> Option<&'a str> {
cell(self.cells, column)
}
#[must_use]
pub fn key(&self) -> Option<&'a str> {
self.at(self.table.key_column)
}
#[must_use]
pub fn columns(&self) -> &'a [String] {
&self.table.columns
}
}
impl fmt::Debug for Row<'_> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
let mut map = f.debug_map();
for (column, value) in self.table.columns.iter().zip(self.cells) {
map.entry(column, &value.as_deref());
}
map.finish()
}
}
#[derive(Debug)]
enum Positions<'a> {
Listed(slice::Iter<'a, usize>),
All(Range<usize>),
}
#[derive(Debug)]
pub struct Rows<'a> {
table: &'a Table,
positions: Positions<'a>,
}
impl<'a> Iterator for Rows<'a> {
type Item = Row<'a>;
fn next(&mut self) -> Option<Self::Item> {
let at = match &mut self.positions {
Positions::Listed(listed) => *listed.next()?,
Positions::All(all) => all.next()?,
};
Some(self.table.row(at))
}
fn size_hint(&self) -> (usize, Option<usize>) {
match &self.positions {
Positions::Listed(listed) => listed.size_hint(),
Positions::All(all) => all.size_hint(),
}
}
}
impl ExactSizeIterator for Rows<'_> {}
fn transfer(source: &TableSource, dest: PathBuf) -> Transfer {
Transfer {
url: source.url.clone(),
fallback_url: None,
checksum_url: source.checksum_url.clone(),
dest,
archive: match source.archive {
SourceArchive::Raw => Archive::Raw,
SourceArchive::Gzip => Archive::Gzip,
},
format: match source.format {
TableFormat::Csv { .. } => DatabaseFormat::Csv,
TableFormat::Json => DatabaseFormat::Json,
},
credential: Credential::None,
}
}
fn cell(cells: &[Cell], at: usize) -> Option<&str> {
cells.get(at).and_then(Option::as_deref)
}
fn as_address(text: &str) -> Option<IpAddr> {
text.parse().ok()
}
fn as_prefix(text: &str) -> Option<IpNet> {
text.parse()
.ok()
.or_else(|| as_address(text).map(IpNet::from))
}
fn keyed_column(
columns: &[String],
rows: &[Vec<Cell>],
names: &[&str],
parses: fn(&str) -> bool,
) -> Option<usize> {
let holds = |at: usize| {
let mut seen = 0;
for row in rows.iter().take(ADDRESS_SAMPLE) {
let Some(text) = cell(row, at) else {
continue;
};
if !parses(text) {
return false;
}
seen += 1;
}
seen > 0
};
let conventional = columns
.iter()
.position(|column| names.contains(&column.to_ascii_lowercase().as_str()));
if let Some(at) = conventional
&& holds(at)
{
return Some(at);
}
(0..columns.len()).find(|&at| holds(at))
}
#[cfg(feature = "metrics")]
mod telemetry {
use super::TableBacking;
const TYPE: &str = "table";
pub(super) fn backing(file: &str, chosen: TableBacking) {
metrics::gauge!(
"enrichment_table_backing",
"type" => TYPE,
"file" => file.to_string(),
"backing" => chosen.label(),
)
.set(1.0);
}
}
#[cfg(not(feature = "metrics"))]
mod telemetry {
use super::TableBacking;
pub(super) const fn backing(_file: &str, _chosen: TableBacking) {}
}
#[cfg(test)]
mod tests;