1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125

/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use crate::error::Result;
use serde::Serialize;
use std::io::{self, Write};

/// A writer that counts the number of bytes it's asked to write, and discards
/// the data. Used to calculate the serialized size of the commands list.
#[derive(Clone, Copy, Default)]
pub struct WriteCount(usize);

impl WriteCount {
    #[inline]
    pub fn len(self) -> usize {
        self.0
    }
}

impl Write for WriteCount {
    #[inline]
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.0 += buf.len();
        Ok(buf.len())
    }

    #[inline]
    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

/// Returns the size of the given value, in bytes, when serialized to JSON.
fn compute_serialized_size<T: Serialize>(value: &T) -> Result<usize> {
    let mut w = WriteCount::default();
    serde_json::to_writer(&mut w, value)?;
    Ok(w.len())
}

/// Truncates `list` to fit within `payload_size_max_bytes` when serialized to
/// JSON.
pub fn shrink_to_fit<T: Serialize>(list: &mut Vec<T>, payload_size_max_bytes: usize) -> Result<()> {
    let size = compute_serialized_size(&list)?;
    // See bug 535326 comment 8 for an explanation of the estimation
    match ((payload_size_max_bytes / 4) * 3).checked_sub(1500) {
        Some(max_serialized_size) => {
            if size > max_serialized_size {
                // Estimate a little more than the direct fraction to maximize packing
                let cutoff = (list.len() * max_serialized_size - 1) / size + 1;
                list.truncate(cutoff + 1);
                // Keep dropping off the last entry until the data fits.
                while compute_serialized_size(&list)? > max_serialized_size {
                    if list.pop().is_none() {
                        break;
                    }
                }
            }
            Ok(())
        }
        None => {
            list.clear();
            Ok(())
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::clients::record::CommandRecord;

    #[test]
    fn test_compute_serialized_size() {
        assert_eq!(compute_serialized_size(&1).unwrap(), 1);
        assert_eq!(compute_serialized_size(&"hi").unwrap(), 4);
        assert_eq!(
            compute_serialized_size(&["hi", "hello", "bye"]).unwrap(),
            20
        );
    }

    #[test]
    fn test_shrink_to_fit() {
        let mut commands = vec![
            CommandRecord {
                name: "wipeEngine".into(),
                args: vec!["bookmarks".into()],
                flow_id: Some("flow".into()),
            },
            CommandRecord {
                name: "resetEngine".into(),
                args: vec!["history".into()],
                flow_id: Some("flow".into()),
            },
            CommandRecord {
                name: "logout".into(),
                args: Vec::new(),
                flow_id: None,
            },
        ];

        // 4096 bytes is enough to fit all three commands.
        shrink_to_fit(&mut commands, 4096).unwrap();
        assert_eq!(commands.len(), 3);

        let sizes = commands
            .iter()
            .map(|c| compute_serialized_size(c).unwrap())
            .collect::<Vec<_>>();
        assert_eq!(sizes, &[61, 60, 30]);

        // `logout` won't fit within 2168 bytes.
        shrink_to_fit(&mut commands, 2168).unwrap();
        assert_eq!(commands.len(), 2);

        // `resetEngine` won't fit within 2084 bytes.
        shrink_to_fit(&mut commands, 2084).unwrap();
        assert_eq!(commands.len(), 1);

        // `wipeEngine` won't fit at all.
        shrink_to_fit(&mut commands, 1024).unwrap();
        assert!(commands.is_empty());
    }
}