vector/enrichment_tables/
mmdb.rs

1//! Handles enrichment tables for `type = mmdb`.
2//! Enrichment data is loaded from any database in [MaxMind][maxmind] format.
3//!
4//! [maxmind]: https://maxmind.com
5use std::{fs, net::IpAddr, path::PathBuf, sync::Arc, time::SystemTime};
6
7use maxminddb::Reader;
8use vector_lib::{
9    configurable::configurable_component,
10    enrichment::{Case, Condition, Error, IndexHandle, Table},
11};
12use vrl::value::{ObjectMap, Value};
13
14use crate::config::{EnrichmentTableConfig, GenerateConfig};
15
16/// Configuration for the `mmdb` enrichment table.
17#[derive(Clone, Debug, Eq, PartialEq)]
18#[configurable_component(enrichment_table("mmdb"))]
19pub struct MmdbConfig {
20    /// Path to the [MaxMind][maxmind] database
21    ///
22    /// [maxmind]: https://maxmind.com
23    pub path: PathBuf,
24}
25
26impl GenerateConfig for MmdbConfig {
27    fn generate_config() -> toml::Value {
28        toml::Value::try_from(Self {
29            path: "/path/to/GeoLite2-City.mmdb".into(),
30        })
31        .unwrap()
32    }
33}
34
35impl EnrichmentTableConfig for MmdbConfig {
36    async fn build(
37        &self,
38        _: &crate::config::GlobalOptions,
39        _: Option<Box<dyn std::any::Any + Send + Sync>>,
40    ) -> crate::Result<Box<dyn Table + Send + Sync>> {
41        Ok(Box::new(Mmdb::new(self.clone())?))
42    }
43}
44
45#[derive(Clone)]
46/// A struct that implements [vector_lib::enrichment::Table] to handle loading enrichment data from a MaxMind database.
47pub struct Mmdb {
48    config: MmdbConfig,
49    dbreader: Arc<maxminddb::Reader<Vec<u8>>>,
50    last_modified: SystemTime,
51}
52
53impl Mmdb {
54    /// Creates a new Mmdb struct from the provided config.
55    pub fn new(config: MmdbConfig) -> crate::Result<Self> {
56        let dbreader = Arc::new(Reader::open_readfile(&config.path)?);
57
58        // Check if we can read database with dummy Ip.
59        let ip = IpAddr::V4(std::net::Ipv4Addr::UNSPECIFIED);
60        let result = dbreader.lookup(ip)?.decode::<ObjectMap>().map(|_| ());
61
62        match result {
63            Ok(_) => Ok(Mmdb {
64                last_modified: fs::metadata(&config.path)?.modified()?,
65                dbreader,
66                config,
67            }),
68            Err(error) => Err(error.into()),
69        }
70    }
71
72    fn lookup(&self, ip: IpAddr, select: Option<&[String]>) -> Option<ObjectMap> {
73        let data = self.dbreader.lookup(ip).ok()?.decode().ok()??;
74
75        if let Some(fields) = select {
76            let mut filtered = Value::from(ObjectMap::new());
77            let mut data_value = Value::from(data);
78            for field in fields {
79                // Unparseable field names in the caller's `select` list are
80                // skipped: this matches the pre-migration behavior where the
81                // JIT `&str` path parser failed to iterate and the field was
82                // omitted from the output entirely.
83                let Ok(field_path) = vrl::path::parse_value_path(field) else {
84                    continue;
85                };
86                filtered.insert(
87                    &field_path,
88                    data_value.remove(&field_path, false).unwrap_or(Value::Null),
89                );
90            }
91            filtered.into_object()
92        } else {
93            Some(data)
94        }
95    }
96}
97
98impl Table for Mmdb {
99    /// Search the enrichment table data with the given condition.
100    /// All conditions must match (AND).
101    ///
102    /// # Errors
103    /// Errors if no rows, or more than 1 row is found.
104    fn find_table_row<'a>(
105        &self,
106        case: Case,
107        condition: &'a [Condition<'a>],
108        select: Option<&[String]>,
109        wildcard: Option<&Value>,
110        index: Option<IndexHandle>,
111    ) -> Result<ObjectMap, Error> {
112        let mut rows = self.find_table_rows(case, condition, select, wildcard, index)?;
113
114        match rows.pop() {
115            Some(row) if rows.is_empty() => Ok(row),
116            Some(_) => Err(Error::MoreThanOneRowFound),
117            None => Err(Error::NoRowsFound),
118        }
119    }
120
121    /// Search the enrichment table data with the given condition.
122    /// All conditions must match (AND).
123    /// Can return multiple matched records
124    fn find_table_rows<'a>(
125        &self,
126        _: Case,
127        condition: &'a [Condition<'a>],
128        select: Option<&[String]>,
129        _wildcard: Option<&Value>,
130        _: Option<IndexHandle>,
131    ) -> Result<Vec<ObjectMap>, Error> {
132        match condition.first() {
133            Some(_) if condition.len() > 1 => Err(Error::OnlyOneConditionAllowed),
134            Some(Condition::Equals { value, .. }) => {
135                let ip = value
136                    .to_string_lossy()
137                    .parse::<IpAddr>()
138                    .map_err(|source| Error::InvalidAddress { source })?;
139                Ok(self
140                    .lookup(ip, select)
141                    .map(|values| vec![values])
142                    .unwrap_or_default())
143            }
144            Some(_) => Err(Error::OnlyEqualityConditionAllowed),
145            None => Err(Error::MissingCondition { kind: "IP" }),
146        }
147    }
148
149    /// Hints to the enrichment table what data is going to be searched to allow it to index the
150    /// data in advance.
151    ///
152    /// # Errors
153    /// Errors if the fields are not in the table.
154    fn add_index(&mut self, _: Case, fields: &[&str]) -> Result<IndexHandle, Error> {
155        match fields.len() {
156            0 => Err(Error::MissingRequiredField { field: "IP" }),
157            1 => Ok(IndexHandle(0)),
158            _ => Err(Error::OnlyOneFieldAllowed),
159        }
160    }
161
162    /// Returns a list of the field names that are in each index
163    fn index_fields(&self) -> Vec<(Case, Vec<String>)> {
164        Vec::new()
165    }
166
167    /// Returns true if the underlying data has changed and the table needs reloading.
168    fn needs_reload(&self) -> bool {
169        matches!(fs::metadata(&self.config.path)
170            .and_then(|metadata| metadata.modified()),
171            Ok(modified) if modified > self.last_modified)
172    }
173}
174
175impl std::fmt::Debug for Mmdb {
176    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
177        write!(f, "Maxmind database {})", self.config.path.display())
178    }
179}
180
181#[cfg(test)]
182mod tests {
183    use vrl::value::Value;
184
185    use super::*;
186
187    #[test]
188    fn city_partial_lookup() {
189        let values = find_select(
190            "2.125.160.216",
191            "tests/data/GeoIP2-City-Test.mmdb",
192            Some(&[
193                "location.latitude".to_string(),
194                "location.longitude".to_string(),
195            ]),
196        )
197        .unwrap();
198
199        let mut expected = ObjectMap::new();
200        expected.insert(
201            "location".into(),
202            ObjectMap::from([
203                ("latitude".into(), Value::from(51.75)),
204                ("longitude".into(), Value::from(-1.25)),
205            ])
206            .into(),
207        );
208
209        assert_eq!(values, expected);
210    }
211
212    #[test]
213    fn city_partial_lookup_skips_unparseable_selects() {
214        // An unparseable path in the select list is silently dropped from the
215        // output rather than surfacing as a null-valued literal field.
216        let values = find_select(
217            "2.125.160.216",
218            "tests/data/GeoIP2-City-Test.mmdb",
219            Some(&[
220                "location.latitude".to_string(),
221                "not a valid path".to_string(),
222            ]),
223        )
224        .unwrap();
225
226        let mut expected = ObjectMap::new();
227        expected.insert(
228            "location".into(),
229            ObjectMap::from([("latitude".into(), Value::from(51.75))]).into(),
230        );
231
232        assert_eq!(values, expected);
233    }
234
235    #[test]
236    fn isp_lookup() {
237        let values = find("208.192.1.2", "tests/data/GeoIP2-ISP-Test.mmdb").unwrap();
238
239        let mut expected = ObjectMap::new();
240        expected.insert("autonomous_system_number".into(), 701i64.into());
241        expected.insert(
242            "autonomous_system_organization".into(),
243            "MCI Communications Services, Inc. d/b/a Verizon Business".into(),
244        );
245        expected.insert("isp".into(), "Verizon Business".into());
246        expected.insert("organization".into(), "Verizon Business".into());
247
248        assert_eq!(values, expected);
249    }
250
251    #[test]
252    fn connection_type_lookup_success() {
253        let values = find(
254            "201.243.200.1",
255            "tests/data/GeoIP2-Connection-Type-Test.mmdb",
256        )
257        .unwrap();
258
259        let mut expected = ObjectMap::new();
260        expected.insert("connection_type".into(), "Corporate".into());
261
262        assert_eq!(values, expected);
263    }
264
265    #[test]
266    fn lookup_missing() {
267        let values = find("10.1.12.1", "tests/data/custom-type.mmdb");
268
269        assert!(values.is_none());
270    }
271
272    #[test]
273    fn custom_mmdb_type() {
274        let values = find("208.192.1.2", "tests/data/custom-type.mmdb").unwrap();
275
276        let mut expected = ObjectMap::new();
277        expected.insert("hostname".into(), "custom".into());
278        expected.insert(
279            "nested".into(),
280            ObjectMap::from([
281                ("hostname".into(), "custom".into()),
282                ("original_cidr".into(), "208.192.1.2/24".into()),
283            ])
284            .into(),
285        );
286
287        assert_eq!(values, expected);
288    }
289
290    fn find(ip: &str, database: &str) -> Option<ObjectMap> {
291        find_select(ip, database, None)
292    }
293
294    fn find_select(ip: &str, database: &str, select: Option<&[String]>) -> Option<ObjectMap> {
295        Mmdb::new(MmdbConfig {
296            path: database.into(),
297        })
298        .unwrap()
299        .find_table_rows(
300            Case::Insensitive,
301            &[Condition::Equals {
302                field: "ip",
303                value: ip.into(),
304            }],
305            select,
306            None,
307            None,
308        )
309        .unwrap()
310        .pop()
311    }
312}
vector/enrichment_tables/mmdb.rs

vector/enrichment_tables/
mmdb.rs
