maxminddb/
encoder.rs

1use std::{collections::BTreeMap, fmt};
2
3use serde::{
4    de::{self, Deserialize, Deserializer, MapAccess, SeqAccess, Visitor},
5    ser::{self, Serialize, SerializeMap, SerializeSeq, SerializeStruct, SerializeTupleStruct},
6};
7
8pub type EncodeResult<T> = Result<T, String>;
9
10/// Error returned while encoding a value for a MaxMind DB.
11#[derive(Debug)]
12pub struct SerError(pub String);
13
14impl fmt::Display for SerError {
15    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
16        f.write_str(&self.0)
17    }
18}
19
20impl core::error::Error for SerError {}
21
22impl ser::Error for SerError {
23    fn custom<T: fmt::Display>(msg: T) -> Self {
24        SerError(msg.to_string())
25    }
26}
27
28const FIRST_SIZE: usize = 29;
29const SECOND_SIZE: usize = FIRST_SIZE + 256;
30const THIRD_SIZE: usize = SECOND_SIZE + (1 << 16);
31
32/// A MaxMind DB value.
33///
34/// This is a fully owned representation. When reading, prefer deserializing
35/// into a borrowed type (for example the structs in [`crate::geoip2`]) to avoid
36/// copying strings and byte arrays out of the database.
37///
38/// [`Value::Pointer`] is only meaningful while encoding: the numeric value is a
39/// raw data-section pointer. Use it only if you are building pointer-based
40/// databases by hand.
41#[derive(Debug, Clone, PartialEq)]
42pub enum Value {
43    Pointer(u32),
44    String(String),
45    Float64(f64),
46    Bytes(Vec<u8>),
47    Uint16(u16),
48    Uint32(u32),
49    Int32(i32),
50    Map(BTreeMap<String, Value>),
51    Uint64(u64),
52    Uint128(u128),
53    Slice(Vec<Value>),
54    Bool(bool),
55    Float32(f32),
56}
57
58impl From<&str> for Value {
59    fn from(s: &str) -> Self {
60        Value::String(s.to_string())
61    }
62}
63
64impl From<String> for Value {
65    fn from(s: String) -> Self {
66        Value::String(s)
67    }
68}
69
70impl From<Vec<u8>> for Value {
71    fn from(v: Vec<u8>) -> Self {
72        Value::Bytes(v)
73    }
74}
75
76impl From<u16> for Value {
77    fn from(v: u16) -> Self {
78        Value::Uint16(v)
79    }
80}
81
82impl From<u32> for Value {
83    fn from(v: u32) -> Self {
84        Value::Uint32(v)
85    }
86}
87
88impl From<i32> for Value {
89    fn from(v: i32) -> Self {
90        Value::Int32(v)
91    }
92}
93
94impl From<u64> for Value {
95    fn from(v: u64) -> Self {
96        Value::Uint64(v)
97    }
98}
99
100impl From<u128> for Value {
101    fn from(v: u128) -> Self {
102        Value::Uint128(v)
103    }
104}
105
106impl From<f64> for Value {
107    fn from(v: f64) -> Self {
108        Value::Float64(v)
109    }
110}
111
112impl From<f32> for Value {
113    fn from(v: f32) -> Self {
114        Value::Float32(v)
115    }
116}
117
118impl From<bool> for Value {
119    fn from(v: bool) -> Self {
120        Value::Bool(v)
121    }
122}
123
124/// Build a [`Value::Map`] from `key => value` pairs.
125///
126/// Keys are converted with `to_string`; values with `Into<Value>`.
127#[macro_export]
128macro_rules! map {
129    ($($key:expr => $val:expr),* $(,)?) => {{
130        let mut __map = ::std::collections::BTreeMap::new();
131        $(
132            __map.insert(
133                ($key).to_string(),
134                ::std::convert::Into::<$crate::encoder::Value>::into($val),
135            );
136        )*
137        $crate::encoder::Value::Map(__map)
138    }};
139}
140
141impl<'de> Deserialize<'de> for Value {
142    fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
143        struct ValueVisitor;
144
145        impl<'de> Visitor<'de> for ValueVisitor {
146            type Value = Value;
147
148            fn expecting(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
149                f.write_str("a MaxMind DB value")
150            }
151
152            fn visit_bool<E: de::Error>(self, v: bool) -> Result<Value, E> {
153                Ok(Value::Bool(v))
154            }
155
156            fn visit_i64<E: de::Error>(self, v: i64) -> Result<Value, E> {
157                if let Ok(v) = i32::try_from(v) {
158                    Ok(Value::Int32(v))
159                } else if v >= 0 {
160                    Ok(Value::Uint64(v as u64))
161                } else {
162                    Err(E::custom(format!(
163                        "negative integer {v} is out of range for a MaxMind DB value (minimum is {})",
164                        i32::MIN
165                    )))
166                }
167            }
168
169            fn visit_u64<E: de::Error>(self, v: u64) -> Result<Value, E> {
170                Ok(Value::Uint64(v))
171            }
172
173            fn visit_f64<E: de::Error>(self, v: f64) -> Result<Value, E> {
174                Ok(Value::Float64(v))
175            }
176
177            fn visit_str<E: de::Error>(self, v: &str) -> Result<Value, E> {
178                Ok(Value::String(v.to_string()))
179            }
180
181            fn visit_string<E: de::Error>(self, v: String) -> Result<Value, E> {
182                Ok(Value::String(v))
183            }
184
185            fn visit_bytes<E: de::Error>(self, v: &[u8]) -> Result<Value, E> {
186                Ok(Value::Bytes(v.to_vec()))
187            }
188
189            fn visit_byte_buf<E: de::Error>(self, v: Vec<u8>) -> Result<Value, E> {
190                Ok(Value::Bytes(v))
191            }
192
193            fn visit_seq<A: SeqAccess<'de>>(self, mut seq: A) -> Result<Value, A::Error> {
194                let mut items = Vec::with_capacity(seq.size_hint().unwrap_or(1));
195                while let Some(elem) = seq.next_element::<Value>()? {
196                    items.push(elem);
197                }
198                Ok(Value::Slice(items))
199            }
200
201            fn visit_map<A: MapAccess<'de>>(self, mut map: A) -> Result<Value, A::Error> {
202                let mut entries = BTreeMap::new();
203                while let Some((key, val)) = map.next_entry::<String, Value>()? {
204                    entries.insert(key, val);
205                }
206                Ok(Value::Map(entries))
207            }
208        }
209
210        deserializer.deserialize_any(ValueVisitor)
211    }
212}
213
214// ---------------------------------------------------------------------------
215// ValueSerializer — converts Serialize into owned Value (used by insert_value)
216// ---------------------------------------------------------------------------
217
218/// Serializes a value into an owned [`Value`].
219pub struct ValueSerializer;
220
221impl ser::Serializer for &ValueSerializer {
222    type Ok = Value;
223    type Error = SerError;
224
225    type SerializeSeq = OwnedMapOrSlice;
226    type SerializeTuple = OwnedMapOrSlice;
227    type SerializeTupleStruct = OwnedMapOrSlice;
228    type SerializeTupleVariant = ser::Impossible<Value, SerError>;
229    type SerializeMap = OwnedMapOrSlice;
230    type SerializeStruct = OwnedMapOrSlice;
231    type SerializeStructVariant = ser::Impossible<Value, SerError>;
232
233    fn serialize_bool(self, v: bool) -> Result<Value, SerError> {
234        Ok(Value::Bool(v))
235    }
236
237    fn serialize_i8(self, v: i8) -> Result<Value, SerError> {
238        Ok(Value::Int32(v as i32))
239    }
240
241    fn serialize_i16(self, v: i16) -> Result<Value, SerError> {
242        Ok(Value::Int32(v as i32))
243    }
244
245    fn serialize_i32(self, v: i32) -> Result<Value, SerError> {
246        Ok(Value::Int32(v))
247    }
248
249    fn serialize_i64(self, v: i64) -> Result<Value, SerError> {
250        if let Ok(v) = i32::try_from(v) {
251            Ok(Value::Int32(v))
252        } else if v >= 0 {
253            Ok(Value::Uint64(v as u64))
254        } else {
255            Err(SerError(format!(
256                "negative integer {v} is out of range for a MaxMind DB value (minimum is {})",
257                i32::MIN
258            )))
259        }
260    }
261
262    fn serialize_u8(self, v: u8) -> Result<Value, SerError> {
263        Ok(Value::Uint16(v as u16))
264    }
265
266    fn serialize_u16(self, v: u16) -> Result<Value, SerError> {
267        Ok(Value::Uint16(v))
268    }
269
270    fn serialize_u32(self, v: u32) -> Result<Value, SerError> {
271        Ok(Value::Uint32(v))
272    }
273
274    fn serialize_u64(self, v: u64) -> Result<Value, SerError> {
275        Ok(Value::Uint64(v))
276    }
277
278    fn serialize_u128(self, v: u128) -> Result<Value, SerError> {
279        Ok(Value::Uint128(v))
280    }
281
282    fn serialize_f32(self, v: f32) -> Result<Value, SerError> {
283        Ok(Value::Float32(v))
284    }
285
286    fn serialize_f64(self, v: f64) -> Result<Value, SerError> {
287        Ok(Value::Float64(v))
288    }
289
290    fn serialize_char(self, v: char) -> Result<Value, SerError> {
291        Ok(Value::String(v.to_string()))
292    }
293
294    fn serialize_str(self, v: &str) -> Result<Value, SerError> {
295        Ok(Value::String(v.to_string()))
296    }
297
298    fn serialize_bytes(self, v: &[u8]) -> Result<Value, SerError> {
299        Ok(Value::Bytes(v.to_vec()))
300    }
301
302    fn serialize_none(self) -> Result<Value, SerError> {
303        Err(SerError("MaxMind DB has no null type".into()))
304    }
305
306    fn serialize_some<T: ?Sized + Serialize>(self, value: &T) -> Result<Value, SerError> {
307        value.serialize(self)
308    }
309
310    fn serialize_unit(self) -> Result<Value, SerError> {
311        Err(SerError("MaxMind DB has no null type".into()))
312    }
313
314    fn serialize_unit_struct(self, _name: &'static str) -> Result<Value, SerError> {
315        Err(SerError("MaxMind DB has no null type".into()))
316    }
317
318    fn serialize_unit_variant(
319        self,
320        _name: &'static str,
321        _variant_index: u32,
322        variant: &'static str,
323    ) -> Result<Value, SerError> {
324        Ok(Value::String(variant.to_string()))
325    }
326
327    fn serialize_newtype_struct<T: ?Sized + Serialize>(
328        self,
329        _name: &'static str,
330        value: &T,
331    ) -> Result<Value, SerError> {
332        value.serialize(self)
333    }
334
335    fn serialize_newtype_variant<T: ?Sized + Serialize>(
336        self,
337        _name: &'static str,
338        _variant_index: u32,
339        _variant: &'static str,
340        value: &T,
341    ) -> Result<Value, SerError> {
342        value.serialize(self)
343    }
344
345    fn serialize_seq(self, _len: Option<usize>) -> Result<OwnedMapOrSlice, SerError> {
346        Ok(OwnedMapOrSlice::slice())
347    }
348
349    fn serialize_tuple(self, len: usize) -> Result<OwnedMapOrSlice, SerError> {
350        self.serialize_seq(Some(len))
351    }
352
353    fn serialize_tuple_struct(self, _name: &'static str, len: usize) -> Result<OwnedMapOrSlice, SerError> {
354        self.serialize_seq(Some(len))
355    }
356
357    fn serialize_map(self, _len: Option<usize>) -> Result<OwnedMapOrSlice, SerError> {
358        Ok(OwnedMapOrSlice::map())
359    }
360
361    fn serialize_struct(self, _name: &'static str, len: usize) -> Result<OwnedMapOrSlice, SerError> {
362        self.serialize_map(Some(len))
363    }
364
365    fn serialize_tuple_variant(
366        self,
367        _name: &'static str,
368        _variant_index: u32,
369        _variant: &'static str,
370        _len: usize,
371    ) -> Result<Self::SerializeTupleVariant, SerError> {
372        Err(SerError("tuple variants not supported".into()))
373    }
374
375    fn serialize_struct_variant(
376        self,
377        _name: &'static str,
378        _variant_index: u32,
379        _variant: &'static str,
380        _len: usize,
381    ) -> Result<Self::SerializeStructVariant, SerError> {
382        Err(SerError("struct variants not supported".into()))
383    }
384}
385
386pub struct OwnedMapOrSlice {
387    pending_key: Option<String>,
388    inner: OwnedMapOrSliceInner,
389}
390
391enum OwnedMapOrSliceInner {
392    Map(BTreeMap<String, Value>),
393    Slice(Vec<Value>),
394}
395
396impl OwnedMapOrSlice {
397    fn map() -> Self {
398        OwnedMapOrSlice {
399            pending_key: None,
400            inner: OwnedMapOrSliceInner::Map(BTreeMap::new()),
401        }
402    }
403    fn slice() -> Self {
404        OwnedMapOrSlice {
405            pending_key: None,
406            inner: OwnedMapOrSliceInner::Slice(Vec::new()),
407        }
408    }
409}
410
411impl SerializeSeq for OwnedMapOrSlice {
412    type Ok = Value;
413    type Error = SerError;
414
415    fn serialize_element<T: ?Sized + Serialize>(&mut self, value: &T) -> Result<(), SerError> {
416        match self.inner {
417            OwnedMapOrSliceInner::Slice(ref mut v) => {
418                v.push(value.serialize(&ValueSerializer)?);
419                Ok(())
420            }
421            OwnedMapOrSliceInner::Map(_) => Err(SerError("expected map entry, got element".into())),
422        }
423    }
424
425    fn end(self) -> Result<Value, SerError> {
426        match self.inner {
427            OwnedMapOrSliceInner::Slice(v) => Ok(Value::Slice(v)),
428            OwnedMapOrSliceInner::Map(m) => Ok(Value::Map(m)),
429        }
430    }
431}
432
433impl ser::SerializeTuple for OwnedMapOrSlice {
434    type Ok = Value;
435    type Error = SerError;
436
437    fn serialize_element<T: ?Sized + Serialize>(&mut self, value: &T) -> Result<(), SerError> {
438        SerializeSeq::serialize_element(self, value)
439    }
440
441    fn end(self) -> Result<Value, SerError> {
442        SerializeSeq::end(self)
443    }
444}
445
446impl SerializeTupleStruct for OwnedMapOrSlice {
447    type Ok = Value;
448    type Error = SerError;
449
450    fn serialize_field<T: ?Sized + Serialize>(&mut self, value: &T) -> Result<(), SerError> {
451        SerializeSeq::serialize_element(self, value)
452    }
453
454    fn end(self) -> Result<Value, SerError> {
455        SerializeSeq::end(self)
456    }
457}
458
459impl SerializeMap for OwnedMapOrSlice {
460    type Ok = Value;
461    type Error = SerError;
462
463    fn serialize_key<T: ?Sized + Serialize>(&mut self, key: &T) -> Result<(), SerError> {
464        match self.inner {
465            OwnedMapOrSliceInner::Slice(_) => Err(SerError("expected element, got map key".into())),
466            OwnedMapOrSliceInner::Map(_) => {
467                let k = key.serialize(&ValueSerializer)?;
468                match k {
469                    Value::String(s) => {
470                        self.pending_key = Some(s);
471                        Ok(())
472                    }
473                    _ => Err(SerError("map key must be a string".into())),
474                }
475            }
476        }
477    }
478
479    fn serialize_value<T: ?Sized + Serialize>(&mut self, value: &T) -> Result<(), SerError> {
480        match self.inner {
481            OwnedMapOrSliceInner::Slice(_) => Err(SerError("expected element, got map value".into())),
482            OwnedMapOrSliceInner::Map(ref mut m) => {
483                let key = self
484                    .pending_key
485                    .take()
486                    .ok_or_else(|| SerError("missing map key".into()))?;
487                let v = value.serialize(&ValueSerializer)?;
488                m.insert(key, v);
489                Ok(())
490            }
491        }
492    }
493
494    fn end(self) -> Result<Value, SerError> {
495        match self.inner {
496            OwnedMapOrSliceInner::Map(m) => Ok(Value::Map(m)),
497            OwnedMapOrSliceInner::Slice(_) => Err(SerError("expected map, got slice".into())),
498        }
499    }
500}
501
502impl SerializeStruct for OwnedMapOrSlice {
503    type Ok = Value;
504    type Error = SerError;
505
506    fn serialize_field<T: ?Sized + Serialize>(&mut self, key: &'static str, value: &T) -> Result<(), SerError> {
507        match self.inner {
508            OwnedMapOrSliceInner::Map(ref mut m) => {
509                let v = value.serialize(&ValueSerializer)?;
510                m.insert(key.to_string(), v);
511                Ok(())
512            }
513            OwnedMapOrSliceInner::Slice(_) => Err(SerError("expected map, got struct".into())),
514        }
515    }
516
517    fn end(self) -> Result<Value, SerError> {
518        match self.inner {
519            OwnedMapOrSliceInner::Map(m) => Ok(Value::Map(m)),
520            OwnedMapOrSliceInner::Slice(_) => Err(SerError("expected map, got slice".into())),
521        }
522    }
523}
524
525// ---------------------------------------------------------------------------
526// ByteSerializer — writes MaxMind DB binary directly into a Vec<u8>,
527//                  avoiding the intermediate Value allocation.
528// ---------------------------------------------------------------------------
529
530/// Serializes a value directly into MaxMind DB binary form, avoiding an
531/// intermediate [`Value`] allocation.
532pub struct ByteSerializer {
533    buf: Vec<u8>,
534}
535
536impl ByteSerializer {
537    pub fn new() -> Self {
538        ByteSerializer {
539            buf: Vec::new(),
540        }
541    }
542
543    pub fn into_bytes(self) -> Vec<u8> {
544        self.buf
545    }
546}
547
548impl Default for ByteSerializer {
549    fn default() -> Self {
550        Self::new()
551    }
552}
553
554/// Largest payload size representable in a MaxMind DB header. The three-byte
555/// size extension covers values below this bound.
556const MAX_PAYLOAD_SIZE: usize = 65_821 + (1 << 24);
557
558/// Reject payloads whose length cannot be encoded.
559fn check_payload_size(len: usize) -> Result<(), SerError> {
560    if len >= MAX_PAYLOAD_SIZE {
561        Err(SerError(format!(
562            "payload of {len} bytes is too large to encode (maximum is {})",
563            MAX_PAYLOAD_SIZE - 1
564        )))
565    } else {
566        Ok(())
567    }
568}
569
570fn ctrl_and_size(w: &mut Vec<u8>, type_num: u8, payload_len: usize) {
571    debug_assert!(payload_len < MAX_PAYLOAD_SIZE, "payload size {payload_len} is not encodable");
572    let first_byte: u8;
573    let second_byte: u8;
574
575    if type_num < 8 {
576        first_byte = type_num << 5;
577        second_byte = 0;
578    } else {
579        first_byte = 0;
580        second_byte = type_num - 7;
581    }
582
583    let size_val = match payload_len {
584        s if s < FIRST_SIZE => s as u8,
585        s if s < SECOND_SIZE => 29u8,
586        s if s < THIRD_SIZE => 30u8,
587        _ => 31u8,
588    };
589    w.push(first_byte | size_val);
590
591    if second_byte != 0 {
592        w.push(second_byte);
593    }
594
595    match payload_len {
596        s if s < FIRST_SIZE => {}
597        s if s < SECOND_SIZE => {
598            w.push((s - FIRST_SIZE) as u8);
599        }
600        s if s < THIRD_SIZE => {
601            let v = s - SECOND_SIZE;
602            w.push((v >> 8) as u8);
603            w.push((v & 0xFF) as u8);
604        }
605        s => {
606            let v = s - THIRD_SIZE;
607            w.push((v >> 16) as u8);
608            w.push((v >> 8) as u8);
609            w.push((v & 0xFF) as u8);
610        }
611    }
612}
613
614struct ByteMapOrSlice<'a> {
615    ser: &'a mut ByteSerializer,
616    state: ByteMOS,
617    pending_key: Option<String>,
618}
619
620enum ByteMOS {
621    Map(BTreeMap<String, Vec<u8>>),
622    Slice(Vec<Vec<u8>>),
623}
624
625impl ByteMapOrSlice<'_> {
626    fn flush_map(self) -> Result<(), SerError> {
627        match self.state {
628            ByteMOS::Map(m) => {
629                check_payload_size(m.len())?;
630                ctrl_and_size(&mut self.ser.buf, 7, m.len());
631                for (k, v) in m {
632                    write_str_raw(&mut self.ser.buf, &k)?;
633                    self.ser.buf.extend_from_slice(&v);
634                }
635                Ok(())
636            }
637            ByteMOS::Slice(s) => {
638                check_payload_size(s.len())?;
639                ctrl_and_size(&mut self.ser.buf, 11, s.len());
640                for v in s {
641                    self.ser.buf.extend_from_slice(&v);
642                }
643                Ok(())
644            }
645        }
646    }
647}
648
649fn write_str_raw(w: &mut Vec<u8>, s: &str) -> Result<(), SerError> {
650    check_payload_size(s.len())?;
651    ctrl_and_size(w, 2, s.len());
652    w.extend_from_slice(s.as_bytes());
653    Ok(())
654}
655
656fn write_bytes_raw(w: &mut Vec<u8>, b: &[u8]) -> Result<(), SerError> {
657    check_payload_size(b.len())?;
658    ctrl_and_size(w, 4, b.len());
659    w.extend_from_slice(b);
660    Ok(())
661}
662
663impl SerializeSeq for ByteMapOrSlice<'_> {
664    type Ok = ();
665    type Error = SerError;
666
667    fn serialize_element<T: ?Sized + Serialize>(&mut self, value: &T) -> Result<(), SerError> {
668        let mut sub = ByteSerializer::new();
669        value.serialize(ByteSerializerRef(&mut sub))?;
670        match self.state {
671            ByteMOS::Slice(ref mut v) => v.push(sub.buf),
672            ByteMOS::Map(_) => return Err(SerError("expected map entry, got element".into())),
673        }
674        Ok(())
675    }
676
677    fn end(self) -> Result<(), SerError> {
678        self.flush_map()
679    }
680}
681
682impl ser::SerializeTuple for ByteMapOrSlice<'_> {
683    type Ok = ();
684    type Error = SerError;
685
686    fn serialize_element<T: ?Sized + Serialize>(&mut self, value: &T) -> Result<(), SerError> {
687        SerializeSeq::serialize_element(self, value)
688    }
689
690    fn end(self) -> Result<(), SerError> {
691        SerializeSeq::end(self)
692    }
693}
694
695impl SerializeTupleStruct for ByteMapOrSlice<'_> {
696    type Ok = ();
697    type Error = SerError;
698
699    fn serialize_field<T: ?Sized + Serialize>(&mut self, value: &T) -> Result<(), SerError> {
700        SerializeSeq::serialize_element(self, value)
701    }
702
703    fn end(self) -> Result<(), SerError> {
704        SerializeSeq::end(self)
705    }
706}
707
708impl SerializeMap for ByteMapOrSlice<'_> {
709    type Ok = ();
710    type Error = SerError;
711
712    fn serialize_key<T: ?Sized + Serialize>(&mut self, key: &T) -> Result<(), SerError> {
713        match self.state {
714            ByteMOS::Slice(_) => Err(SerError("expected element, got map key".into())),
715            ByteMOS::Map(_) => {
716                // Keys are stored as decoded strings and re-encoded when the map
717                // is flushed, so only the value form is needed here.
718                match key.serialize(&ValueSerializer)? {
719                    Value::String(s_val) => {
720                        self.pending_key = Some(s_val);
721                        Ok(())
722                    }
723                    _ => Err(SerError("map key must be a string".into())),
724                }
725            }
726        }
727    }
728
729    fn serialize_value<T: ?Sized + Serialize>(&mut self, value: &T) -> Result<(), SerError> {
730        match self.state {
731            ByteMOS::Slice(_) => Err(SerError("expected element, got map value".into())),
732            ByteMOS::Map(ref mut m) => {
733                let key = self
734                    .pending_key
735                    .take()
736                    .ok_or_else(|| SerError("missing map key".into()))?;
737                let mut sub = ByteSerializer::new();
738                value.serialize(ByteSerializerRef(&mut sub))?;
739                m.insert(key, sub.buf);
740                Ok(())
741            }
742        }
743    }
744
745    fn end(self) -> Result<(), SerError> {
746        self.flush_map()
747    }
748}
749
750impl SerializeStruct for ByteMapOrSlice<'_> {
751    type Ok = ();
752    type Error = SerError;
753
754    fn serialize_field<T: ?Sized + Serialize>(&mut self, key: &'static str, value: &T) -> Result<(), SerError> {
755        match self.state {
756            ByteMOS::Map(ref mut m) => {
757                let mut sub = ByteSerializer::new();
758                value.serialize(ByteSerializerRef(&mut sub))?;
759                m.insert(key.to_string(), sub.buf);
760                Ok(())
761            }
762            ByteMOS::Slice(_) => Err(SerError("expected map, got struct".into())),
763        }
764    }
765
766    fn end(self) -> Result<(), SerError> {
767        self.flush_map()
768    }
769}
770
771struct ByteSerializerRef<'a>(&'a mut ByteSerializer);
772
773impl<'a> ser::Serializer for ByteSerializerRef<'a> {
774    type Ok = ();
775    type Error = SerError;
776
777    type SerializeSeq = ByteMapOrSlice<'a>;
778    type SerializeTuple = ByteMapOrSlice<'a>;
779    type SerializeTupleStruct = ByteMapOrSlice<'a>;
780    type SerializeTupleVariant = ser::Impossible<(), SerError>;
781    type SerializeMap = ByteMapOrSlice<'a>;
782    type SerializeStruct = ByteMapOrSlice<'a>;
783    type SerializeStructVariant = ser::Impossible<(), SerError>;
784
785    fn serialize_bool(self, v: bool) -> Result<(), SerError> {
786        let payload_len = if v { 1 } else { 0 };
787        ctrl_and_size(&mut self.0.buf, 14, payload_len);
788        Ok(())
789    }
790
791    fn serialize_i8(self, v: i8) -> Result<(), SerError> {
792        self.serialize_i32(v as i32)
793    }
794
795    fn serialize_i16(self, v: i16) -> Result<(), SerError> {
796        self.serialize_i32(v as i32)
797    }
798
799    fn serialize_i32(self, v: i32) -> Result<(), SerError> {
800        let u = v as u32;
801        let size = 4 - (u.leading_zeros() as usize / 8);
802        ctrl_and_size(&mut self.0.buf, 8, size);
803        for i in (0..size).rev() {
804            self.0.buf.push((u >> (8 * i)) as u8);
805        }
806        Ok(())
807    }
808
809    fn serialize_i64(self, v: i64) -> Result<(), SerError> {
810        if let Ok(v) = i32::try_from(v) {
811            self.serialize_i32(v)
812        } else if v >= 0 {
813            self.serialize_u64(v as u64)
814        } else {
815            Err(SerError(format!(
816                "negative integer {v} is out of range for a MaxMind DB value (minimum is {})",
817                i32::MIN
818            )))
819        }
820    }
821
822    fn serialize_u8(self, v: u8) -> Result<(), SerError> {
823        self.serialize_u16(v as u16)
824    }
825
826    fn serialize_u16(self, v: u16) -> Result<(), SerError> {
827        let size = 2 - (v.leading_zeros() as usize / 8);
828        ctrl_and_size(&mut self.0.buf, 5, size);
829        for i in (0..size).rev() {
830            self.0.buf.push((v >> (8 * i)) as u8);
831        }
832        Ok(())
833    }
834
835    fn serialize_u32(self, v: u32) -> Result<(), SerError> {
836        let size = 4 - (v.leading_zeros() as usize / 8);
837        ctrl_and_size(&mut self.0.buf, 6, size);
838        for i in (0..size).rev() {
839            self.0.buf.push((v >> (8 * i)) as u8);
840        }
841        Ok(())
842    }
843
844    fn serialize_u64(self, v: u64) -> Result<(), SerError> {
845        let size = 8 - (v.leading_zeros() as usize / 8);
846        ctrl_and_size(&mut self.0.buf, 9, size);
847        for i in (0..size).rev() {
848            self.0.buf.push((v >> (8 * i)) as u8);
849        }
850        Ok(())
851    }
852
853    fn serialize_u128(self, v: u128) -> Result<(), SerError> {
854        let size = 16 - (v.leading_zeros() as usize / 8);
855        ctrl_and_size(&mut self.0.buf, 10, size);
856        for i in (0..size).rev() {
857            self.0.buf.push((v >> (8 * i)) as u8);
858        }
859        Ok(())
860    }
861
862    fn serialize_f32(self, v: f32) -> Result<(), SerError> {
863        ctrl_and_size(&mut self.0.buf, 15, 4);
864        self.0.buf.extend_from_slice(&v.to_bits().to_be_bytes());
865        Ok(())
866    }
867
868    fn serialize_f64(self, v: f64) -> Result<(), SerError> {
869        ctrl_and_size(&mut self.0.buf, 3, 8);
870        self.0.buf.extend_from_slice(&v.to_bits().to_be_bytes());
871        Ok(())
872    }
873
874    fn serialize_char(self, v: char) -> Result<(), SerError> {
875        self.serialize_str(&v.to_string())
876    }
877
878    fn serialize_str(self, v: &str) -> Result<(), SerError> {
879        write_str_raw(&mut self.0.buf, v)?;
880        Ok(())
881    }
882
883    fn serialize_bytes(self, v: &[u8]) -> Result<(), SerError> {
884        write_bytes_raw(&mut self.0.buf, v)?;
885        Ok(())
886    }
887
888    fn serialize_none(self) -> Result<(), SerError> {
889        Err(SerError("MaxMind DB has no null type".into()))
890    }
891
892    fn serialize_some<T: ?Sized + Serialize>(self, value: &T) -> Result<(), SerError> {
893        value.serialize(self)
894    }
895
896    fn serialize_unit(self) -> Result<(), SerError> {
897        Err(SerError("MaxMind DB has no null type".into()))
898    }
899
900    fn serialize_unit_struct(self, _name: &'static str) -> Result<(), SerError> {
901        Err(SerError("MaxMind DB has no null type".into()))
902    }
903
904    fn serialize_unit_variant(
905        self,
906        _name: &'static str,
907        _variant_index: u32,
908        variant: &'static str,
909    ) -> Result<(), SerError> {
910        self.serialize_str(variant)
911    }
912
913    fn serialize_newtype_struct<T: ?Sized + Serialize>(self, _name: &'static str, value: &T) -> Result<(), SerError> {
914        value.serialize(self)
915    }
916
917    fn serialize_newtype_variant<T: ?Sized + Serialize>(
918        self,
919        _name: &'static str,
920        _variant_index: u32,
921        _variant: &'static str,
922        value: &T,
923    ) -> Result<(), SerError> {
924        value.serialize(self)
925    }
926
927    fn serialize_seq(self, _len: Option<usize>) -> Result<ByteMapOrSlice<'a>, SerError> {
928        Ok(ByteMapOrSlice {
929            ser: self.0,
930            state: ByteMOS::Slice(Vec::new()),
931            pending_key: None,
932        })
933    }
934
935    fn serialize_tuple(self, len: usize) -> Result<ByteMapOrSlice<'a>, SerError> {
936        self.serialize_seq(Some(len))
937    }
938
939    fn serialize_tuple_struct(self, _name: &'static str, len: usize) -> Result<ByteMapOrSlice<'a>, SerError> {
940        self.serialize_seq(Some(len))
941    }
942
943    fn serialize_map(self, _len: Option<usize>) -> Result<ByteMapOrSlice<'a>, SerError> {
944        Ok(ByteMapOrSlice {
945            ser: self.0,
946            state: ByteMOS::Map(BTreeMap::new()),
947            pending_key: None,
948        })
949    }
950
951    fn serialize_struct(self, _name: &'static str, len: usize) -> Result<ByteMapOrSlice<'a>, SerError> {
952        self.serialize_map(Some(len))
953    }
954
955    fn serialize_tuple_variant(
956        self,
957        _name: &'static str,
958        _variant_index: u32,
959        _variant: &'static str,
960        _len: usize,
961    ) -> Result<Self::SerializeTupleVariant, SerError> {
962        Err(SerError("tuple variants not supported".into()))
963    }
964
965    fn serialize_struct_variant(
966        self,
967        _name: &'static str,
968        _variant_index: u32,
969        _variant: &'static str,
970        _len: usize,
971    ) -> Result<Self::SerializeStructVariant, SerError> {
972        Err(SerError("struct variants not supported".into()))
973    }
974}
975
976/// Serialize a `Serialize` value directly into MaxMind DB binary format.
977/// Returns the encoded bytes. This avoids allocating an intermediate `Value`.
978pub fn encode_serialize(v: &impl Serialize) -> Result<Vec<u8>, SerError> {
979    let mut ser = ByteSerializer::new();
980    v.serialize(ByteSerializerRef(&mut ser))?;
981    Ok(ser.buf)
982}
983
984fn encode_pointer_raw(w: &mut Vec<u8>, pointer: u32) {
985    match pointer {
986        p if p < 2048 => {
987            w.push(0b00100000 | ((p >> 8) as u8 & 0x07));
988            w.push((p & 0xFF) as u8);
989        }
990        p if p < 526336 => {
991            let v = p - 2048;
992            w.push(0b00101000 | ((v >> 16) as u8 & 0x07));
993            w.push((v >> 8) as u8);
994            w.push((v & 0xFF) as u8);
995        }
996        p if p < 134744064 => {
997            let v = p - 526336;
998            w.push(0b00110000 | ((v >> 24) as u8 & 0x07));
999            w.push((v >> 16) as u8);
1000            w.push((v >> 8) as u8);
1001            w.push((v & 0xFF) as u8);
1002        }
1003        p => {
1004            w.push(0b00111000);
1005            w.push((p >> 24) as u8);
1006            w.push((p >> 16) as u8);
1007            w.push((p >> 8) as u8);
1008            w.push((p & 0xFF) as u8);
1009        }
1010    }
1011}
1012
1013impl ByteSerializerRef<'_> {
1014    fn serialize_value(self, v: &Value) -> Result<(), SerError> {
1015        match v {
1016            Value::Pointer(p) => {
1017                encode_pointer_raw(&mut self.0.buf, *p);
1018                Ok(())
1019            }
1020            Value::String(s) => {
1021                write_str_raw(&mut self.0.buf, s)?;
1022                Ok(())
1023            }
1024            Value::Float64(f) => {
1025                ctrl_and_size(&mut self.0.buf, 3, 8);
1026                self.0.buf.extend_from_slice(&f.to_bits().to_be_bytes());
1027                Ok(())
1028            }
1029            Value::Bytes(b) => {
1030                write_bytes_raw(&mut self.0.buf, b)?;
1031                Ok(())
1032            }
1033            Value::Uint16(n) => {
1034                let size = 2 - (n.leading_zeros() as usize / 8);
1035                ctrl_and_size(&mut self.0.buf, 5, size);
1036                for i in (0..size).rev() {
1037                    self.0.buf.push((n >> (8 * i)) as u8);
1038                }
1039                Ok(())
1040            }
1041            Value::Uint32(n) => {
1042                let size = 4 - (n.leading_zeros() as usize / 8);
1043                ctrl_and_size(&mut self.0.buf, 6, size);
1044                for i in (0..size).rev() {
1045                    self.0.buf.push((n >> (8 * i)) as u8);
1046                }
1047                Ok(())
1048            }
1049            Value::Int32(n) => {
1050                let u = *n as u32;
1051                let size = 4 - (u.leading_zeros() as usize / 8);
1052                ctrl_and_size(&mut self.0.buf, 8, size);
1053                for i in (0..size).rev() {
1054                    self.0.buf.push((u >> (8 * i)) as u8);
1055                }
1056                Ok(())
1057            }
1058            Value::Uint64(n) => {
1059                let size = 8 - (n.leading_zeros() as usize / 8);
1060                ctrl_and_size(&mut self.0.buf, 9, size);
1061                for i in (0..size).rev() {
1062                    self.0.buf.push((n >> (8 * i)) as u8);
1063                }
1064                Ok(())
1065            }
1066            Value::Uint128(n) => {
1067                let size = 16 - (n.leading_zeros() as usize / 8);
1068                ctrl_and_size(&mut self.0.buf, 10, size);
1069                for i in (0..size).rev() {
1070                    self.0.buf.push((n >> (8 * i)) as u8);
1071                }
1072                Ok(())
1073            }
1074            Value::Map(m) => {
1075                check_payload_size(m.len())?;
1076                ctrl_and_size(&mut self.0.buf, 7, m.len());
1077                for (k, val) in m.iter() {
1078                    write_str_raw(&mut self.0.buf, k)?;
1079                    ByteSerializerRef(&mut *self.0).serialize_value(val)?;
1080                }
1081                Ok(())
1082            }
1083            Value::Slice(s) => {
1084                check_payload_size(s.len())?;
1085                ctrl_and_size(&mut self.0.buf, 11, s.len());
1086                for val in s {
1087                    ByteSerializerRef(&mut *self.0).serialize_value(val)?;
1088                }
1089                Ok(())
1090            }
1091            Value::Bool(b) => {
1092                let payload_len = if *b { 1 } else { 0 };
1093                ctrl_and_size(&mut self.0.buf, 14, payload_len);
1094                Ok(())
1095            }
1096            Value::Float32(f) => {
1097                ctrl_and_size(&mut self.0.buf, 15, 4);
1098                self.0.buf.extend_from_slice(&f.to_bits().to_be_bytes());
1099                Ok(())
1100            }
1101        }
1102    }
1103}
1104
1105/// Encode `v` into MaxMind DB binary form.
1106///
1107/// # Errors
1108///
1109/// Returns an error if a payload is larger than the format can encode
1110/// (`16_843_036` bytes).
1111pub fn encode_value(v: &Value) -> EncodeResult<Vec<u8>> {
1112    let mut ser = ByteSerializer::new();
1113    ByteSerializerRef(&mut ser).serialize_value(v).map_err(|e| e.0)?;
1114    Ok(ser.into_bytes())
1115}
1116
1117/// Encode `v` and return the number of bytes it would occupy.
1118///
1119/// Returns `0` if `v` cannot be encoded; use [`encode_value`] if you need the
1120/// error.
1121pub fn encoded_size(v: &Value) -> usize {
1122    encode_value(v).map(|b| b.len()).unwrap_or(0)
1123}
maxminddb/encoder.rs

maxminddb/
encoder.rs
