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