Support more color escape codes (#757)

* Support text cursor color escape codes
* Support reset color index escape code
* Support multiple colors in set color index escape code

1 files changed, 81 insertions, 27 deletions

diff --git a/src/ansi.rs b/src/ansi.rs
index 76064f18..2b196145 100644
--- a/src/ansi.rs
+++ b/src/ansi.rs
@@ -78,6 +78,25 @@ fn parse_rgb_color(color: &[u8]) -> Option<Rgb> {
     }
 }
 
+fn parse_number(input: &[u8]) -> Option<u8> {
+    if input.is_empty() {
+        return None;
+    }
+    let mut num: u8 = 0;
+    for c in input {
+        let c = *c as char;
+        if let Some(digit) = c.to_digit(10) {
+            num = match num.checked_mul(10).and_then(|v| v.checked_add(digit as u8)) {
+                Some(v) => v,
+                None => return None,
+            }
+        } else {
+            return None;
+        }
+    }
+    Some(num)
+}
+
 /// The processor wraps a `vte::Parser` to ultimately call methods on a Handler
 pub struct Processor {
     state: ProcessorState,
@@ -730,37 +749,20 @@ impl<'a, H, W> vte::Perform for Performer<'a, H, W>
 
             // Set color index
             b"4" => {
-                if params.len() < 3 {
+                if params.len() == 1 || params.len() % 2 == 0 {
                     return unhandled!();
                 }
 
-                // Parse index
-                let index = {
-                    let raw = params[1];
-                    if raw.len() > 3 {
-                        return unhandled!();
-                    }
-                    let mut index: u8 = 0;
-                    for c in raw {
-                        let c = *c as char;
-                        if let Some(digit) = c.to_digit(10) {
-                            index *= 10;
-                            index += digit as u8;
+                for chunk in params[1..].chunks(2) {
+                    if let Some(index) = parse_number(chunk[0]) {
+                        if let Some(color) = parse_rgb_color(chunk[1]) {
+                            self.handler.set_color(index as usize, color);
+                        } else {
+                            unhandled!();
                         }
+                    } else {
+                        unhandled!();
                     }
-
-                    index
-                };
-
-                // Check that index is a valid u8
-                if index & 0xff != index {
-                    return unhandled!();
-                }
-
-                if let Some(color) = parse_rgb_color(params[2]) {
-                    self.handler.set_color(index as usize, color);
-                } else {
-                    unhandled!();
                 }
             }
 
@@ -790,6 +792,38 @@ impl<'a, H, W> vte::Perform for Performer<'a, H, W>
                 }
             }
 
+            // Set text cursor color
+            b"12" => {
+                if params.len() < 2 {
+                    return unhandled!();
+                }
+
+                if let Some(color) = parse_rgb_color(params[1]) {
+                    self.handler.set_color(NamedColor::Cursor as usize, color);
+                } else {
+                    unhandled!()
+                }
+            }
+
+            // Reset color index
+            b"104" => {
+                // Reset all color indexes when no parameters are given
+                if params.len() == 1 {
+                    for i in 0..256 {
+                        self.handler.reset_color(i);
+                    }
+                    return;
+                }
+
+                // Reset color indexes given as parameters
+                for param in &params[1..] {
+                    match parse_number(param) {
+                        Some(index) => self.handler.reset_color(index as usize),
+                        None => unhandled!(),
+                    }
+                }
+            }
+
             // Reset foreground color
             b"110" => {
                 self.handler.reset_color(NamedColor::Foreground as usize);
@@ -800,6 +834,11 @@ impl<'a, H, W> vte::Perform for Performer<'a, H, W>
                 self.handler.reset_color(NamedColor::Background as usize);
             }
 
+            // Reset text cursor color
+            b"112" => {
+                self.handler.reset_color(NamedColor::Cursor as usize);
+            }
+
             _ => {
                 unhandled!();
             }
@@ -1310,7 +1349,7 @@ pub mod C1 {
 mod tests {
     use std::io;
     use index::{Line, Column};
-    use super::{Processor, Handler, Attr, TermInfo, Color, StandardCharset, CharsetIndex, parse_rgb_color};
+    use super::{Processor, Handler, Attr, TermInfo, Color, StandardCharset, CharsetIndex, parse_rgb_color, parse_number};
     use ::Rgb;
 
     /// The /dev/null of io::Write
@@ -1488,4 +1527,19 @@ mod tests {
     fn parse_valid_rgb_color2() {
         assert_eq!(parse_rgb_color(b"#11aaff"), Some(Rgb { r: 0x11, g: 0xaa, b: 0xff }));
     }
+
+    #[test]
+    fn parse_invalid_number() {
+        assert_eq!(parse_number(b"1abc"), None);
+    }
+
+    #[test]
+    fn parse_valid_number() {
+        assert_eq!(parse_number(b"123"), Some(123));
+    }
+
+    #[test]
+    fn parse_number_too_large() {
+        assert_eq!(parse_number(b"321"), None);
+    }
 }