Ensure source buffer is cleared after reading in MaybeTextCodec. (#2168)

This commit is contained in:
Jason Gedge 2020-07-13 19:24:52 -04:00 committed by GitHub
parent 97cedeb324
commit 8551e06d9e
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23

View file

@ -5,8 +5,12 @@ use nu_errors::ShellError;
extern crate encoding_rs;
use encoding_rs::{CoderResult, Decoder, Encoding, UTF_8};
#[cfg(not(test))]
const OUTPUT_BUFFER_SIZE: usize = 8192;
#[cfg(test)]
const OUTPUT_BUFFER_SIZE: usize = 4;
#[derive(Debug, Eq, PartialEq)]
pub enum StringOrBinary {
String(String),
Binary(Vec<u8>),
@ -28,7 +32,6 @@ impl MaybeTextCodec {
}
impl Default for MaybeTextCodec {
// The default MaybeTextCodec uses a UTF_8 decoder
fn default() -> Self {
MaybeTextCodec {
decoder: UTF_8.new_decoder(),
@ -56,7 +59,6 @@ impl futures_codec::Encoder for MaybeTextCodec {
}
}
// TODO: Write some tests
impl futures_codec::Decoder for MaybeTextCodec {
type Item = StringOrBinary;
type Error = ShellError;
@ -68,36 +70,58 @@ impl futures_codec::Decoder for MaybeTextCodec {
let mut s = String::with_capacity(OUTPUT_BUFFER_SIZE);
let (res, read, replacements) = self.decoder.decode_to_string(src, &mut s, false);
// If we had to make replacements when converting to utf8, fallback to binary
if replacements {
return Ok(Some(StringOrBinary::Binary(src.to_vec())));
}
let (res, _read, replacements) = self.decoder.decode_to_string(src, &mut s, false);
match res {
CoderResult::InputEmpty => {
src.clear();
Ok(Some(StringOrBinary::String(s)))
}
CoderResult::OutputFull => {
// If the original buffer size is too small,
// We continue to allocate new Strings and append them to the result until the input buffer is smaller than the allocated String
let mut starting_index = read;
let result = if replacements {
// If we had to make replacements when converting to utf8, fall back to binary
StringOrBinary::Binary(src.to_vec())
} else {
// If original buffer size is too small, we continue to allocate new Strings and append
// them to the result until the input buffer is smaller than the allocated String
if let CoderResult::OutputFull = res {
let mut buffer = String::with_capacity(OUTPUT_BUFFER_SIZE);
loop {
let mut more = String::with_capacity(OUTPUT_BUFFER_SIZE);
let (res, read, _replacements) =
let (res, _read, _replacements) =
self.decoder
.decode_to_string(&src[starting_index..], &mut more, false);
s.push_str(&more);
// Our input buffer is smaller than out allocated String, we can stop now
.decode_to_string(&src[s.len()..], &mut buffer, false);
s.push_str(&buffer);
if let CoderResult::InputEmpty = res {
break;
}
starting_index += read;
buffer.clear();
}
src.clear();
Ok(Some(StringOrBinary::String(s)))
}
}
StringOrBinary::String(s)
};
src.clear();
Ok(Some(result))
}
}
#[cfg(test)]
mod tests {
use super::{MaybeTextCodec, StringOrBinary};
use bytes::BytesMut;
use futures_codec::Decoder;
// TODO: Write some more tests
#[test]
fn should_consume_all_bytes_from_source_when_temporary_buffer_overflows() {
let mut maybe_text = MaybeTextCodec::new(None);
let mut bytes = BytesMut::from("0123456789");
let text = maybe_text.decode(&mut bytes);
assert_eq!(
Ok(Some(StringOrBinary::String("0123456789".to_string()))),
text
);
assert!(bytes.is_empty());
}
}