mirror of
https://github.com/nushell/nushell
synced 2025-01-15 22:54:16 +00:00
Ensure source buffer is cleared after reading in MaybeTextCodec. (#2168)
This commit is contained in:
parent
97cedeb324
commit
8551e06d9e
1 changed files with 49 additions and 25 deletions
|
@ -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());
|
||||
}
|
||||
}
|
||||
|
|
Loading…
Reference in a new issue