maxminddb/
writer.rs

1use std::{
2    collections::{BTreeMap, HashMap},
3    hash::{Hash, Hasher},
4    io::Write,
5    net::IpAddr,
6    time::{SystemTime, UNIX_EPOCH},
7};
8
9use ipnetwork::IpNetwork;
10use serde::Serialize;
11
12use crate::encoder::{self, Value, encode_serialize};
13
14/// The result of a writer operation.
15///
16/// The writer reports errors as strings rather than as
17/// [`MaxMindDBError`](crate::MaxMindDBError) because it constructs databases
18/// rather than reading them.
19pub type WriterResult<T> = Result<T, String>;
20
21fn map_io_err(e: std::io::Error) -> String {
22    e.to_string()
23}
24
25const DATA_SECTION_SEPARATOR: &[u8; 16] = &[0u8; 16];
26const METADATA_START_MARKER: &[u8] = b"\xab\xcd\xefMaxMind.com";
27
28#[derive(Debug, Clone)]
29enum RecordType {
30    Empty,
31    Data,
32    Node,
33}
34
35#[derive(Debug, Clone)]
36struct DataRecord {
37    key: u64,
38}
39
40#[derive(Debug, Clone)]
41struct ChildRecord {
42    record_type: RecordType,
43    data: Option<DataRecord>,
44    node: Option<Box<Node>>,
45}
46
47impl ChildRecord {
48    fn empty() -> Self {
49        ChildRecord {
50            record_type: RecordType::Empty,
51            data: None,
52            node: None,
53        }
54    }
55
56    fn data(key: u64) -> Self {
57        ChildRecord {
58            record_type: RecordType::Data,
59            data: Some(DataRecord {
60                key,
61            }),
62            node: None,
63        }
64    }
65
66    fn node(node: Node) -> Self {
67        ChildRecord {
68            record_type: RecordType::Node,
69            data: None,
70            node: Some(Box::new(node)),
71        }
72    }
73}
74
75#[derive(Debug, Clone)]
76struct Node {
77    node_num: Option<usize>,
78    children: [ChildRecord; 2],
79}
80
81impl Node {
82    fn new() -> Self {
83        Node {
84            node_num: None,
85            children: [ChildRecord::empty(), ChildRecord::empty()],
86        }
87    }
88
89    fn insert(&mut self, ip: &[u8], prefix_len: usize, depth: usize, data_key: u64) -> WriterResult<()> {
90        if depth == prefix_len {
91            self.children[0] = ChildRecord::data(data_key);
92            self.children[1] = ChildRecord::data(data_key);
93            return Ok(());
94        }
95
96        let bit = bit_at(ip, depth);
97        let child = &mut self.children[bit as usize];
98
99        match &child.record_type {
100            RecordType::Empty => {
101                let mut subtree = Node::new();
102                subtree.insert(ip, prefix_len, depth + 1, data_key)?;
103                *child = ChildRecord::node(subtree);
104            }
105            RecordType::Data => {
106                let existing = child.data.as_ref().unwrap();
107                let mut subtree = Node::new();
108                subtree.children[0] = ChildRecord::data(existing.key);
109                subtree.children[1] = ChildRecord::data(existing.key);
110                subtree.insert(ip, prefix_len, depth + 1, data_key)?;
111                *child = ChildRecord::node(subtree);
112            }
113            RecordType::Node => {
114                child
115                    .node
116                    .as_mut()
117                    .unwrap()
118                    .insert(ip, prefix_len, depth + 1, data_key)?;
119            }
120        }
121
122        Ok(())
123    }
124
125    fn finalize(&mut self, next_node: &mut usize) {
126        self.node_num = Some(*next_node);
127        *next_node += 1;
128
129        for i in 0..2 {
130            if let Some(node) = self.children[i].node.as_mut() {
131                node.finalize(next_node);
132            }
133        }
134    }
135
136    fn write_node<W: Write>(
137        &self,
138        writer: &mut W,
139        node_count: usize,
140        data_offsets: &BTreeMap<u64, usize>,
141        record_size: u16,
142    ) -> WriterResult<()> {
143        let left_val = self.resolve_record_value(&self.children[0], node_count, data_offsets)?;
144        let right_val = self.resolve_record_value(&self.children[1], node_count, data_offsets)?;
145
146        let max_record = 1u64 << record_size;
147        if left_val as u64 >= max_record || right_val as u64 >= max_record {
148            return Err(format!(
149                "record value ({}, {}) exceeds max for {} bit record size",
150                left_val, right_val, record_size
151            ));
152        }
153
154        match record_size {
155            24 => writer
156                .write_all(&[
157                    ((left_val >> 16) & 0xFF) as u8,
158                    ((left_val >> 8) & 0xFF) as u8,
159                    (left_val & 0xFF) as u8,
160                    ((right_val >> 16) & 0xFF) as u8,
161                    ((right_val >> 8) & 0xFF) as u8,
162                    (right_val & 0xFF) as u8,
163                ])
164                .map_err(map_io_err)?,
165            28 => writer
166                .write_all(&[
167                    ((left_val >> 16) & 0xFF) as u8,
168                    ((left_val >> 8) & 0xFF) as u8,
169                    (left_val & 0xFF) as u8,
170                    ((((left_val >> 24) & 0x0F) << 4) | ((right_val >> 24) & 0x0F)) as u8,
171                    ((right_val >> 16) & 0xFF) as u8,
172                    ((right_val >> 8) & 0xFF) as u8,
173                    (right_val & 0xFF) as u8,
174                ])
175                .map_err(map_io_err)?,
176            32 => writer
177                .write_all(&[
178                    ((left_val >> 24) & 0xFF) as u8,
179                    ((left_val >> 16) & 0xFF) as u8,
180                    ((left_val >> 8) & 0xFF) as u8,
181                    (left_val & 0xFF) as u8,
182                    ((right_val >> 24) & 0xFF) as u8,
183                    ((right_val >> 16) & 0xFF) as u8,
184                    ((right_val >> 8) & 0xFF) as u8,
185                    (right_val & 0xFF) as u8,
186                ])
187                .map_err(map_io_err)?,
188            s => return Err(format!("unsupported record size: {s}")),
189        }
190
191        for child in &self.children {
192            if let Some(node) = &child.node {
193                node.write_node(writer, node_count, data_offsets, record_size)?;
194            }
195        }
196
197        Ok(())
198    }
199
200    fn resolve_record_value(
201        &self,
202        record: &ChildRecord,
203        node_count: usize,
204        data_offsets: &BTreeMap<u64, usize>,
205    ) -> WriterResult<usize> {
206        match &record.record_type {
207            RecordType::Empty => Ok(node_count),
208            RecordType::Data => {
209                let d = record.data.as_ref().unwrap();
210                let offset = data_offsets.get(&d.key).unwrap();
211                Ok(node_count + 16 + offset)
212            }
213            RecordType::Node => Ok(record.node.as_ref().unwrap().node_num.unwrap()),
214        }
215    }
216}
217
218fn bit_at(ip: &[u8], depth: usize) -> u8 {
219    (ip[depth >> 3] >> (7 - (depth & 7))) & 1
220}
221
222/// Builds a MaxMind DB file from inserted networks.
223///
224/// The writer accumulates encoded data records, deduplicating equal payloads,
225/// and serializes a search tree plus metadata when [`Writer::write_to`] is
226/// called.
227///
228/// # Example
229///
230/// ```
231/// use maxminddb::writer::{Writer, WriterOptions};
232///
233/// let opts = WriterOptions {
234///     database_type: "Test".to_string(),
235///     ip_version: 4,
236///     record_size: 24,
237///     ..Default::default()
238/// };
239/// let mut writer = Writer::new(opts).unwrap();
240/// writer.insert("10.0.0.0/8".parse().unwrap(), "example").unwrap();
241///
242/// let mut buf = Vec::new();
243/// writer.write_to(&mut buf).unwrap();
244/// let reader = maxminddb::Reader::from_source(buf).unwrap();
245/// assert_eq!(reader.lookup::<String>("10.0.0.1".parse().unwrap()).unwrap(), "example");
246/// ```
247#[derive(Debug)]
248pub struct Writer {
249    root: Node,
250    data_map: BTreeMap<u64, Vec<u8>>,
251    /// Maps a hash of an encoded payload to the data keys with that hash, used
252    /// to deduplicate equal payloads without an O(n) scan per insert.
253    data_index: HashMap<u64, Vec<u64>>,
254    node_count: usize,
255    ip_version: u16,
256    record_size: u16,
257    database_type: String,
258    description: BTreeMap<String, String>,
259    languages: Vec<String>,
260    build_epoch: u64,
261    next_data_key: u64,
262}
263
264/// Configuration for a [`Writer`].
265#[derive(Clone)]
266pub struct WriterOptions {
267    /// IP version of the database: `4` or `6`. For an IPv6 database, IPv4
268    /// networks are stored in the IPv4-mapped subtree.
269    pub ip_version: u16,
270    /// Search tree record size in bits: `24`, `28`, or `32`. Larger records
271    /// allow more nodes at the cost of file size.
272    pub record_size: u16,
273    /// Value for the `database_type` metadata field.
274    pub database_type: String,
275    /// Localized descriptions for the `description` metadata field.
276    pub description: BTreeMap<String, String>,
277    /// Languages present in `description`.
278    pub languages: Vec<String>,
279    /// Build time as a Unix timestamp for the `build_epoch` metadata field.
280    pub build_epoch: u64,
281}
282
283impl Default for WriterOptions {
284    fn default() -> Self {
285        WriterOptions {
286            ip_version: 6,
287            record_size: 28,
288            database_type: String::new(),
289            description: BTreeMap::new(),
290            languages: Vec::new(),
291            build_epoch: SystemTime::now()
292                .duration_since(UNIX_EPOCH)
293                .unwrap_or_default()
294                .as_secs(),
295        }
296    }
297}
298
299impl Writer {
300    /// Create a writer.
301    ///
302    /// # Errors
303    ///
304    /// Returns an error if `ip_version` is not `4` or `6`, or if `record_size`
305    /// is not `24`, `28`, or `32`.
306    pub fn new(opts: WriterOptions) -> WriterResult<Self> {
307        if opts.ip_version != 4 && opts.ip_version != 6 {
308            return Err(format!("invalid ip_version: {} (must be 4 or 6)", opts.ip_version));
309        }
310        if opts.record_size != 24 && opts.record_size != 28 && opts.record_size != 32 {
311            return Err(format!("invalid record_size: {} (must be 24, 28, or 32)", opts.record_size));
312        }
313        Ok(Writer {
314            root: Node::new(),
315            data_map: BTreeMap::new(),
316            data_index: HashMap::new(),
317            node_count: 0,
318            ip_version: opts.ip_version,
319            record_size: opts.record_size,
320            database_type: opts.database_type,
321            description: opts.description,
322            languages: opts.languages,
323            build_epoch: opts.build_epoch,
324            next_data_key: 1,
325        })
326    }
327
328    /// Stores an encoded payload, returning the existing key if an equal
329    /// payload was already stored.
330    fn intern_value(&mut self, encoded: Vec<u8>) -> u64 {
331        let mut hasher = std::collections::hash_map::DefaultHasher::new();
332        encoded.hash(&mut hasher);
333        let digest = hasher.finish();
334
335        if let Some(keys) = self.data_index.get(&digest) {
336            for key in keys {
337                if self.data_map.get(key).is_some_and(|existing| existing == &encoded) {
338                    return *key;
339                }
340            }
341        }
342
343        let key = self.next_data_key;
344        self.next_data_key += 1;
345        self.data_map.insert(key, encoded);
346        self.data_index.entry(digest).or_default().push(key);
347        key
348    }
349
350    /// Insert a fully encoded [`Value`] for `network`.
351    ///
352    /// Equal values are stored once in the data section and referenced by every
353    /// network that uses them.
354    ///
355    /// # Errors
356    ///
357    /// Returns an error if an IPv6 network is inserted into an IPv4 database or
358    /// if the value cannot be encoded.
359    ///
360    /// # Overwriting
361    ///
362    /// Inserting a network replaces the subtree at that prefix. Inserting a
363    /// *less specific* network after a more specific one therefore discards the
364    /// more specific entries, so insert least specific networks first.
365    pub fn insert_value(&mut self, network: IpNetwork, value: Value) -> WriterResult<()> {
366        let prefix_len = network.prefix() as usize;
367        let ip = ip_to_bytes(network.network());
368
369        if self.ip_version == 4 && ip.len() != 4 {
370            return Err("cannot insert IPv6 network into an IPv4 tree".to_string());
371        }
372
373        let (ip_bytes, actual_prefix_len) = if self.ip_version == 6 && ip.len() == 4 {
374            let mut v6 = [0u8; 16];
375            v6[12..].copy_from_slice(&ip);
376            (v6.to_vec(), prefix_len + 96)
377        } else {
378            (ip, prefix_len)
379        };
380
381        let encoded = encoder::encode_value(&value)?;
382
383        let data_key = self.intern_value(encoded);
384
385        self.node_count = 0;
386        self.root.insert(&ip_bytes, actual_prefix_len, 0, data_key)?;
387
388        Ok(())
389    }
390
391    /// Insert a `Serialize` value for `network`.
392    ///
393    /// This is the same as [`Writer::insert_value`] but encodes `value` for you.
394    ///
395    /// # Errors
396    ///
397    /// Returns an error if `value` cannot be encoded, if an IPv6 network is
398    /// inserted into an IPv4 database, or if the value is out of range for the
399    /// MaxMind DB format (for example a negative integer below `i32::MIN`).
400    ///
401    /// # Overwriting
402    ///
403    /// See [`Writer::insert_value`]: insert least specific networks first.
404    pub fn insert<T: Serialize>(&mut self, network: IpNetwork, value: T) -> WriterResult<()> {
405        let encoded = encode_serialize(&value).map_err(|e| format!("serialization failed: {e}"))?;
406
407        let prefix_len = network.prefix() as usize;
408        let ip = ip_to_bytes(network.network());
409
410        if self.ip_version == 4 && ip.len() != 4 {
411            return Err("cannot insert IPv6 network into an IPv4 tree".to_string());
412        }
413
414        let (ip_bytes, actual_prefix_len) = if self.ip_version == 6 && ip.len() == 4 {
415            let mut v6 = [0u8; 16];
416            v6[12..].copy_from_slice(&ip);
417            (v6.to_vec(), prefix_len + 96)
418        } else {
419            (ip, prefix_len)
420        };
421
422        let data_key = self.intern_value(encoded);
423
424        self.node_count = 0;
425        self.root.insert(&ip_bytes, actual_prefix_len, 0, data_key)?;
426
427        Ok(())
428    }
429
430    /// Serialize the database into `writer`, returning the number of bytes
431    /// written.
432    ///
433    /// This finalizes node numbering, so call it once, after all inserts.
434    ///
435    /// # Errors
436    ///
437    /// Returns an error if a record value does not fit the configured record
438    /// size, if the node count exceeds `u32`, or if writing to `writer` fails.
439    pub fn write_to<W: Write>(&mut self, writer: &mut W) -> WriterResult<u64> {
440        let mut next = 1usize;
441        for child in &mut self.root.children {
442            if let Some(node) = child.node.as_mut() {
443                node.finalize(&mut next);
444            }
445        }
446        self.node_count = next;
447
448        let search_tree_size_bytes = self
449            .node_count
450            .checked_mul(self.record_size as usize)
451            .map(|v| v / 4)
452            .ok_or_else(|| "search tree size overflow".to_string())?;
453        let node_count_u32 =
454            u32::try_from(self.node_count).map_err(|_| format!("node count {} exceeds u32", self.node_count))?;
455
456        let mut data_offsets: BTreeMap<u64, usize> = BTreeMap::new();
457        let mut data_section = Vec::new();
458        for (key, encoded) in &self.data_map {
459            let offset = data_section.len();
460            data_section.extend_from_slice(encoded);
461            data_offsets.insert(*key, offset);
462        }
463
464        let node_count = self.node_count;
465        self.root
466            .write_node(writer, node_count, &data_offsets, self.record_size)?;
467
468        writer.write_all(DATA_SECTION_SEPARATOR).map_err(map_io_err)?;
469        writer.write_all(&data_section).map_err(map_io_err)?;
470        writer.write_all(METADATA_START_MARKER).map_err(map_io_err)?;
471
472        let metadata = build_metadata_value(
473            node_count_u32,
474            self.record_size,
475            self.ip_version,
476            &self.database_type,
477            &self.description,
478            &self.languages,
479            self.build_epoch,
480        );
481        let metadata_bytes = encoder::encode_value(&metadata)?;
482        const METADATA_MAX_SIZE: usize = 128 * 1024;
483        if metadata_bytes.len() > METADATA_MAX_SIZE {
484            return Err(format!(
485                "encoded metadata is {} bytes; the MaxMind DB maximum is {} bytes",
486                metadata_bytes.len(),
487                METADATA_MAX_SIZE
488            ));
489        }
490        writer.write_all(&metadata_bytes).map_err(map_io_err)?;
491
492        Ok((search_tree_size_bytes
493            + DATA_SECTION_SEPARATOR.len()
494            + data_section.len()
495            + METADATA_START_MARKER.len()
496            + metadata_bytes.len()) as u64)
497    }
498}
499
500fn build_metadata_value(
501    node_count: u32,
502    record_size: u16,
503    ip_version: u16,
504    database_type: &str,
505    description: &BTreeMap<String, String>,
506    languages: &[String],
507    build_epoch: u64,
508) -> Value {
509    let mut desc_map = BTreeMap::new();
510    for (k, v) in description {
511        desc_map.insert(k.clone(), Value::String(v.clone()));
512    }
513
514    let mut lang_array = Vec::new();
515    for l in languages {
516        lang_array.push(Value::String(l.clone()));
517    }
518
519    let mut m = BTreeMap::new();
520    m.insert("binary_format_major_version".to_string(), Value::Uint16(2));
521    m.insert("binary_format_minor_version".to_string(), Value::Uint16(0));
522    m.insert("build_epoch".to_string(), Value::Uint64(build_epoch));
523    m.insert("database_type".to_string(), Value::String(database_type.to_string()));
524    m.insert("description".to_string(), Value::Map(desc_map));
525    m.insert("ip_version".to_string(), Value::Uint16(ip_version));
526    m.insert("languages".to_string(), Value::Slice(lang_array));
527    m.insert("node_count".to_string(), Value::Uint32(node_count));
528    m.insert("record_size".to_string(), Value::Uint16(record_size));
529    Value::Map(m)
530}
531
532fn ip_to_bytes(address: IpAddr) -> Vec<u8> {
533    match address {
534        IpAddr::V4(a) => a.octets().to_vec(),
535        IpAddr::V6(a) => a.octets().to_vec(),
536    }
537}
538
539#[cfg(test)]
540mod tests {
541    use super::*;
542
543    #[test]
544    fn test_write_and_read_simple_v4() {
545        let opts = WriterOptions {
546            database_type: "Test".to_string(),
547            ip_version: 4,
548            record_size: 24,
549            ..Default::default()
550        };
551        let mut tree = Writer::new(opts).unwrap();
552        let mut m = BTreeMap::new();
553        m.insert("country".to_string(), Value::String("US".to_string()));
554        tree.insert_value("1.2.3.0/24".parse().unwrap(), Value::Map(m)).unwrap();
555        let mut buf = Vec::new();
556        tree.write_to(&mut buf).unwrap();
557        let reader = crate::Reader::from_source(buf).unwrap();
558        #[derive(serde::Deserialize, Debug)]
559        struct TestRecord {
560            country: String,
561        }
562        let result: TestRecord = reader.lookup("1.2.3.4".parse().unwrap()).unwrap();
563        assert_eq!(result.country, "US");
564    }
565
566    #[test]
567    fn test_write_and_read_multiple_networks() {
568        let opts = WriterOptions {
569            database_type: "Test".to_string(),
570            ip_version: 4,
571            record_size: 24,
572            ..Default::default()
573        };
574        let mut tree = Writer::new(opts).unwrap();
575        let mut m1 = BTreeMap::new();
576        m1.insert("country".to_string(), Value::String("US".to_string()));
577        tree.insert_value("1.0.0.0/8".parse().unwrap(), Value::Map(m1)).unwrap();
578        let mut m2 = BTreeMap::new();
579        m2.insert("country".to_string(), Value::String("FR".to_string()));
580        tree.insert_value("2.0.0.0/8".parse().unwrap(), Value::Map(m2)).unwrap();
581        let mut buf = Vec::new();
582        tree.write_to(&mut buf).unwrap();
583        let reader = crate::Reader::from_source(buf).unwrap();
584        #[derive(serde::Deserialize, Debug)]
585        struct TestRecord {
586            country: String,
587        }
588        let r1: TestRecord = reader.lookup("1.2.3.4".parse().unwrap()).unwrap();
589        assert_eq!(r1.country, "US");
590        let r2: TestRecord = reader.lookup("2.3.4.5".parse().unwrap()).unwrap();
591        assert_eq!(r2.country, "FR");
592        assert!(reader.lookup::<TestRecord>("3.4.5.6".parse().unwrap()).is_err());
593    }
594
595    #[test]
596    fn test_v6_database() {
597        let opts = WriterOptions {
598            database_type: "Test".to_string(),
599            ip_version: 6,
600            record_size: 28,
601            ..Default::default()
602        };
603        let mut tree = Writer::new(opts).unwrap();
604        let mut m = BTreeMap::new();
605        m.insert("country".to_string(), Value::String("FR".to_string()));
606        tree.insert_value("2a00:1450:4000::/36".parse().unwrap(), Value::Map(m))
607            .unwrap();
608        let mut buf = Vec::new();
609        tree.write_to(&mut buf).unwrap();
610        let reader = crate::Reader::from_source(buf).unwrap();
611        #[derive(serde::Deserialize, Debug)]
612        struct TestRecord {
613            country: String,
614        }
615        let r: TestRecord = reader.lookup("2a00:1450:4000::".parse().unwrap()).unwrap();
616        assert_eq!(r.country, "FR");
617    }
618
619    #[test]
620    fn test_all_record_sizes() {
621        for record_size in [24, 28, 32] {
622            let opts = WriterOptions {
623                database_type: "Test".to_string(),
624                ip_version: 4,
625                record_size,
626                ..Default::default()
627            };
628            let mut tree = Writer::new(opts).unwrap();
629            let mut m = BTreeMap::new();
630            m.insert("ip".to_string(), Value::String("test".to_string()));
631            tree.insert_value("10.0.0.0/8".parse().unwrap(), Value::Map(m)).unwrap();
632            let mut buf = Vec::new();
633            tree.write_to(&mut buf).unwrap();
634            let reader = crate::Reader::from_source(buf.clone()).unwrap();
635            #[derive(serde::Deserialize, Debug)]
636            struct TestRecord {
637                ip: String,
638            }
639            let r: TestRecord = reader.lookup("10.0.0.1".parse().unwrap()).unwrap();
640            assert_eq!(r.ip, "test", "failed for record_size={record_size}");
641        }
642    }
643
644    #[test]
645    fn test_write_with_metadata() {
646        let mut desc = BTreeMap::new();
647        desc.insert("en".to_string(), "Test Database".to_string());
648        let opts = WriterOptions {
649            database_type: "TestIP".to_string(),
650            ip_version: 4,
651            record_size: 24,
652            description: desc,
653            languages: vec!["en".to_string()],
654            build_epoch: 1000000,
655        };
656        let mut tree = Writer::new(opts).unwrap();
657        let mut m = BTreeMap::new();
658        m.insert("data".to_string(), Value::String("hello".to_string()));
659        tree.insert_value("192.168.0.0/16".parse().unwrap(), Value::Map(m))
660            .unwrap();
661        let mut buf = Vec::new();
662        tree.write_to(&mut buf).unwrap();
663        let reader = crate::Reader::from_source(buf).unwrap();
664        assert_eq!(reader.metadata.database_type, "TestIP");
665        assert_eq!(reader.metadata.build_epoch, 1000000);
666        assert_eq!(reader.metadata.description["en"], "Test Database");
667    }
668
669    #[test]
670    fn test_bool_value() {
671        let opts = WriterOptions {
672            database_type: "Test".to_string(),
673            ip_version: 4,
674            record_size: 24,
675            ..Default::default()
676        };
677        let mut tree = Writer::new(opts).unwrap();
678        let mut m = BTreeMap::new();
679        m.insert("active".to_string(), Value::Bool(true));
680        m.insert("inactive".to_string(), Value::Bool(false));
681        tree.insert_value("10.0.0.0/8".parse().unwrap(), Value::Map(m)).unwrap();
682        let mut buf = Vec::new();
683        tree.write_to(&mut buf).unwrap();
684        let reader = crate::Reader::from_source(buf).unwrap();
685        #[derive(serde::Deserialize, Debug)]
686        struct TestRecord {
687            active: bool,
688            inactive: bool,
689        }
690        let r: TestRecord = reader.lookup("10.0.0.1".parse().unwrap()).unwrap();
691        assert!(r.active);
692        assert!(!r.inactive);
693    }
694
695    #[test]
696    fn test_slice_value() {
697        let opts = WriterOptions {
698            database_type: "Test".to_string(),
699            ip_version: 4,
700            record_size: 24,
701            ..Default::default()
702        };
703        let mut tree = Writer::new(opts).unwrap();
704        let mut m = BTreeMap::new();
705        m.insert(
706            "tags".to_string(),
707            Value::Slice(vec![
708                Value::String("a".to_string()),
709                Value::String("b".to_string()),
710                Value::String("c".to_string()),
711            ]),
712        );
713        tree.insert_value("10.0.0.0/8".parse().unwrap(), Value::Map(m)).unwrap();
714        let mut buf = Vec::new();
715        tree.write_to(&mut buf).unwrap();
716        let reader = crate::Reader::from_source(buf).unwrap();
717        #[derive(serde::Deserialize, Debug)]
718        struct TestRecord {
719            tags: Vec<String>,
720        }
721        let r: TestRecord = reader.lookup("10.0.0.1".parse().unwrap()).unwrap();
722        assert_eq!(r.tags, vec!["a", "b", "c"]);
723    }
724
725    #[test]
726    fn test_integer_types() {
727        let opts = WriterOptions {
728            database_type: "Test".to_string(),
729            ip_version: 4,
730            record_size: 24,
731            ..Default::default()
732        };
733        let mut tree = Writer::new(opts).unwrap();
734        let mut m = BTreeMap::new();
735        m.insert("u16".to_string(), Value::Uint16(100));
736        m.insert("u32".to_string(), Value::Uint32(100000));
737        m.insert("i32".to_string(), Value::Int32(-42));
738        m.insert("u64".to_string(), Value::Uint64(1 << 40));
739        m.insert("u128".to_string(), Value::Uint128(1u128 << 100));
740        tree.insert_value("10.0.0.0/8".parse().unwrap(), Value::Map(m)).unwrap();
741        let mut buf = Vec::new();
742        tree.write_to(&mut buf).unwrap();
743        let reader = crate::Reader::from_source(buf).unwrap();
744        #[derive(serde::Deserialize, Debug)]
745        struct TestRecord {
746            u16: u16,
747            u32: u32,
748            i32: i32,
749            u64: u64,
750            u128: u128,
751        }
752        let r: TestRecord = reader.lookup("10.0.0.1".parse().unwrap()).unwrap();
753        assert_eq!(r.u16, 100);
754        assert_eq!(r.u32, 100000);
755        assert_eq!(r.i32, -42);
756        assert_eq!(r.u64, 1 << 40);
757        assert_eq!(r.u128, 1u128 << 100);
758    }
759
760    #[test]
761    fn test_pointer_encoding_size3() {
762        let v = Value::Pointer(200000000);
763        let buf = encoder::encode_value(&v).unwrap();
764        assert_eq!(buf[0], 0b00111000, "size=3 pointer control byte should be 0x38");
765        let ptr_val = u32::from_be_bytes([buf[1], buf[2], buf[3], buf[4]]);
766        assert_eq!(ptr_val, 200000000);
767    }
768
769    #[test]
770    fn test_pointer_encoding_size0() {
771        let v = Value::Pointer(500);
772        let buf = encoder::encode_value(&v).unwrap();
773        assert_eq!(buf[0], 0b00100001);
774        assert_eq!(buf.len(), 2);
775        let decoded = ((buf[0] as u32 & 0x07) << 8) | buf[1] as u32;
776        assert_eq!(decoded, 500);
777    }
778
779    #[test]
780    fn test_pointer_encoding_size1() {
781        let v = Value::Pointer(10000);
782        let buf = encoder::encode_value(&v).unwrap();
783        assert_eq!((buf[0] >> 5) & 0x07, 0b001);
784        assert_eq!((buf[0] >> 3) & 0x03, 0b01);
785        assert_eq!(buf.len(), 3);
786    }
787
788    #[test]
789    fn test_empty_database() {
790        let opts = WriterOptions {
791            database_type: "Empty".to_string(),
792            ip_version: 4,
793            record_size: 24,
794            ..Default::default()
795        };
796        let mut tree = Writer::new(opts).unwrap();
797        let mut buf = Vec::new();
798        let n = tree.write_to(&mut buf).unwrap();
799        assert!(n > 0);
800        let reader = crate::Reader::from_source(buf).unwrap();
801        assert_eq!(reader.metadata.database_type, "Empty");
802        assert_eq!(reader.metadata.node_count, 1);
803    }
804
805    #[test]
806    fn test_overlapping_networks() {
807        let opts = WriterOptions {
808            database_type: "Test".to_string(),
809            ip_version: 4,
810            record_size: 24,
811            ..Default::default()
812        };
813        let mut tree = Writer::new(opts).unwrap();
814        let mut m1 = BTreeMap::new();
815        m1.insert("data".to_string(), Value::String("broad".to_string()));
816        tree.insert_value("10.0.0.0/8".parse().unwrap(), Value::Map(m1))
817            .unwrap();
818
819        let mut m2 = BTreeMap::new();
820        m2.insert("data".to_string(), Value::String("specific".to_string()));
821        tree.insert_value("10.1.0.0/16".parse().unwrap(), Value::Map(m2))
822            .unwrap();
823
824        let mut buf = Vec::new();
825        tree.write_to(&mut buf).unwrap();
826        let reader = crate::Reader::from_source(buf).unwrap();
827
828        #[derive(serde::Deserialize, Debug)]
829        struct TestRecord {
830            data: String,
831        }
832
833        let r1: TestRecord = reader.lookup("10.0.0.1".parse().unwrap()).unwrap();
834        assert_eq!(r1.data, "broad");
835
836        let r2: TestRecord = reader.lookup("10.1.0.1".parse().unwrap()).unwrap();
837        assert_eq!(r2.data, "specific");
838    }
839
840    #[test]
841    fn test_nested_map() {
842        let opts = WriterOptions {
843            database_type: "Test".to_string(),
844            ip_version: 4,
845            record_size: 24,
846            ..Default::default()
847        };
848        let mut tree = Writer::new(opts).unwrap();
849        let mut inner = BTreeMap::new();
850        inner.insert("code".to_string(), Value::Uint16(1));
851        inner.insert("name".to_string(), Value::String("US".to_string()));
852
853        let mut m = BTreeMap::new();
854        m.insert("country".to_string(), Value::Map(inner));
855        m.insert("active".to_string(), Value::Bool(true));
856        tree.insert_value("10.0.0.0/8".parse().unwrap(), Value::Map(m)).unwrap();
857
858        let mut buf = Vec::new();
859        tree.write_to(&mut buf).unwrap();
860        let reader = crate::Reader::from_source(buf).unwrap();
861
862        #[derive(serde::Deserialize, Debug)]
863        struct Country {
864            code: u16,
865            name: String,
866        }
867        #[derive(serde::Deserialize, Debug)]
868        struct TestRecord {
869            country: Country,
870            active: bool,
871        }
872
873        let r: TestRecord = reader.lookup("10.0.0.1".parse().unwrap()).unwrap();
874        assert_eq!(r.country.code, 1);
875        assert_eq!(r.country.name, "US");
876        assert!(r.active);
877    }
878
879    #[test]
880    fn test_data_deduplication() {
881        let opts = WriterOptions {
882            database_type: "Test".to_string(),
883            ip_version: 4,
884            record_size: 24,
885            ..Default::default()
886        };
887
888        let mut tree = Writer::new(opts.clone()).unwrap();
889        let shared_data = {
890            let mut m = BTreeMap::new();
891            m.insert("val".to_string(), Value::String("shared".to_string()));
892            Value::Map(m)
893        };
894        tree.insert_value("10.0.0.0/8".parse().unwrap(), shared_data.clone())
895            .unwrap();
896        tree.insert_value("11.0.0.0/8".parse().unwrap(), shared_data).unwrap();
897        let mut buf_dedup = Vec::new();
898        tree.write_to(&mut buf_dedup).unwrap();
899
900        let mut tree2 = Writer::new(opts).unwrap();
901        let mut m1 = BTreeMap::new();
902        m1.insert("val".to_string(), Value::String("first".to_string()));
903        let mut m2 = BTreeMap::new();
904        m2.insert("val".to_string(), Value::String("second".to_string()));
905        tree2
906            .insert_value("10.0.0.0/8".parse().unwrap(), Value::Map(m1))
907            .unwrap();
908        tree2
909            .insert_value("11.0.0.0/8".parse().unwrap(), Value::Map(m2))
910            .unwrap();
911        let mut buf_no_dedup = Vec::new();
912        tree2.write_to(&mut buf_no_dedup).unwrap();
913
914        assert!(
915            buf_dedup.len() < buf_no_dedup.len(),
916            "dedup size {} should be < no-dedup size {}",
917            buf_dedup.len(),
918            buf_no_dedup.len()
919        );
920
921        let r1 = crate::Reader::from_source(buf_dedup).unwrap();
922        let r2 = crate::Reader::from_source(buf_no_dedup).unwrap();
923        assert_eq!(r1.metadata.node_count, r2.metadata.node_count);
924    }
925
926    #[test]
927    fn test_ipv4_in_ipv6_tree() {
928        let opts = WriterOptions {
929            database_type: "Test".to_string(),
930            ip_version: 6,
931            record_size: 28,
932            ..Default::default()
933        };
934        let mut tree = Writer::new(opts).unwrap();
935        let mut m = BTreeMap::new();
936        m.insert("v4in6".to_string(), Value::Bool(true));
937        tree.insert_value("10.0.0.0/8".parse().unwrap(), Value::Map(m)).unwrap();
938
939        let mut buf = Vec::new();
940        tree.write_to(&mut buf).unwrap();
941        let reader = crate::Reader::from_source(buf).unwrap();
942
943        #[derive(serde::Deserialize, Debug)]
944        struct TestRecord {
945            v4in6: bool,
946        }
947
948        let r: TestRecord = reader.lookup("10.0.0.1".parse().unwrap()).unwrap();
949        assert!(r.v4in6);
950    }
951
952    #[test]
953    fn test_record_size_24_node_format() {
954        let opts = WriterOptions {
955            database_type: "Test".to_string(),
956            ip_version: 4,
957            record_size: 24,
958            ..Default::default()
959        };
960        let mut tree = Writer::new(opts).unwrap();
961        let mut m = BTreeMap::new();
962        m.insert("x".to_string(), Value::Uint16(1));
963        tree.insert_value("1.0.0.0/8".parse().unwrap(), Value::Map(m)).unwrap();
964        let mut buf = Vec::new();
965        tree.write_to(&mut buf).unwrap();
966
967        let reader = crate::Reader::from_source(buf.clone()).unwrap();
968        let node_count = reader.metadata.node_count;
969        let nc = node_count as u64;
970        let left = ((buf[0] as u64) << 16) | ((buf[1] as u64) << 8) | (buf[2] as u64);
971        let right = ((buf[3] as u64) << 16) | ((buf[4] as u64) << 8) | (buf[5] as u64);
972        assert!(left <= nc || left >= nc + 16, "left={left} nc={nc}");
973        assert!(right <= nc || right >= nc + 16, "right={right} nc={nc}");
974    }
975
976    #[test]
977    fn test_record_size_28_node_format() {
978        let opts = WriterOptions {
979            database_type: "Test".to_string(),
980            ip_version: 4,
981            record_size: 28,
982            ..Default::default()
983        };
984        let mut tree = Writer::new(opts).unwrap();
985        let mut m = BTreeMap::new();
986        m.insert("x".to_string(), Value::Uint16(1));
987        tree.insert_value("1.0.0.0/8".parse().unwrap(), Value::Map(m)).unwrap();
988        let mut buf = Vec::new();
989        tree.write_to(&mut buf).unwrap();
990
991        let node_count = crate::Reader::from_source(buf.clone()).unwrap().metadata.node_count as usize;
992        assert!(buf.len() > node_count * 7);
993    }
994
995    #[test]
996    fn test_record_size_32_node_format() {
997        let opts = WriterOptions {
998            database_type: "Test".to_string(),
999            ip_version: 4,
1000            record_size: 32,
1001            ..Default::default()
1002        };
1003        let mut tree = Writer::new(opts).unwrap();
1004        let mut m = BTreeMap::new();
1005        m.insert("x".to_string(), Value::Uint16(1));
1006        tree.insert_value("1.0.0.0/8".parse().unwrap(), Value::Map(m)).unwrap();
1007        let mut buf = Vec::new();
1008        tree.write_to(&mut buf).unwrap();
1009
1010        let node_count = crate::Reader::from_source(buf.clone()).unwrap().metadata.node_count as usize;
1011        assert!(buf.len() > node_count * 8);
1012    }
1013
1014    #[test]
1015    fn test_search_tree_size_calculation() {
1016        let opts = WriterOptions {
1017            database_type: "Test".to_string(),
1018            ip_version: 4,
1019            record_size: 24,
1020            ..Default::default()
1021        };
1022        let mut tree = Writer::new(opts).unwrap();
1023        let mut m = BTreeMap::new();
1024        m.insert("x".to_string(), Value::Uint16(1));
1025        tree.insert_value("10.0.0.0/8".parse().unwrap(), Value::Map(m)).unwrap();
1026        let mut buf = Vec::new();
1027        tree.write_to(&mut buf).unwrap();
1028
1029        let reader = crate::Reader::from_source(&buf[..]).unwrap();
1030        let node_count = reader.metadata.node_count as usize;
1031        let record_size = reader.metadata.record_size as usize;
1032
1033        // The search tree occupies node_count * record_size / 4 bytes, and must
1034        // fit in the file ahead of the 16-byte separator and the data section.
1035        let tree_size = node_count * record_size / 4;
1036        assert!(tree_size > 0);
1037        assert!(tree_size + 16 < buf.len());
1038    }
1039
1040    #[test]
1041    fn test_metadata_binary_format_version() {
1042        let opts = WriterOptions {
1043            database_type: "Test".to_string(),
1044            ip_version: 4,
1045            record_size: 24,
1046            ..Default::default()
1047        };
1048        let mut tree = Writer::new(opts).unwrap();
1049        let mut m = BTreeMap::new();
1050        m.insert("x".to_string(), Value::Uint16(1));
1051        tree.insert_value("10.0.0.0/8".parse().unwrap(), Value::Map(m)).unwrap();
1052        let mut buf = Vec::new();
1053        tree.write_to(&mut buf).unwrap();
1054        let reader = crate::Reader::from_source(buf).unwrap();
1055        assert_eq!(reader.metadata.binary_format_major_version, 2);
1056        assert_eq!(reader.metadata.binary_format_minor_version, 0);
1057    }
1058
1059    #[test]
1060    fn test_metadata_description_optional() {
1061        let opts = WriterOptions {
1062            database_type: "Test".to_string(),
1063            ip_version: 4,
1064            record_size: 24,
1065            description: BTreeMap::new(),
1066            ..Default::default()
1067        };
1068        let mut tree = Writer::new(opts).unwrap();
1069        let mut m = BTreeMap::new();
1070        m.insert("x".to_string(), Value::Uint16(1));
1071        tree.insert_value("10.0.0.0/8".parse().unwrap(), Value::Map(m)).unwrap();
1072        let mut buf = Vec::new();
1073        tree.write_to(&mut buf).unwrap();
1074        let reader = crate::Reader::from_source(buf).unwrap();
1075        assert!(reader.metadata.description.is_empty());
1076    }
1077
1078    #[test]
1079    fn test_languages_optional() {
1080        let opts = WriterOptions {
1081            database_type: "Test".to_string(),
1082            ip_version: 4,
1083            record_size: 24,
1084            languages: vec![],
1085            ..Default::default()
1086        };
1087        let mut tree = Writer::new(opts).unwrap();
1088        let mut m = BTreeMap::new();
1089        m.insert("x".to_string(), Value::Uint16(1));
1090        tree.insert_value("10.0.0.0/8".parse().unwrap(), Value::Map(m)).unwrap();
1091        let mut buf = Vec::new();
1092        tree.write_to(&mut buf).unwrap();
1093        let reader = crate::Reader::from_source(buf).unwrap();
1094        assert!(reader.metadata.languages.is_empty());
1095    }
1096
1097    #[test]
1098    fn test_insert_multiple_cidrs_then_lookup() {
1099        let opts = WriterOptions {
1100            database_type: "Test".to_string(),
1101            ip_version: 4,
1102            record_size: 24,
1103            ..Default::default()
1104        };
1105        let mut tree = Writer::new(opts).unwrap();
1106        for i in 0u8..10 {
1107            let cidr = format!("{i}.0.0.0/8");
1108            let mut m = BTreeMap::new();
1109            m.insert("net".to_string(), Value::String(cidr.clone()));
1110            tree.insert_value(cidr.parse().unwrap(), Value::Map(m)).unwrap();
1111        }
1112        let mut buf = Vec::new();
1113        tree.write_to(&mut buf).unwrap();
1114        let reader = crate::Reader::from_source(buf).unwrap();
1115        #[derive(serde::Deserialize, Debug)]
1116        struct TestRecord {
1117            net: String,
1118        }
1119        for i in 0u8..10 {
1120            let ip = format!("{i}.0.0.1");
1121            let r: TestRecord = reader.lookup(ip.parse().unwrap()).unwrap();
1122            assert_eq!(r.net, format!("{i}.0.0.0/8"));
1123        }
1124    }
1125
1126    #[test]
1127    fn test_large_string_value() {
1128        let opts = WriterOptions {
1129            database_type: "Test".to_string(),
1130            ip_version: 4,
1131            record_size: 24,
1132            ..Default::default()
1133        };
1134        let mut tree = Writer::new(opts).unwrap();
1135        let long_str = "x".repeat(100);
1136        let mut m = BTreeMap::new();
1137        m.insert("long".to_string(), Value::String(long_str.clone()));
1138        tree.insert_value("10.0.0.0/8".parse().unwrap(), Value::Map(m)).unwrap();
1139        let mut buf = Vec::new();
1140        tree.write_to(&mut buf).unwrap();
1141        let reader = crate::Reader::from_source(buf).unwrap();
1142        #[derive(serde::Deserialize, Debug)]
1143        struct TestRecord {
1144            long: String,
1145        }
1146        let r: TestRecord = reader.lookup("10.0.0.1".parse().unwrap()).unwrap();
1147        assert_eq!(r.long, long_str);
1148    }
1149
1150    #[test]
1151    fn test_metadata_start_marker_last_occurrence() {
1152        let opts = WriterOptions {
1153            database_type: "Test".to_string(),
1154            ip_version: 4,
1155            record_size: 24,
1156            ..Default::default()
1157        };
1158        let mut tree = Writer::new(opts).unwrap();
1159        let mut m = BTreeMap::new();
1160        m.insert("x".to_string(), Value::Uint16(1));
1161        tree.insert_value("10.0.0.0/8".parse().unwrap(), Value::Map(m)).unwrap();
1162        let mut buf = Vec::new();
1163        tree.write_to(&mut buf).unwrap();
1164
1165        let marker = b"\xab\xcd\xefMaxMind.com";
1166        let all_positions: Vec<usize> = buf
1167            .windows(marker.len())
1168            .enumerate()
1169            .filter(|(_, w)| *w == marker)
1170            .map(|(i, _)| i)
1171            .collect();
1172        assert_eq!(all_positions.len(), 1);
1173        let last_pos = all_positions[0];
1174        assert!(last_pos + marker.len() <= buf.len());
1175    }
1176
1177    #[test]
1178    fn test_write_to_is_idempotent() {
1179        let opts = WriterOptions {
1180            database_type: "Test".to_string(),
1181            ip_version: 4,
1182            record_size: 24,
1183            ..Default::default()
1184        };
1185        let mut tree = Writer::new(opts).unwrap();
1186        let mut m = BTreeMap::new();
1187        m.insert("x".to_string(), Value::Uint16(42));
1188        tree.insert_value("10.0.0.0/8".parse().unwrap(), Value::Map(m)).unwrap();
1189
1190        let mut buf1 = Vec::new();
1191        let n1 = tree.write_to(&mut buf1).unwrap();
1192
1193        let mut buf2 = Vec::new();
1194        let n2 = tree.write_to(&mut buf2).unwrap();
1195
1196        assert_eq!(n1, n2);
1197        assert_eq!(buf1.len(), buf2.len());
1198        assert_eq!(buf1, buf2);
1199    }
1200
1201    #[test]
1202    fn test_zero_length_prefix() {
1203        let opts = WriterOptions {
1204            database_type: "Test".to_string(),
1205            ip_version: 4,
1206            record_size: 24,
1207            ..Default::default()
1208        };
1209        let mut tree = Writer::new(opts).unwrap();
1210        let mut m = BTreeMap::new();
1211        m.insert("all".to_string(), Value::Bool(true));
1212        tree.insert_value("0.0.0.0/0".parse().unwrap(), Value::Map(m)).unwrap();
1213        let mut buf = Vec::new();
1214        tree.write_to(&mut buf).unwrap();
1215        let reader = crate::Reader::from_source(buf).unwrap();
1216        #[derive(serde::Deserialize, Debug)]
1217        struct TestRecord {
1218            all: bool,
1219        }
1220        let r: TestRecord = reader.lookup("255.255.255.255".parse().unwrap()).unwrap();
1221        assert!(r.all);
1222    }
1223
1224    #[test]
1225    fn test_exact_prefix_length() {
1226        let opts = WriterOptions {
1227            database_type: "Test".to_string(),
1228            ip_version: 4,
1229            record_size: 24,
1230            ..Default::default()
1231        };
1232        let mut tree = Writer::new(opts).unwrap();
1233        let mut m = BTreeMap::new();
1234        m.insert("host".to_string(), Value::String("single".to_string()));
1235        tree.insert_value("10.0.0.1/32".parse().unwrap(), Value::Map(m))
1236            .unwrap();
1237        let mut buf = Vec::new();
1238        tree.write_to(&mut buf).unwrap();
1239        let reader = crate::Reader::from_source(buf).unwrap();
1240        #[derive(serde::Deserialize, Debug)]
1241        struct TestRecord {
1242            host: String,
1243        }
1244        let r: TestRecord = reader.lookup("10.0.0.1".parse().unwrap()).unwrap();
1245        assert_eq!(r.host, "single");
1246        assert!(reader.lookup::<TestRecord>("10.0.0.2".parse().unwrap()).is_err());
1247    }
1248
1249    #[test]
1250    fn test_negative_int32() {
1251        let opts = WriterOptions {
1252            database_type: "Test".to_string(),
1253            ip_version: 4,
1254            record_size: 24,
1255            ..Default::default()
1256        };
1257        let mut tree = Writer::new(opts).unwrap();
1258        let mut m = BTreeMap::new();
1259        m.insert("neg".to_string(), Value::Int32(-1));
1260        m.insert("min".to_string(), Value::Int32(i32::MIN));
1261        m.insert("max".to_string(), Value::Int32(i32::MAX));
1262        tree.insert_value("10.0.0.0/8".parse().unwrap(), Value::Map(m)).unwrap();
1263        let mut buf = Vec::new();
1264        tree.write_to(&mut buf).unwrap();
1265        let reader = crate::Reader::from_source(buf).unwrap();
1266        #[derive(serde::Deserialize, Debug)]
1267        struct TestRecord {
1268            neg: i32,
1269            min: i32,
1270            max: i32,
1271        }
1272        let r: TestRecord = reader.lookup("10.0.0.1".parse().unwrap()).unwrap();
1273        assert_eq!(r.neg, -1);
1274        assert_eq!(r.min, i32::MIN);
1275        assert_eq!(r.max, i32::MAX);
1276    }
1277
1278    #[test]
1279    fn test_float32_value() {
1280        let opts = WriterOptions {
1281            database_type: "Test".to_string(),
1282            ip_version: 4,
1283            record_size: 24,
1284            ..Default::default()
1285        };
1286        let mut tree = Writer::new(opts).unwrap();
1287        let mut m = BTreeMap::new();
1288        m.insert("f".to_string(), Value::Float32(1.5));
1289        tree.insert_value("10.0.0.0/8".parse().unwrap(), Value::Map(m)).unwrap();
1290        let mut buf = Vec::new();
1291        tree.write_to(&mut buf).unwrap();
1292        let reader = crate::Reader::from_source(buf).unwrap();
1293        #[derive(serde::Deserialize, Debug)]
1294        struct TestRecord {
1295            f: f32,
1296        }
1297        let r: TestRecord = reader.lookup("10.0.0.1".parse().unwrap()).unwrap();
1298        assert!((r.f - 1.5).abs() < 1e-6);
1299    }
1300
1301    #[test]
1302    fn test_map_with_many_keys() {
1303        let opts = WriterOptions {
1304            database_type: "Test".to_string(),
1305            ip_version: 4,
1306            record_size: 24,
1307            ..Default::default()
1308        };
1309        let mut tree = Writer::new(opts).unwrap();
1310        let mut m = BTreeMap::new();
1311        for i in 0..20 {
1312            m.insert(format!("key{i}").to_string(), Value::Uint16(i as u16));
1313        }
1314        tree.insert_value("10.0.0.0/8".parse().unwrap(), Value::Map(m)).unwrap();
1315        let mut buf = Vec::new();
1316        tree.write_to(&mut buf).unwrap();
1317        let reader = crate::Reader::from_source(buf).unwrap();
1318        #[derive(serde::Deserialize, Debug)]
1319        struct ManyKeys {}
1320        let _: ManyKeys = reader.lookup("10.0.0.1".parse().unwrap()).unwrap();
1321    }
1322
1323    #[test]
1324    fn test_write_error_on_invalid_record_size() {
1325        let opts = WriterOptions {
1326            database_type: "Test".to_string(),
1327            ip_version: 4,
1328            record_size: 16,
1329            ..Default::default()
1330        };
1331        let result = Writer::new(opts);
1332        assert!(result.is_err());
1333    }
1334
1335    #[test]
1336    fn test_write_error_on_invalid_ip_version() {
1337        let result = Writer::new(WriterOptions {
1338            database_type: "Test".to_string(),
1339            ip_version: 5,
1340            record_size: 24,
1341            ..Default::default()
1342        });
1343        assert!(result.is_err());
1344        assert!(result.unwrap_err().contains("ip_version"));
1345    }
1346
1347    #[test]
1348    fn test_write_error_on_invalid_record_size_in_new() {
1349        let result = Writer::new(WriterOptions {
1350            database_type: "Test".to_string(),
1351            ip_version: 4,
1352            record_size: 20,
1353            ..Default::default()
1354        });
1355        assert!(result.is_err());
1356        assert!(result.unwrap_err().contains("record_size"));
1357    }
1358
1359    #[test]
1360    fn test_write_error_on_record_size_36() {
1361        let result = Writer::new(WriterOptions {
1362            database_type: "Test".to_string(),
1363            ip_version: 4,
1364            record_size: 36,
1365            ..Default::default()
1366        });
1367        assert!(result.is_err());
1368        assert!(result.unwrap_err().contains("record_size"));
1369    }
1370
1371    #[test]
1372    fn test_write_error_ipv6_in_ipv4_tree() {
1373        let mut tree = Writer::new(WriterOptions {
1374            database_type: "Test".to_string(),
1375            ip_version: 4,
1376            record_size: 24,
1377            ..Default::default()
1378        })
1379        .unwrap();
1380        let result = tree.insert_value("::1/128".parse().unwrap(), Value::Bool(true));
1381        assert!(result.is_err());
1382        assert!(result.unwrap_err().contains("IPv6"));
1383    }
1384
1385    #[test]
1386    fn test_float64_value() {
1387        let opts = WriterOptions {
1388            database_type: "Test".to_string(),
1389            ip_version: 4,
1390            record_size: 24,
1391            ..Default::default()
1392        };
1393        let mut tree = Writer::new(opts).unwrap();
1394        let mut m = BTreeMap::new();
1395        m.insert("f64".to_string(), Value::Float64(std::f64::consts::PI));
1396        tree.insert_value("10.0.0.0/8".parse().unwrap(), Value::Map(m)).unwrap();
1397        let mut buf = Vec::new();
1398        tree.write_to(&mut buf).unwrap();
1399        let reader = crate::Reader::from_source(buf).unwrap();
1400        #[derive(serde::Deserialize, Debug)]
1401        struct TestRecord {
1402            f64: f64,
1403        }
1404        let r: TestRecord = reader.lookup("10.0.0.1".parse().unwrap()).unwrap();
1405        assert!((r.f64 - std::f64::consts::PI).abs() < 1e-15);
1406    }
1407
1408    #[test]
1409    fn test_bytes_value() {
1410        let opts = WriterOptions {
1411            database_type: "Test".to_string(),
1412            ip_version: 4,
1413            record_size: 24,
1414            ..Default::default()
1415        };
1416        let mut tree = Writer::new(opts).unwrap();
1417        let mut m = BTreeMap::new();
1418        m.insert("raw".to_string(), Value::Bytes(vec![0xde, 0xad, 0xbe, 0xef]));
1419        tree.insert_value("10.0.0.0/8".parse().unwrap(), Value::Map(m)).unwrap();
1420        let mut buf = Vec::new();
1421        tree.write_to(&mut buf).unwrap();
1422        let reader = crate::Reader::from_source(buf).unwrap();
1423        assert!(reader.metadata.node_count > 0);
1424    }
1425
1426    #[test]
1427    fn test_pointer_encoding_size2() {
1428        let v = Value::Pointer(1000000);
1429        let buf = encoder::encode_value(&v).unwrap();
1430        assert_eq!((buf[0] >> 5) & 0x07, 0b001);
1431        assert_eq!((buf[0] >> 3) & 0x03, 0b10);
1432        assert_eq!(buf.len(), 4);
1433    }
1434
1435    #[test]
1436    fn test_metadata_size_limit() {
1437        let mut desc = BTreeMap::new();
1438        desc.insert("en".to_string(), "x".repeat(65536));
1439        let opts = WriterOptions {
1440            database_type: "LargeMeta".to_string(),
1441            ip_version: 4,
1442            record_size: 24,
1443            description: desc,
1444            ..Default::default()
1445        };
1446        let mut tree = Writer::new(opts).unwrap();
1447        let mut m = BTreeMap::new();
1448        m.insert("x".to_string(), Value::Uint16(1));
1449        tree.insert_value("10.0.0.0/8".parse().unwrap(), Value::Map(m)).unwrap();
1450        let mut buf = Vec::new();
1451        tree.write_to(&mut buf).unwrap();
1452        let marker = b"\xab\xcd\xefMaxMind.com";
1453        let marker_pos = buf.windows(marker.len()).rposition(|w| w == marker).unwrap();
1454        let metadata_section_size = buf.len() - marker_pos;
1455        assert!(
1456            metadata_section_size <= 128 * 1024,
1457            "metadata section {} bytes exceeds 128 KiB limit",
1458            metadata_section_size
1459        );
1460    }
1461
1462    #[test]
1463    fn test_encoded_size_matches_actual() {
1464        let cases: Vec<Value> = vec![
1465            Value::Pointer(0),
1466            Value::Pointer(2047),
1467            Value::Pointer(2048),
1468            Value::Pointer(526335),
1469            Value::Pointer(526336),
1470            Value::Pointer(134744063),
1471            Value::Pointer(134744064),
1472            Value::Pointer(u32::MAX),
1473            Value::String("hello".to_string()),
1474            Value::String("x".repeat(29)),
1475            Value::String("x".repeat(300)),
1476            Value::Float64(0.0),
1477            Value::Float64(-1.5),
1478            Value::Float64(f64::MAX),
1479            Value::Bytes(vec![]),
1480            Value::Bytes(vec![0; 42]),
1481            Value::Uint16(0),
1482            Value::Uint16(1),
1483            Value::Uint16(u16::MAX),
1484            Value::Uint32(0),
1485            Value::Uint32(1 << 20),
1486            Value::Uint32(u32::MAX),
1487            Value::Int32(0),
1488            Value::Int32(1),
1489            Value::Int32(-1),
1490            Value::Int32(i32::MIN),
1491            Value::Int32(i32::MAX),
1492            Value::Uint64(0),
1493            Value::Uint64(1 << 56),
1494            Value::Uint64(u64::MAX),
1495            Value::Uint128(0),
1496            Value::Uint128(1u128 << 120),
1497            Value::Uint128(u128::MAX),
1498            Value::Bool(true),
1499            Value::Bool(false),
1500            Value::Float32(0.0),
1501            Value::Float32(1.25),
1502            Value::Float32(f32::MAX),
1503        ];
1504
1505        for v in cases {
1506            let encoded = encoder::encode_value(&v).unwrap();
1507            let size = encoder::encoded_size(&v);
1508            assert_eq!(
1509                encoded.len(),
1510                size,
1511                "encoded_size mismatch for {v:?}: got {} encoded but size said {size}",
1512                encoded.len()
1513            );
1514        }
1515    }
1516
1517    #[test]
1518    fn test_encoded_size_matches_actual_empty_string() {
1519        let v = Value::String(String::new());
1520        let encoded = encoder::encode_value(&v).unwrap();
1521        assert_eq!(encoded.len(), encoder::encoded_size(&v));
1522    }
1523
1524    #[test]
1525    fn test_values_after_insert_and_write_are_preserved() {
1526        let v = Value::String("hello".to_string());
1527
1528        let opts = WriterOptions {
1529            database_type: "Test".to_string(),
1530            ip_version: 4,
1531            record_size: 24,
1532            ..Default::default()
1533        };
1534        let mut tree = Writer::new(opts).unwrap();
1535        tree.insert_value("0.0.0.0/0".parse().unwrap(), v.clone()).unwrap();
1536        let mut buf = Vec::new();
1537        tree.write_to(&mut buf).unwrap();
1538        let reader = crate::Reader::from_source(buf).unwrap();
1539        let val: Value = reader.lookup("1.2.3.4".parse().unwrap()).unwrap();
1540        assert_eq!(val, v);
1541    }
1542
1543    #[test]
1544    fn test_u128_roundtrip() {
1545        let opts = WriterOptions {
1546            database_type: "Test".to_string(),
1547            ip_version: 4,
1548            record_size: 24,
1549            ..Default::default()
1550        };
1551        let mut tree = Writer::new(opts).unwrap();
1552        let mut m = BTreeMap::new();
1553        let big = 0xdeadbeefcafebabeu128 << 64 | 0x1234567890abcdefu128;
1554        m.insert("v".to_string(), Value::Uint128(big));
1555        tree.insert_value("10.0.0.0/8".parse().unwrap(), Value::Map(m)).unwrap();
1556        let mut buf = Vec::new();
1557        tree.write_to(&mut buf).unwrap();
1558        let reader = crate::Reader::from_source(buf).unwrap();
1559        #[derive(serde::Deserialize, Debug)]
1560        struct R {
1561            v: u128,
1562        }
1563        let r: R = reader.lookup("10.0.0.1".parse().unwrap()).unwrap();
1564        assert_eq!(r.v, big);
1565    }
1566
1567    #[test]
1568    fn test_zero_values() {
1569        let opts = WriterOptions {
1570            database_type: "Test".to_string(),
1571            ip_version: 4,
1572            record_size: 24,
1573            ..Default::default()
1574        };
1575        let mut tree = Writer::new(opts).unwrap();
1576        let mut m = BTreeMap::new();
1577        m.insert("u16".to_string(), Value::Uint16(0));
1578        m.insert("u32".to_string(), Value::Uint32(0));
1579        m.insert("u64".to_string(), Value::Uint64(0));
1580        m.insert("u128".to_string(), Value::Uint128(0));
1581        tree.insert_value("10.0.0.0/8".parse().unwrap(), Value::Map(m)).unwrap();
1582        let mut buf = Vec::new();
1583        tree.write_to(&mut buf).unwrap();
1584        let reader = crate::Reader::from_source(buf).unwrap();
1585        #[derive(serde::Deserialize, Debug)]
1586        struct Zeros {
1587            u16: u16,
1588            u32: u32,
1589            u64: u64,
1590            u128: u128,
1591        }
1592        let r: Zeros = reader.lookup("10.0.0.1".parse().unwrap()).unwrap();
1593        assert_eq!(r.u16, 0);
1594        assert_eq!(r.u32, 0);
1595        assert_eq!(r.u64, 0);
1596        assert_eq!(r.u128, 0);
1597    }
1598
1599    #[test]
1600    fn test_insert_with_serialize_struct() {
1601        #[derive(serde::Serialize)]
1602        struct Record {
1603            country: String,
1604        }
1605
1606        let opts = WriterOptions {
1607            database_type: "Test".to_string(),
1608            ip_version: 4,
1609            record_size: 24,
1610            ..Default::default()
1611        };
1612        let mut tree = Writer::new(opts).unwrap();
1613        tree.insert(
1614            "1.0.0.0/8".parse().unwrap(),
1615            Record {
1616                country: "US".into(),
1617            },
1618        )
1619        .unwrap();
1620        let mut buf = Vec::new();
1621        tree.write_to(&mut buf).unwrap();
1622        let reader = crate::Reader::from_source(buf).unwrap();
1623        #[derive(serde::Deserialize, Debug)]
1624        struct Record2 {
1625            country: String,
1626        }
1627        let result: Record2 = reader.lookup("1.0.0.1".parse().unwrap()).unwrap();
1628        assert_eq!(result.country, "US");
1629    }
1630
1631    #[test]
1632    fn test_insert_with_serialize_btreemap() {
1633        let opts = WriterOptions {
1634            database_type: "Test".to_string(),
1635            ip_version: 4,
1636            record_size: 24,
1637            ..Default::default()
1638        };
1639        let mut tree = Writer::new(opts).unwrap();
1640        let mut map = BTreeMap::new();
1641        map.insert("value".to_string(), "hello".to_string());
1642        tree.insert("10.0.0.0/8".parse().unwrap(), map).unwrap();
1643        let mut buf = Vec::new();
1644        tree.write_to(&mut buf).unwrap();
1645        let reader = crate::Reader::from_source(buf).unwrap();
1646        #[derive(serde::Deserialize, Debug)]
1647        struct R {
1648            value: String,
1649        }
1650        let result: R = reader.lookup("10.0.0.1".parse().unwrap()).unwrap();
1651        assert_eq!(result.value, "hello");
1652    }
1653
1654    #[test]
1655    fn test_insert_rejects_negative_out_of_range_integer() {
1656        let opts = WriterOptions {
1657            database_type: "Test".to_string(),
1658            ip_version: 4,
1659            record_size: 24,
1660            ..Default::default()
1661        };
1662        let mut tree = Writer::new(opts).unwrap();
1663        let err = tree.insert("10.0.0.0/8".parse().unwrap(), i64::MIN).unwrap_err();
1664        assert!(err.contains("out of range"), "unexpected error: {err}");
1665    }
1666
1667    #[test]
1668    fn test_insert_accepts_large_positive_integer_as_u64() {
1669        let opts = WriterOptions {
1670            database_type: "Test".to_string(),
1671            ip_version: 4,
1672            record_size: 24,
1673            ..Default::default()
1674        };
1675        let mut tree = Writer::new(opts).unwrap();
1676        let big = i64::MAX;
1677        tree.insert("10.0.0.0/8".parse().unwrap(), big).unwrap();
1678        let mut buf = Vec::new();
1679        tree.write_to(&mut buf).unwrap();
1680
1681        let reader = crate::Reader::from_source(buf).unwrap();
1682        // The value is encoded as a bare u64, so decode it as such.
1683        let decoded: u64 = reader.lookup("10.0.0.1".parse().unwrap()).unwrap();
1684        assert_eq!(decoded, big as u64);
1685    }
1686
1687    #[test]
1688    fn test_encode_rejects_oversized_payload() {
1689        let too_big = crate::encoder::Value::String("x".repeat(16_843_037));
1690        let err = crate::encoder::encode_value(&too_big).unwrap_err();
1691        assert!(err.contains("too large"), "unexpected error: {err}");
1692    }
1693
1694    #[test]
1695    fn test_metadata_size_limit_is_enforced() {
1696        let mut description = BTreeMap::new();
1697        // One description string above the 128 KiB metadata budget.
1698        description.insert("en".to_string(), "x".repeat(200 * 1024));
1699        let opts = WriterOptions {
1700            database_type: "Test".to_string(),
1701            ip_version: 4,
1702            record_size: 24,
1703            description,
1704            languages: vec!["en".to_string()],
1705            build_epoch: 1,
1706        };
1707        let mut tree = Writer::new(opts).unwrap();
1708        tree.insert("10.0.0.0/8".parse().unwrap(), "v").unwrap();
1709        let err = tree.write_to(&mut Vec::new()).unwrap_err();
1710        assert!(err.contains("metadata"), "unexpected error: {err}");
1711    }
1712}
maxminddb/writer.rs

maxminddb/
writer.rs
