mongodb/
bson_util.rs

1use std::{
2    collections::HashSet,
3    convert::TryFrom,
4    io::{Read, Write},
5};
6
7use serde::Serialize;
8
9use crate::{
10    bson::{
11        oid::ObjectId,
12        rawdoc,
13        Bson,
14        Document,
15        RawArrayBuf,
16        RawBson,
17        RawBsonRef,
18        RawDocumentBuf,
19    },
20    checked::Checked,
21    error::{Error, ErrorKind, Result},
22    runtime::SyncLittleEndianRead,
23};
24
25/// Coerce numeric types into an `i64` if it would be lossless to do so. If this Bson is not numeric
26/// or the conversion would be lossy (e.g. 1.5 -> 1), this returns `None`.
27#[allow(clippy::cast_possible_truncation)]
28pub(crate) fn get_int(val: &Bson) -> Option<i64> {
29    match *val {
30        Bson::Int32(i) => Some(i64::from(i)),
31        Bson::Int64(i) => Some(i),
32        Bson::Double(f) if (f - (f as i64 as f64)).abs() <= f64::EPSILON => Some(f as i64),
33        _ => None,
34    }
35}
36
37/// Coerce numeric types into an `i64` if it would be lossless to do so. If this Bson is not numeric
38/// or the conversion would be lossy (e.g. 1.5 -> 1), this returns `None`.
39pub(crate) fn get_int_raw(val: RawBsonRef<'_>) -> Option<i64> {
40    match val {
41        RawBsonRef::Int32(i) => get_int(&Bson::Int32(i)),
42        RawBsonRef::Int64(i) => get_int(&Bson::Int64(i)),
43        RawBsonRef::Double(i) => get_int(&Bson::Double(i)),
44        _ => None,
45    }
46}
47
48/// Coerce numeric types into an `u64` if it would be lossless to do so. If this Bson is not numeric
49/// or the conversion would be lossy (e.g. 1.5 -> 1), this returns `None`.
50#[allow(clippy::cast_possible_truncation)]
51pub(crate) fn get_u64(val: &Bson) -> Option<u64> {
52    match *val {
53        Bson::Int32(i) => u64::try_from(i).ok(),
54        Bson::Int64(i) => u64::try_from(i).ok(),
55        Bson::Double(f) if (f - (f as u64 as f64)).abs() <= f64::EPSILON => Some(f as u64),
56        _ => None,
57    }
58}
59
60pub(crate) fn to_bson_array(docs: &[Document]) -> Bson {
61    Bson::Array(docs.iter().map(|doc| Bson::Document(doc.clone())).collect())
62}
63
64pub(crate) fn to_raw_bson_array(docs: &[Document]) -> Result<RawBson> {
65    let mut array = RawArrayBuf::new();
66    for doc in docs {
67        array.push(RawDocumentBuf::from_document(doc)?);
68    }
69    Ok(RawBson::Array(array))
70}
71pub(crate) fn to_raw_bson_array_ser<T: Serialize>(values: &[T]) -> Result<RawBson> {
72    let mut array = RawArrayBuf::new();
73    for value in values {
74        array.push(bson::to_raw_document_buf(value)?);
75    }
76    Ok(RawBson::Array(array))
77}
78
79pub(crate) fn first_key(document: &Document) -> Option<&str> {
80    document.keys().next().map(String::as_str)
81}
82
83pub(crate) fn update_document_check(update: &Document) -> Result<()> {
84    match first_key(update) {
85        Some(key) => {
86            if !key.starts_with('$') {
87                Err(ErrorKind::InvalidArgument {
88                    message: "update document must only contain update modifiers".to_string(),
89                }
90                .into())
91            } else {
92                Ok(())
93            }
94        }
95        None => Err(ErrorKind::InvalidArgument {
96            message: "update document must not be empty".to_string(),
97        }
98        .into()),
99    }
100}
101
102pub(crate) fn replacement_document_check(replacement: &Document) -> Result<()> {
103    if let Some(key) = first_key(replacement) {
104        if key.starts_with('$') {
105            return Err(ErrorKind::InvalidArgument {
106                message: "replacement document must not contain update modifiers".to_string(),
107            }
108            .into());
109        }
110    }
111    Ok(())
112}
113
114pub(crate) fn replacement_raw_document_check(replacement: &RawDocumentBuf) -> Result<()> {
115    if let Some((key, _)) = replacement.iter().next().transpose()? {
116        if key.starts_with('$') {
117            return Err(ErrorKind::InvalidArgument {
118                message: "replacement document must not contain update modifiers".to_string(),
119            }
120            .into());
121        };
122    }
123    Ok(())
124}
125
126/// The size in bytes of the provided document's entry in a BSON array at the given index.
127pub(crate) fn array_entry_size_bytes(index: usize, doc_len: usize) -> Result<usize> {
128    //   * type (1 byte)
129    //   * number of decimal digits in key
130    //   * null terminator for the key (1 byte)
131    //   * size of value
132
133    (Checked::new(1) + num_decimal_digits(index) + 1 + doc_len).get()
134}
135
136pub(crate) fn vec_to_raw_array_buf(docs: Vec<RawDocumentBuf>) -> RawArrayBuf {
137    let mut array = RawArrayBuf::new();
138    for doc in docs {
139        array.push(doc);
140    }
141    array
142}
143
144/// The number of digits in `n` in base 10.
145/// Useful for calculating the size of an array entry in BSON.
146fn num_decimal_digits(mut n: usize) -> usize {
147    let mut digits = 0;
148
149    loop {
150        n /= 10;
151        digits += 1;
152
153        if n == 0 {
154            return digits;
155        }
156    }
157}
158
159/// Read a document's raw BSON bytes from the provided reader.
160pub(crate) fn read_document_bytes<R: Read>(mut reader: R) -> Result<Vec<u8>> {
161    let length = Checked::new(reader.read_i32_sync()?);
162
163    let mut bytes = Vec::with_capacity(length.try_into()?);
164    bytes.write_all(&length.try_into::<u32>()?.to_le_bytes())?;
165
166    reader
167        .take((length - 4).try_into()?)
168        .read_to_end(&mut bytes)?;
169
170    Ok(bytes)
171}
172
173pub(crate) fn extend_raw_document_buf(
174    this: &mut RawDocumentBuf,
175    other: RawDocumentBuf,
176) -> Result<()> {
177    let mut keys: HashSet<String> = HashSet::new();
178    for elem in this.iter_elements() {
179        keys.insert(elem?.key().to_owned());
180    }
181    for result in other.iter() {
182        let (k, v) = result?;
183        if keys.contains(k) {
184            return Err(Error::internal(format!(
185                "duplicate raw document key {:?}",
186                k
187            )));
188        }
189        this.append(k, v.to_raw_bson());
190    }
191    Ok(())
192}
193
194pub(crate) fn append_ser(
195    this: &mut RawDocumentBuf,
196    key: impl AsRef<str>,
197    value: impl Serialize,
198) -> Result<()> {
199    #[derive(Serialize)]
200    struct Helper<T> {
201        value: T,
202    }
203    let raw_doc = bson::to_raw_document_buf(&Helper { value })?;
204    this.append_ref(
205        key,
206        raw_doc
207            .get("value")?
208            .ok_or_else(|| Error::internal("no value"))?,
209    );
210    Ok(())
211}
212
213/// Returns the _id field of this document, prepending the field to the document if one is not
214/// already present.
215pub(crate) fn get_or_prepend_id_field(doc: &mut RawDocumentBuf) -> Result<Bson> {
216    match doc.get("_id")? {
217        Some(id) => Ok(id.try_into()?),
218        None => {
219            let id = ObjectId::new();
220            let mut new_bytes = rawdoc! { "_id": id }.into_bytes();
221
222            // Remove the trailing null byte (which will be replaced by the null byte in the given
223            // document) and append the document's elements
224            new_bytes.pop();
225            new_bytes.extend(&doc.as_bytes()[4..]);
226
227            let new_length: i32 = Checked::new(new_bytes.len()).try_into()?;
228            new_bytes[0..4].copy_from_slice(&new_length.to_le_bytes());
229
230            *doc = RawDocumentBuf::from_bytes(new_bytes)?;
231
232            Ok(id.into())
233        }
234    }
235}
236
237#[cfg(test)]
238mod test {
239    use crate::bson_util::num_decimal_digits;
240
241    #[test]
242    fn num_digits() {
243        assert_eq!(num_decimal_digits(0), 1);
244        assert_eq!(num_decimal_digits(1), 1);
245        assert_eq!(num_decimal_digits(10), 2);
246        assert_eq!(num_decimal_digits(15), 2);
247        assert_eq!(num_decimal_digits(100), 3);
248        assert_eq!(num_decimal_digits(125), 3);
249    }
250}
mongodb/bson_util.rs

mongodb/
bson_util.rs
