mmdb_writer/writer.rs
1//! [`Writer`] — the public entry point. Accumulates `(network, value)` inserts and produces
2//! the final `.mmdb` byte sequence.
3//!
4//! ## File layout ([MMDB spec] §2)
5//!
6//! ```text
7//! [ search-tree bytes ]
8//! [ 16 × 0x00 data-section separator ]
9//! [ data-section bytes ]
10//! [ \xab\xcd\xef MaxMind.com metadata marker (14 bytes) ]
11//! [ metadata map bytes ]
12//! ```
13//!
14//! A reader scans backward from the end for the last occurrence of the metadata marker, then
15//! decodes the map that follows.
16//!
17//! [MMDB spec]: https://maxmind.github.io/MaxMind-DB/
18
19use std::collections::{BTreeMap, HashMap};
20use std::io;
21use std::net::IpAddr;
22use std::time::{SystemTime, UNIX_EPOCH};
23
24use ipnet::IpNet;
25
26use crate::data_section::{DataOffset, DataSection};
27use crate::error::Error;
28use crate::metadata::Metadata;
29use crate::net::{IpVersion, alias_networks, range_to_networks, to_tree_prefix};
30use crate::options::{Ipv4Aliasing, MergeStrategy, MetadataPointers, ReservedNetworks};
31use crate::pool::{ValueId, ValuePool};
32use crate::record_size::RecordSize;
33use crate::reserved;
34use crate::tree::Tree;
35use crate::value::Value;
36
37/// 16 bytes of `0x00` separating the tree section from the data section.
38const DATA_SECTION_SEPARATOR: [u8; 16] = [0; 16];
39
40/// Marker that precedes the metadata section. Readers scan for its last occurrence.
41const METADATA_MARKER: &[u8; 14] = b"\xab\xcd\xefMaxMind.com";
42
43fn default_languages() -> Vec<String> {
44 vec!["en".to_owned()]
45}
46
47/// Builds a MaxMind DB and serializes it.
48///
49/// Construct one with [`Writer::new`] for the common case, or [`Writer::builder`] to set
50/// options such as the [`IpVersion`], descriptions, or a fixed [`RecordSize`]. Then add data
51/// with [`insert`](Writer::insert) / [`insert_value`](Writer::insert_value) and produce bytes
52/// with [`to_bytes`](Writer::to_bytes) or [`write_to`](Writer::write_to).
53///
54/// ```
55/// use ipnet::IpNet;
56/// use mmdb_writer::{Value, Writer};
57///
58/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
59/// let mut writer = Writer::new("Example-DB");
60/// writer.insert_value(
61/// "192.0.2.0/24".parse::<IpNet>()?,
62/// Value::map([("hello", Value::from("world"))]),
63/// )?;
64/// let bytes = writer.to_bytes()?;
65/// assert!(!bytes.is_empty());
66/// # Ok(())
67/// # }
68/// ```
69#[derive(Debug)]
70pub struct Writer {
71 database_type: String,
72 description: BTreeMap<String, String>,
73 languages: Vec<String>,
74 ip_version: IpVersion,
75 record_size: Option<RecordSize>,
76 ipv4_aliasing: Ipv4Aliasing,
77 reserved_networks: ReservedNetworks,
78 metadata_pointers: MetadataPointers,
79 build_epoch: Option<SystemTime>,
80 tree: Tree,
81 pool: ValuePool,
82}
83
84#[bon::bon]
85impl Writer {
86 /// Start building a [`Writer`] with options.
87 ///
88 /// `database_type` names the database for readers (conventionally `Vendor-Dataset`).
89 /// Every other option has a default, so the terminal [`build`](WriterBuilder::build) can
90 /// follow immediately.
91 ///
92 /// ```
93 /// use mmdb_writer::{IpVersion, RecordSize, Writer};
94 ///
95 /// let writer = Writer::builder("Example-DB")
96 /// .ip_version(IpVersion::V4)
97 /// .record_size(RecordSize::Bits32)
98 /// .languages(["en", "de"])
99 /// .build();
100 /// ```
101 #[builder(builder_type = WriterBuilder, finish_fn = build)]
102 pub fn builder(
103 #[builder(start_fn, into)] database_type: String,
104 /// Locales the descriptions cover. Defaults to `["en"]`.
105 #[builder(default = default_languages(), with = |langs: impl IntoIterator<Item: Into<String>>| langs.into_iter().map(Into::into).collect())]
106 languages: Vec<String>,
107 /// Per-language description strings as `(language code, text)` pairs. Defaults to
108 /// empty.
109 ///
110 /// ```
111 /// # use mmdb_writer::Writer;
112 /// let w = Writer::builder("DB")
113 /// .description(&[("en", "An example database"), ("de", "Eine Beispieldatenbank")])
114 /// .build();
115 /// ```
116 #[builder(default, with = |entries: &[(&str, &str)]| entries
117 .iter()
118 .map(|(lang, text)| ((*lang).to_owned(), (*text).to_owned()))
119 .collect())]
120 description: BTreeMap<String, String>,
121 /// IP version of the database. Defaults to [`IpVersion::V6`].
122 #[builder(default)]
123 ip_version: IpVersion,
124 /// Fixed tree record size. Defaults to automatic selection (smallest that fits).
125 record_size: Option<RecordSize>,
126 /// Whether to install IPv4 aliases (V6 only). Defaults to
127 /// [`Ipv4Aliasing::Enabled`].
128 #[builder(default)]
129 ipv4_aliasing: Ipv4Aliasing,
130 /// Whether reserved networks are writable. Defaults to
131 /// [`ReservedNetworks::Included`] (note: this differs from the Go writer).
132 #[builder(default)]
133 reserved_networks: ReservedNetworks,
134 /// Whether the metadata section may use pointers. Defaults to
135 /// [`MetadataPointers::Enabled`].
136 #[builder(default)]
137 metadata_pointers: MetadataPointers,
138 /// Build timestamp written to the metadata. Defaults to the current time; set it for
139 /// reproducible output.
140 build_epoch: Option<SystemTime>,
141 ) -> Self {
142 let mut tree = Tree::new();
143 if reserved_networks.is_excluded() {
144 for network in reserved::networks(ip_version) {
145 let (bits, prefix_len) =
146 to_tree_prefix(network, ip_version).expect("reserved network fits the tree");
147 tree.paint_reserved(bits, prefix_len)
148 .expect("reserved network fits the tree");
149 }
150 }
151 Self {
152 database_type,
153 description,
154 languages,
155 ip_version,
156 record_size,
157 ipv4_aliasing,
158 reserved_networks,
159 metadata_pointers,
160 build_epoch,
161 tree,
162 pool: ValuePool::new(),
163 }
164 }
165}
166
167impl Writer {
168 /// Create a [`Writer`] with all default options and the given database type.
169 #[must_use]
170 pub fn new(database_type: impl Into<String>) -> Self {
171 Self::builder(database_type).build()
172 }
173
174 /// Insert any [`Serialize`](serde::Serialize) value at every address in `network`.
175 ///
176 /// The value is projected onto the MMDB type system: structs and maps become MMDB maps,
177 /// sequences become arrays, `Option::None` and unit values are dropped, and enums are
178 /// serialized in serde's externally tagged form. Later inserts win, as with
179 /// [`insert_value`](Self::insert_value).
180 ///
181 /// # Errors
182 ///
183 /// Returns [`Error::UnsupportedValue`] if the value uses a type MMDB cannot represent
184 /// (such as `i64`), [`Error::Serialize`] if serialization fails, or
185 /// [`Error::Ipv6InIpv4Tree`] for an IPv6 network in an IPv4 database.
186 ///
187 /// ```
188 /// use ipnet::IpNet;
189 /// use mmdb_writer::Writer;
190 /// use serde::Serialize;
191 ///
192 /// #[derive(Serialize)]
193 /// struct Asn { autonomous_system_number: u32, autonomous_system_organization: String }
194 ///
195 /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
196 /// let mut writer = Writer::new("ASN-DB");
197 /// writer.insert(
198 /// "192.0.2.0/24".parse::<IpNet>()?,
199 /// &Asn { autonomous_system_number: 64_512, autonomous_system_organization: "Example".into() },
200 /// )?;
201 /// # Ok(())
202 /// # }
203 /// ```
204 #[cfg(feature = "serde")]
205 pub fn insert<N: Into<IpNet>, T: serde::Serialize + ?Sized>(
206 &mut self,
207 network: N,
208 value: &T,
209 ) -> Result<(), Error> {
210 let value = crate::ser::to_value(value)?;
211 self.insert_value(network, value)
212 }
213
214 /// Insert a [`Value`] at every address in `network`.
215 ///
216 /// Later inserts win: a more-specific network inserted afterward overrides the covered
217 /// addresses, and a less-specific one overwrites everything it covers. Host bits in
218 /// `network` are ignored.
219 ///
220 /// # Errors
221 ///
222 /// Returns [`Error::Ipv6InIpv4Tree`] if `network` is IPv6 but the database is
223 /// [`IpVersion::V4`].
224 pub fn insert_value<N: Into<IpNet>>(&mut self, network: N, value: Value) -> Result<(), Error> {
225 let net = network.into();
226 self.ensure_insertable(net)?;
227 let (bits, prefix_len) = to_tree_prefix(net, self.ip_version)?;
228 let id = self.pool.intern(value);
229 self.tree.insert(bits, prefix_len, &mut |_| Some(id))?;
230 Ok(())
231 }
232
233 /// Reject inserts that target aliased or (when excluded) reserved space, matching the Go
234 /// writer's `AliasedNetworkError` / `ReservedNetworkError`.
235 fn ensure_insertable(&self, net: IpNet) -> Result<(), Error> {
236 use ipnet::IpNet as N;
237 // `contains` is true when `net` is equal to or inside the blocking network. A `net`
238 // that *contains* a blocked network is allowed — it is carved out at build time.
239 let contains = |outer: &N, inner: &N| outer.contains(inner);
240 if self.ip_version == IpVersion::V6
241 && self.ipv4_aliasing.is_enabled()
242 && alias_networks().iter().any(|a| contains(a, &net))
243 {
244 return Err(Error::AliasedNetwork(net));
245 }
246 if self.reserved_networks.is_excluded()
247 && reserved::networks(self.ip_version)
248 .iter()
249 .any(|r| contains(r, &net))
250 {
251 return Err(Error::ReservedNetwork(net));
252 }
253 Ok(())
254 }
255
256 /// Insert into `network` by computing each covered leaf's new value from its current one.
257 ///
258 /// The operation receives the value currently covering a leaf (`None` where there is
259 /// none) and returns the value to store, or `None` to clear it. Because a network can
260 /// cover many existing leaves, the operation may be called more than once per insert —
261 /// once per distinct value it paints over. This mirrors the Go writer's `InsertFunc`.
262 ///
263 /// ```
264 /// use ipnet::IpNet;
265 /// use mmdb_writer::{Value, Writer};
266 ///
267 /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
268 /// let mut w = Writer::new("Counter");
269 /// let mut bump = |net: &str| {
270 /// w.insert_with(net.parse::<IpNet>().unwrap(), |existing| {
271 /// let n = match existing {
272 /// Some(Value::U32(n)) => *n + 1,
273 /// _ => 1,
274 /// };
275 /// Some(Value::from(n))
276 /// })
277 /// .unwrap();
278 /// };
279 /// bump("192.0.2.0/24");
280 /// bump("192.0.2.0/25"); // overlaps: sees the existing 1, stores 2
281 /// # Ok(())
282 /// # }
283 /// ```
284 ///
285 /// # Errors
286 ///
287 /// Returns [`Error::Ipv6InIpv4Tree`] for an IPv6 network in an IPv4 database.
288 pub fn insert_with<N, F>(&mut self, network: N, mut op: F) -> Result<(), Error>
289 where
290 N: Into<IpNet>,
291 F: FnMut(Option<&Value>) -> Option<Value>,
292 {
293 let net = network.into();
294 self.ensure_insertable(net)?;
295 let (bits, prefix_len) = to_tree_prefix(net, self.ip_version)?;
296 let tree = &mut self.tree;
297 let pool = &mut self.pool;
298 tree.insert(bits, prefix_len, &mut |old_id| {
299 let old_value = old_id.map(|id| pool.get(id).clone());
300 op(old_value.as_ref()).map(|new_value| pool.intern(new_value))
301 })
302 }
303
304 /// Insert a [`Value`] into `network`, combining it with any existing value per `strategy`.
305 ///
306 /// With [`MergeStrategy::Replace`] this is [`insert_value`](Self::insert_value); the other
307 /// strategies merge maps (and, for [`DeepMerge`](MergeStrategy::DeepMerge), nested maps
308 /// and arrays) rather than overwriting.
309 ///
310 /// # Errors
311 ///
312 /// Returns [`Error::Ipv6InIpv4Tree`] for an IPv6 network in an IPv4 database.
313 pub fn insert_value_merged<N: Into<IpNet>>(
314 &mut self,
315 network: N,
316 value: Value,
317 strategy: MergeStrategy,
318 ) -> Result<(), Error> {
319 match strategy {
320 MergeStrategy::Replace => self.insert_value(network, value),
321 MergeStrategy::TopLevelMerge => self.insert_with(network, |existing| {
322 Some(match existing {
323 Some(old) => Value::merge_top_level(old, &value),
324 None => value.clone(),
325 })
326 }),
327 MergeStrategy::DeepMerge => self.insert_with(network, |existing| {
328 Some(match existing {
329 Some(old) => Value::merge_deep(old, &value),
330 None => value.clone(),
331 })
332 }),
333 }
334 }
335
336 /// Insert any [`Serialize`](serde::Serialize) value into `network`, merging per `strategy`.
337 ///
338 /// The serde equivalent of [`insert_value_merged`](Self::insert_value_merged).
339 ///
340 /// # Errors
341 ///
342 /// As [`insert`](Self::insert) plus [`insert_value_merged`](Self::insert_value_merged).
343 #[cfg(feature = "serde")]
344 pub fn insert_merged<N: Into<IpNet>, T: serde::Serialize + ?Sized>(
345 &mut self,
346 network: N,
347 value: &T,
348 strategy: MergeStrategy,
349 ) -> Result<(), Error> {
350 let value = crate::ser::to_value(value)?;
351 self.insert_value_merged(network, value, strategy)
352 }
353
354 /// Insert a [`Value`] at every address in the inclusive range `[start, end]`.
355 ///
356 /// The range is decomposed into the minimal set of CIDR networks and each is inserted
357 /// (with [`insert_value`](Self::insert_value) semantics). Useful for data sources that
358 /// express coverage as start–end pairs rather than CIDRs.
359 ///
360 /// # Errors
361 ///
362 /// Returns [`Error::InvalidRange`] if `start` and `end` are different IP families or
363 /// `start` is above `end`, or [`Error::Ipv6InIpv4Tree`] for an IPv6 range in an IPv4
364 /// database.
365 pub fn insert_range(&mut self, start: IpAddr, end: IpAddr, value: &Value) -> Result<(), Error> {
366 for network in range_to_networks(start, end)? {
367 self.insert_value(network, value.clone())?;
368 }
369 Ok(())
370 }
371
372 /// Look up the value currently covering `ip`, if any.
373 ///
374 /// Reflects the most-specific matching insert, including the effect of merges and
375 /// removals. Handy for tests and debugging. IPv4 aliases are not consulted (they are
376 /// installed only at serialization time), so query IPv4 addresses directly.
377 #[must_use]
378 pub fn get(&self, ip: IpAddr) -> Option<&Value> {
379 let (bits, _) = to_tree_prefix(IpNet::from(ip), self.ip_version).ok()?;
380 let id = self.tree.get(bits, self.ip_version.tree_depth())?;
381 Some(self.pool.get(id))
382 }
383
384 /// Serialize the database to a byte vector.
385 ///
386 /// This does not consume the writer, so more data can be inserted afterward and the
387 /// database re-serialized.
388 ///
389 /// # Errors
390 ///
391 /// Returns [`Error::TreeTooLarge`] if the tree and data section exceed what the chosen
392 /// [`RecordSize`] can address.
393 pub fn to_bytes(&self) -> Result<Vec<u8>, Error> {
394 self.build()
395 }
396
397 /// Serialize the database, writing it to `writer`.
398 ///
399 /// # Errors
400 ///
401 /// Returns [`Error::TreeTooLarge`] as [`to_bytes`](Self::to_bytes) does, or
402 /// [`Error::Io`] if writing fails.
403 pub fn write_to<W: io::Write>(&self, mut writer: W) -> Result<(), Error> {
404 let bytes = self.build()?;
405 writer.write_all(&bytes)?;
406 Ok(())
407 }
408
409 fn build(&self) -> Result<Vec<u8>, Error> {
410 // Compaction and aliasing mutate the tree; work on a copy so `self` is untouched.
411 let mut tree = self.tree.clone();
412 if self.ip_version == IpVersion::V6 && self.ipv4_aliasing.is_enabled() {
413 tree.install_ipv4_aliases()?;
414 }
415 tree.compact();
416
417 // Encode only the values still reachable from the tree.
418 let ids = tree.reachable_data_ids();
419 let mut data = DataSection::new();
420 let mut id_to_offset: HashMap<ValueId, DataOffset> = HashMap::with_capacity(ids.len());
421 for id in ids {
422 let offset = data.push(self.pool.get(id));
423 id_to_offset.insert(id, offset);
424 }
425 let data_len = data.len();
426 let data_bytes = data.into_bytes();
427
428 let record_size = pick_record_size(&tree, data_len, self.record_size)?;
429 let tree_bytes = tree.serialize(record_size, &id_to_offset)?;
430
431 let metadata = Metadata {
432 database_type: &self.database_type,
433 description: &self.description,
434 languages: &self.languages,
435 ip_version: self.ip_version,
436 record_size,
437 node_count: tree.node_count(),
438 build_epoch: self.build_epoch_secs(),
439 disable_pointers: self.metadata_pointers.is_disabled(),
440 };
441 let metadata_bytes = metadata.to_bytes();
442
443 let mut out = Vec::with_capacity(
444 tree_bytes.len()
445 + DATA_SECTION_SEPARATOR.len()
446 + data_bytes.len()
447 + METADATA_MARKER.len()
448 + metadata_bytes.len(),
449 );
450 out.extend_from_slice(&tree_bytes);
451 out.extend_from_slice(&DATA_SECTION_SEPARATOR);
452 out.extend_from_slice(&data_bytes);
453 out.extend_from_slice(METADATA_MARKER);
454 out.extend_from_slice(&metadata_bytes);
455 Ok(out)
456 }
457
458 fn build_epoch_secs(&self) -> u64 {
459 self.build_epoch
460 .unwrap_or_else(SystemTime::now)
461 .duration_since(UNIX_EPOCH)
462 .map_or(0, |d| d.as_secs())
463 }
464}
465
466/// Choose the smallest record size that fits the compacted tree plus data section, unless one
467/// was pinned explicitly.
468fn pick_record_size(
469 tree: &Tree,
470 data_len: usize,
471 requested: Option<RecordSize>,
472) -> Result<RecordSize, Error> {
473 if let Some(size) = requested {
474 return Ok(size);
475 }
476 for candidate in RecordSize::ASCENDING {
477 if tree.fits_record_size(candidate, data_len) {
478 return Ok(candidate);
479 }
480 }
481 Err(Error::TreeTooLarge {
482 node_count: tree.node_count(),
483 max: RecordSize::Bits32.max_value(),
484 record_size: RecordSize::Bits32,
485 })
486}