From 9aa7d4c6a6a0b39a44c466a44e6b12329ce2ef75 Mon Sep 17 00:00:00 2001
From: Etienne Millon <me@emillon.org>
Date: Thu, 31 Mar 2022 10:02:30 +0200
Subject: [PATCH] Ansi_color.parse: handle CSI n K

CSI n K is used to clear the current line. It was not properly parsed by
`Ansi_color.parse`. It is now ignored.

Closes #5528
Signed-off-by: Etienne Millon <me@emillon.org>
---
 CHANGES.md                                  |  3 ++
 otherlibs/stdune/ansi_color.ml              | 18 ++++++++--
 test/blackbox-tests/test-cases/github5528.t | 37 +++++++++++++++++++++
 3 files changed, 55 insertions(+), 3 deletions(-)
 create mode 100644 test/blackbox-tests/test-cases/github5528.t

diff --git a/CHANGES.md b/CHANGES.md
index a1813b46f0d..724ebfebb4e 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -91,6 +91,9 @@
 
 - Allow include statement in install stanza (#6139, fixes #256, @gridbugs)
 
+- Handle CSI n K code in ANSI escape codes from commands. (#6214, fixes #5528,
+  @emillon)
+
 3.4.1 (26-07-2022)
 ------------------
 
diff --git a/otherlibs/stdune/ansi_color.ml b/otherlibs/stdune/ansi_color.ml
index 9a4842d56e6..621915fe026 100644
--- a/otherlibs/stdune/ansi_color.ml
+++ b/otherlibs/stdune/ansi_color.ml
@@ -179,6 +179,17 @@ let strip str =
   in
   loop 0
 
+let index_from_any str start chars =
+  let n = String.length str in
+  let rec go i =
+    if i >= n then None
+    else
+      match List.find chars ~f:(fun c -> Char.equal str.[i] c) with
+      | None -> go (i + 1)
+      | Some c -> Some (i, c)
+  in
+  go start
+
 let parse_line str styles =
   let len = String.length str in
   let add_chunk acc ~styles ~pos ~len =
@@ -201,9 +212,9 @@ let parse_line str styles =
       let seq_start = seq_start + 2 in
       if seq_start >= len || str.[seq_start - 1] <> '[' then (styles, acc)
       else
-        match String.index_from str seq_start 'm' with
+        match index_from_any str seq_start [ 'm'; 'K' ] with
         | None -> (styles, acc)
-        | Some seq_end ->
+        | Some (seq_end, 'm') ->
           let styles =
             if seq_start = seq_end then
               (* Some commands output "\027[m", which seems to be interpreted
@@ -223,7 +234,8 @@ let parse_line str styles =
                      else s :: styles)
               |> List.rev
           in
-          loop styles (seq_end + 1) acc)
+          loop styles (seq_end + 1) acc
+        | Some (seq_end, _) -> loop styles (seq_end + 1) acc)
   in
   loop styles 0 Pp.nop
 
diff --git a/test/blackbox-tests/test-cases/github5528.t b/test/blackbox-tests/test-cases/github5528.t
new file mode 100644
index 00000000000..0b7eb8f30ad
--- /dev/null
+++ b/test/blackbox-tests/test-cases/github5528.t
@@ -0,0 +1,37 @@
+  $ cat > dune-project <<EOF
+  > (lang dune 1.0)
+  > EOF
+
+  $ cat > dune <<EOF
+  > (test
+  >   (name t))
+  > EOF
+
+  $ cat > t.ml <<EOF
+  > type color =  Normal | Cyan
+  > 
+  > let int_of_color = function
+  >   | Normal -> 0
+  >   | Cyan   -> 6
+  > 
+  > let in_color c pp out x =
+  >   let n = int_of_color c in
+  >   Printf.fprintf out "\x1b[3%dm" n;
+  >   pp out x;
+  >   Printf.fprintf out "\x1b[0m"
+  > 
+  > let reset_line = "\x1b[2K\r"
+  > 
+  > let () =
+  >   Printf.printf "%sVery Secret!\n%!" reset_line;
+  >   Printf.printf "%s\n%!" (String.make 15 '-');
+  >   Printf.printf "%a\n%!" (in_color Cyan output_string) "Can you see it?"
+  > EOF
+
+  $ dune runtest -f
+  
Very Secret!
+  ---------------
+  Can you see it?
+
+  $ dune exec ./t.exe
+  Can you see it?