ensure that all bytes coming from the terminal actually get written

it's not actually safe to just skip over bytes when we see a reset
escape sequence, because there might be other escape sequences changing
the terminal state that don't get cleared by a reset. for instance,
starting vim sends "switch to alternate screen" followed by "reset", and
if those two come in the same packet, we were previously dropping the
"switch to alternate screen" escape sequence, causing incorrect display.

4 files changed, 135 insertions, 94 deletions

diff --git a/src/cmd/record.rs b/src/cmd/record.rs
index 25c11f7..6702ba9 100644
--- a/src/cmd/record.rs
+++ b/src/cmd/record.rs
@@ -101,12 +101,7 @@ impl RecordSession {
     }
 
     fn record_bytes(&mut self, buf: &[u8]) {
-        let truncated = self.buffer.append(buf);
-        if truncated > self.sent_local {
-            self.sent_local = 0;
-        } else {
-            self.sent_local -= truncated;
-        }
+        self.sent_local -= self.buffer.append(buf, self.sent_local);
     }
 }
 
diff --git a/src/cmd/stream.rs b/src/cmd/stream.rs
index c1e1efa..c1d62c6 100644
--- a/src/cmd/stream.rs
+++ b/src/cmd/stream.rs
@@ -178,15 +178,14 @@ impl<S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Send + 'static>
     }
 
     fn record_bytes(&mut self, buf: &[u8]) {
-        let truncated = self.buffer.append(buf);
-        if truncated > self.sent_local {
-            self.sent_local = 0;
-        } else {
-            self.sent_local -= truncated;
-        }
-        if truncated > self.sent_remote {
-            self.sent_remote = 0;
+        let written = if self.connected {
+            self.sent_local.min(self.sent_remote)
         } else {
+            self.sent_local
+        };
+        let truncated = self.buffer.append(buf, written);
+        self.sent_local -= truncated;
+        if self.connected {
             self.sent_remote -= truncated;
         }
     }
diff --git a/src/server.rs b/src/server.rs
index f58cda4..6d38220 100644
--- a/src/server.rs
+++ b/src/server.rs
@@ -542,7 +542,7 @@ impl<S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Send + 'static>
     ) -> Result<()> {
         let saved_data = conn.state.saved_data_mut().unwrap();
 
-        saved_data.append(data);
+        saved_data.append(data, saved_data.len());
         for watch_conn in self.watchers_mut() {
             let watch_id = watch_conn.state.watch_id().unwrap();
             if conn.id == watch_id {
diff --git a/src/term.rs b/src/term.rs
index ba6efdd..fd82604 100644
--- a/src/term.rs
+++ b/src/term.rs
@@ -56,9 +56,47 @@ impl Buffer {
         self_
     }
 
-    pub fn append(&mut self, mut buf: &[u8]) -> usize {
-        let mut truncated = 0;
+    pub fn append(&mut self, buf: &[u8], written: usize) -> usize {
+        self.contents.extend_from_slice(buf);
+
+        if let Some(title) = self.find_window_title(buf) {
+            self.title = title;
+        }
+
+        if written == 0 {
+            return written;
+        }
+
+        if self.has_reset(buf) {
+            self.truncate_at(written);
+            return written;
+        }
+
+        if self.contents.len() > self.max_size {
+            let to_drop = self.contents.len() - self.max_size / 2;
+            let to_drop = to_drop.min(written);
+            if to_drop > 0 {
+                self.truncate_at(to_drop);
+                return to_drop;
+            }
+        }
 
+        0
+    }
+
+    pub fn contents(&self) -> &[u8] {
+        &self.contents
+    }
+
+    pub fn len(&self) -> usize {
+        self.contents.len()
+    }
+
+    pub fn title(&self) -> &str {
+        &self.title
+    }
+
+    fn find_window_title(&self, buf: &[u8]) -> Option<String> {
         let mut found = None;
         for window_title in WINDOW_TITLE {
             if let Some(i) = twoway::rfind_bytes(buf, window_title.0) {
@@ -78,58 +116,27 @@ impl Buffer {
                 }
             }
         }
-        if let Some((_, title)) = found {
-            self.title = title.to_string();
-        }
+        found.map(|(_, title)| title.to_string())
+    }
 
-        let mut found = None;
+    fn has_reset(&self, buf: &[u8]) -> bool {
         for reset in RESET {
-            if let Some(i) = twoway::rfind_bytes(buf, reset) {
-                let len = reset.len();
-                if let Some((i2, len2)) = found {
-                    if i + len > i2 + len2 || (i + len == i2 + len2 && i < i2)
-                    {
-                        found = Some((i, len));
-                    }
-                } else {
-                    found = Some((i, len));
-                }
+            if twoway::find_bytes(buf, reset).is_some() {
+                return true;
             }
         }
-        if let Some((i, _)) = found {
-            truncated = self.contents.len();
-            self.contents.clear();
-            buf = &buf[i..];
-        }
-
-        let prev_len = self.contents.len();
-        self.contents.extend_from_slice(buf);
-        if self.contents.len() > self.max_size {
-            let new_contents = self
-                .contents
-                .split_off(self.contents.len() - self.max_size / 2);
-            truncated = std::cmp::min(self.contents.len(), prev_len);
-            self.contents = new_contents;
-        }
-
-        truncated
-    }
-
-    pub fn contents(&self) -> &[u8] {
-        &self.contents
-    }
-
-    pub fn len(&self) -> usize {
-        self.contents.len()
+        false
     }
 
-    pub fn title(&self) -> &str {
-        &self.title
+    fn truncate_at(&mut self, i: usize) {
+        let new_contents = self.contents.split_off(i);
+        self.contents = new_contents;
     }
 }
 
 #[cfg(test)]
 #[allow(clippy::cognitive_complexity)]
+#[allow(clippy::redundant_clone)]
 #[allow(clippy::shadow_unrelated)]
 mod test {
     use super::*;
@@ -141,13 +148,13 @@ mod test {
         assert_eq!(buffer.len(), 0);
         assert_eq!(buffer.title(), "");
 
-        let n = buffer.append(b"foo");
+        let n = buffer.append(b"foo", 0);
         assert_eq!(buffer.contents(), b"foo");
         assert_eq!(buffer.len(), 3);
         assert_eq!(buffer.title(), "");
         assert_eq!(n, 0);
 
-        let n = buffer.append(b"bar");
+        let n = buffer.append(b"bar", 3);
         assert_eq!(buffer.contents(), b"foobar");
         assert_eq!(buffer.len(), 6);
         assert_eq!(buffer.title(), "");
@@ -158,66 +165,66 @@ mod test {
     fn test_clear() {
         let mut buffer = Buffer::new(100);
 
-        let n = buffer.append(b"foo");
+        let n = buffer.append(b"foo", 0);
         assert_eq!(buffer.len(), 3);
         assert_eq!(buffer.contents(), b"foo");
         assert_eq!(buffer.title(), "");
         assert_eq!(n, 0);
 
-        let n = buffer.append(b"\x1b[3J\x1b[H\x1b[2J");
+        let n = buffer.append(b"\x1b[3J\x1b[H\x1b[2J", 3);
         assert_eq!(buffer.len(), 11);
         assert_eq!(buffer.contents(), b"\x1b[3J\x1b[H\x1b[2J");
         assert_eq!(buffer.title(), "");
         assert_eq!(n, 3);
 
-        let n = buffer.append(b"bar");
+        let n = buffer.append(b"bar", 11);
         assert_eq!(buffer.len(), 14);
         assert_eq!(buffer.contents(), b"\x1b[3J\x1b[H\x1b[2Jbar");
         assert_eq!(buffer.title(), "");
         assert_eq!(n, 0);
 
-        let n = buffer.append(b"baz\x1bcquux");
-        assert_eq!(buffer.len(), 6);
-        assert_eq!(buffer.contents(), b"\x1bcquux");
+        let n = buffer.append(b"baz\x1bcquux", 14);
+        assert_eq!(buffer.len(), 9);
+        assert_eq!(buffer.contents(), b"baz\x1bcquux");
         assert_eq!(buffer.title(), "");
         assert_eq!(n, 14);
 
-        let n = buffer.append(b"blorg\x1b[H\x1b[J");
-        assert_eq!(buffer.len(), 6);
-        assert_eq!(buffer.contents(), b"\x1b[H\x1b[J");
+        let n = buffer.append(b"blorg\x1b[H\x1b[J", 9);
+        assert_eq!(buffer.len(), 11);
+        assert_eq!(buffer.contents(), b"blorg\x1b[H\x1b[J");
         assert_eq!(buffer.title(), "");
-        assert_eq!(n, 6);
+        assert_eq!(n, 9);
 
-        let n = buffer.append(b"\x1b[H\x1b[2Jabc");
+        let n = buffer.append(b"\x1b[H\x1b[2Jabc", 11);
         assert_eq!(buffer.len(), 10);
         assert_eq!(buffer.contents(), b"\x1b[H\x1b[2Jabc");
         assert_eq!(buffer.title(), "");
-        assert_eq!(n, 6);
+        assert_eq!(n, 11);
 
-        let n = buffer.append(b"first\x1bcsecond\x1b[H\x1b[2Jthird");
-        assert_eq!(buffer.len(), 12);
-        assert_eq!(buffer.contents(), b"\x1b[H\x1b[2Jthird");
+        let n = buffer.append(b"first\x1bcsecond\x1b[H\x1b[2Jthird", 10);
+        assert_eq!(buffer.len(), 25);
+        assert_eq!(buffer.contents(), b"first\x1bcsecond\x1b[H\x1b[2Jthird");
         assert_eq!(buffer.title(), "");
         assert_eq!(n, 10);
 
-        let n = buffer.append(b"first\x1b[H\x1b[2Jsecond\x1bcthird");
-        assert_eq!(buffer.len(), 7);
-        assert_eq!(buffer.contents(), b"\x1bcthird");
+        let n = buffer.append(b"first\x1b[H\x1b[2Jsecond\x1bcthird", 25);
+        assert_eq!(buffer.len(), 25);
+        assert_eq!(buffer.contents(), b"first\x1b[H\x1b[2Jsecond\x1bcthird");
         assert_eq!(buffer.title(), "");
-        assert_eq!(n, 12);
+        assert_eq!(n, 25);
     }
 
     #[test]
     fn test_title() {
         let mut buffer = Buffer::new(100);
 
-        let n = buffer.append(b"\x1b]0;this is a title\x07");
+        let n = buffer.append(b"\x1b]0;this is a title\x07", 0);
         assert_eq!(buffer.len(), 20);
         assert_eq!(buffer.contents(), b"\x1b]0;this is a title\x07");
         assert_eq!(buffer.title(), "this is a title");
         assert_eq!(n, 0);
 
-        let n = buffer.append(b"\x1b]2;this is another title\x07");
+        let n = buffer.append(b"\x1b]2;this is another title\x07", 20);
         assert_eq!(buffer.len(), 46);
         assert_eq!(
             buffer.contents(),
@@ -226,14 +233,14 @@ mod test {
         assert_eq!(buffer.title(), "this is another title");
         assert_eq!(n, 0);
 
-        let n = buffer.append(b"\x1bcfoo");
+        let n = buffer.append(b"\x1bcfoo", 46);
         assert_eq!(buffer.len(), 5);
         assert_eq!(buffer.contents(), b"\x1bcfoo");
         assert_eq!(buffer.title(), "this is another title");
         assert_eq!(n, 46);
 
         let n = buffer
-            .append(b"\x1bcabc\x1b]0;title1\x07def\x1b]2;title2\x07ghi");
+            .append(b"\x1bcabc\x1b]0;title1\x07def\x1b]2;title2\x07ghi", 5);
         assert_eq!(buffer.len(), 33);
         assert_eq!(
             buffer.contents(),
@@ -243,7 +250,7 @@ mod test {
         assert_eq!(n, 5);
 
         let n = buffer
-            .append(b"\x1bcabc\x1b]2;title3\x07def\x1b]0;title4\x07ghi");
+            .append(b"\x1bcabc\x1b]2;title3\x07def\x1b]0;title4\x07ghi", 33);
         assert_eq!(buffer.len(), 33);
         assert_eq!(
             buffer.contents(),
@@ -257,13 +264,13 @@ mod test {
     fn test_size_limit() {
         let mut buffer = Buffer::new(100);
 
-        let n = buffer.append(b"foobarbazq");
+        let n = buffer.append(b"foobarbazq", 0);
         assert_eq!(buffer.len(), 10);
         assert_eq!(buffer.contents(), b"foobarbazq");
         assert_eq!(buffer.title(), "");
         assert_eq!(n, 0);
 
-        let n = buffer.append("0123456789".repeat(9).as_ref());
+        let n = buffer.append("0123456789".repeat(9).as_ref(), 10);
         assert_eq!(buffer.len(), 100);
         assert_eq!(
             buffer.contents(),
@@ -272,7 +279,7 @@ mod test {
         assert_eq!(buffer.title(), "");
         assert_eq!(n, 0);
 
-        let n = buffer.append(b"z");
+        let n = buffer.append(b"z", 100);
         assert_eq!(buffer.len(), 50);
         assert_eq!(
             buffer.contents(),
@@ -281,17 +288,57 @@ mod test {
         assert_eq!(buffer.title(), "");
         assert_eq!(n, 51);
 
-        let n = buffer.append("abcdefghij".repeat(15).as_ref());
-        assert_eq!(buffer.len(), 50);
+        let n = buffer.append("abcdefghij".repeat(15).as_ref(), 50);
+        assert_eq!(buffer.len(), 150);
         assert_eq!(
             buffer.contents(),
-            &b"abcdefghijabcdefghijabcdefghijabcdefghijabcdefghij"[..]
+            &b"abcdefghijabcdefghijabcdefghijabcdefghijabcdefghijabcdefghijabcdefghijabcdefghijabcdefghijabcdefghijabcdefghijabcdefghijabcdefghijabcdefghijabcdefghij"[..]
         );
         assert_eq!(buffer.title(), "");
-        // the return value of append represents how many characters from the
-        // beginning of the previous .contents() value no longer exist, so it
-        // should never return a greater value than the previous length of
-        // .contents()
         assert_eq!(n, 50);
     }
+
+    #[test]
+    fn test_written() {
+        let mut buffer = Buffer::new(100);
+
+        let n = buffer.append("abcdefghij".repeat(15).as_ref(), 0);
+        assert_eq!(buffer.len(), 150);
+        assert_eq!(
+            buffer.contents(),
+            &b"abcdefghijabcdefghijabcdefghijabcdefghijabcdefghijabcdefghijabcdefghijabcdefghijabcdefghijabcdefghijabcdefghijabcdefghijabcdefghijabcdefghijabcdefghij"[..]
+        );
+        assert_eq!(buffer.title(), "");
+        assert_eq!(n, 0);
+
+        let mut buffer2 = buffer.clone();
+        let n = buffer2.append(b"z", 0);
+        assert_eq!(buffer2.len(), 151);
+        assert_eq!(
+            buffer2.contents(),
+            &b"abcdefghijabcdefghijabcdefghijabcdefghijabcdefghijabcdefghijabcdefghijabcdefghijabcdefghijabcdefghijabcdefghijabcdefghijabcdefghijabcdefghijabcdefghijz"[..]
+        );
+        assert_eq!(buffer2.title(), "");
+        assert_eq!(n, 0);
+
+        let mut buffer2 = buffer.clone();
+        let n = buffer2.append(b"z", 20);
+        assert_eq!(buffer2.len(), 131);
+        assert_eq!(
+            buffer2.contents(),
+            &b"abcdefghijabcdefghijabcdefghijabcdefghijabcdefghijabcdefghijabcdefghijabcdefghijabcdefghijabcdefghijabcdefghijabcdefghijabcdefghijz"[..]
+        );
+        assert_eq!(buffer2.title(), "");
+        assert_eq!(n, 20);
+
+        let mut buffer2 = buffer.clone();
+        let n = buffer2.append(b"z", 130);
+        assert_eq!(buffer2.len(), 50);
+        assert_eq!(
+            buffer2.contents(),
+            &b"bcdefghijabcdefghijabcdefghijabcdefghijabcdefghijz"[..]
+        );
+        assert_eq!(buffer2.title(), "");
+        assert_eq!(n, 101);
+    }
 }