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

vector/enrichment_tables/
mmdb.rs
