protocol/
util.rs

1//! Helper functions for dealing with sets or lists of parcels.
2
3use crate::{hint, Parcel, Error, ErrorKind, TryFromIntError, Settings};
4use crate::types::Integer;
5
6use std::io::prelude::*;
7use std::io;
8
9/// The integer type that we will use to send length prefixes.
10pub type SizeType = u32;
11
12/// Reads a string of specified length from a stream.
13pub fn read_string(byte_count: usize,
14                   read: &mut dyn Read,
15                   settings: &Settings)
16    -> Result<String, Error> {
17    let bytes: Vec<u8> = read_items(byte_count, read, settings)?.collect();
18    String::from_utf8(bytes).map_err(Into::into)
19}
20
21/// Reads a specified number of items from a stream.
22pub fn read_items<T>(item_count: usize,
23                     read: &mut dyn Read,
24                     settings: &Settings)
25    -> Result<impl Iterator<Item=T>, Error>
26    where T: Parcel {
27    let mut elements = Vec::with_capacity(item_count);
28
29    for _ in 0..item_count {
30        let element = T::read(read, settings)?;
31        elements.push(element);
32    }
33    Ok(elements.into_iter())
34}
35
36/// Writes an iterator of parcels to the stream.
37///
38/// Does not include a length prefix.
39pub fn write_items<'a,T>(items: impl IntoIterator<Item=&'a T>,
40                         write: &mut dyn Write,
41                         settings: &Settings)
42    -> Result<(), Error>
43    where T: Parcel + 'a {
44    for item in items.into_iter() {
45        item.write(write, settings)?;
46    }
47    Ok(())
48}
49
50/// Reads a length-prefixed list from a stream.
51pub fn read_list<T>(read: &mut dyn Read,
52                    settings: &Settings,
53                    hints: &mut hint::Hints)
54    -> Result<Vec<T>, Error>
55    where T: Parcel {
56    self::read_list_ext::<SizeType, T>(read, settings, hints)
57}
58
59/// Writes a length-prefixed list to a stream.
60pub fn write_list<'a,T,I>(elements: I,
61                          write: &mut dyn Write,
62                          settings: &Settings,
63                          hints: &mut hint::Hints)
64    -> Result<(), Error>
65    where T: Parcel+'a,
66          I: IntoIterator<Item=&'a T> {
67    self::write_list_ext::<SizeType, T, I>(elements, write, settings, hints)
68}
69
70/// Reads a length-prefixed list from a stream.
71pub fn read_list_ext<S,T>(read: &mut dyn Read,
72                          settings: &Settings,
73                          hints: &mut hint::Hints)
74    -> Result<Vec<T>, Error>
75    where S: Integer,
76          T: Parcel {
77    match hints.current_field_length() {
78        Some(length) => {
79            match length.kind {
80                hint::LengthPrefixKind::Bytes => {
81                    let byte_count = length.length;
82
83                    // First, read all bytes of the list without processing them.
84                    let bytes: Vec<u8> = read_items(byte_count, read, settings)?.collect();
85                    let mut read_back_bytes = io::Cursor::new(bytes);
86
87                    // Then, parse the items until we reach the end of the buffer stream.
88                    let mut items = Vec::new();
89                    // FIXME: potential DoS vector, should timeout.
90                    while read_back_bytes.position() < byte_count as u64 {
91                        let item = match T::read(&mut read_back_bytes, settings).map_err(|e| e.0) {
92                            Ok(item) => item,
93                            Err(ErrorKind::Io(ref io)) if io.kind() == io::ErrorKind::UnexpectedEof => {
94                                // FIXME: make this a client error.
95                                panic!("length prefix in bytes does not match actual size");
96                            },
97                            Err(e) => return Err(e.into()),
98                        };
99                        items.push(item);
100                    }
101
102                    Ok(items)
103                },
104                hint::LengthPrefixKind::Elements => {
105                    read_items(length.length, read, settings).map(|i| i.collect())
106                },
107            }
108        },
109        None => {
110            // We do not know the length in the field in advance, therefore there
111            // the length prefix is not disjoint.
112            let size = S::read(read, settings)?;
113            let size: usize = size.to_usize().ok_or(TryFromIntError{ })?;
114
115            read_items(size, read, settings).map(|i| i.collect())
116        },
117    }
118}
119
120/// Writes a length-prefixed list to a stream.
121pub fn write_list_ext<'a,S,T,I>(elements: I,
122                                write: &mut dyn Write,
123                                settings: &Settings,
124                                hints: &mut hint::Hints)
125    -> Result<(), Error>
126    where S: Integer,
127          T: Parcel+'a,
128          I: IntoIterator<Item=&'a T> {
129    let elements: Vec<_> = elements.into_iter().collect();
130
131    match hints.current_field_length() {
132        // If there is an existing length prefix, don't bother sending another.
133        Some(_length) => {
134            ()
135        },
136        // The length is not known, send a prefix.
137        _ => {
138            let length = S::from_usize(elements.len()).ok_or(TryFromIntError{ })?;
139            length.write(write, settings)?;
140
141        },
142    }
143    write_items(elements.into_iter(), write, settings)?;
144
145    Ok(())
146}
147
protocol/util.rs

protocol/
util.rs
