From 8551e06d9ea1d25d97989bc4874571c0bcc7d410 Mon Sep 17 00:00:00 2001
From: Jason Gedge <jason.gedge@shopify.com>
Date: Mon, 13 Jul 2020 19:24:52 -0400
Subject: [PATCH] Ensure source buffer is cleared after reading in
 MaybeTextCodec. (#2168)

---
 .../commands/classified/maybe_text_codec.rs   | 74 ++++++++++++-------
 1 file changed, 49 insertions(+), 25 deletions(-)

diff --git a/crates/nu-cli/src/commands/classified/maybe_text_codec.rs b/crates/nu-cli/src/commands/classified/maybe_text_codec.rs
index e7d11cfd2..8aba39b2e 100644
--- a/crates/nu-cli/src/commands/classified/maybe_text_codec.rs
+++ b/crates/nu-cli/src/commands/classified/maybe_text_codec.rs
@@ -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());
     }
 }