From 12b0f632bd5bdf40aab902d5eb5cf6c789da2349 Mon Sep 17 00:00:00 2001
From: Kovid Goyal <kovid@kovidgoyal.net>
Date: Mon, 22 Aug 2022 11:21:23 +0530
Subject: [PATCH] Allow OSC to be terminated by BEL

Also prevent changing streaming state unless in reset state
---
 tools/utils/escape_code_parser.go      | 34 +++++++++++++++++++++-----
 tools/utils/escape_code_parser_test.go | 10 +++++---
 2 files changed, 35 insertions(+), 9 deletions(-)

diff --git a/tools/utils/escape_code_parser.go b/tools/utils/escape_code_parser.go
index 45d317bc3..b15991d03 100644
--- a/tools/utils/escape_code_parser.go
+++ b/tools/utils/escape_code_parser.go
@@ -1,6 +1,9 @@
 package utils
 
-import "bytes"
+import (
+	"bytes"
+	"fmt"
+)
 
 type parser_state uint8
 type csi_state uint8
@@ -13,6 +16,7 @@ const (
 	esc
 	csi
 	st
+	st_or_bel
 	esc_st
 	c1_st
 	bracketed_paste
@@ -40,7 +44,7 @@ type EscapeCodeParser struct {
 
 	// Whether to send escape code bytes as soon as they are received or to
 	// buffer and send full escape codes
-	Streaming bool
+	streaming bool
 
 	// Callbacks
 	HandleRune func(rune)
@@ -52,6 +56,18 @@ type EscapeCodeParser struct {
 	HandleAPC  func([]byte)
 }
 
+func (self *EscapeCodeParser) SetStreaming(streaming bool) error {
+	if self.state != normal || len(self.current_buffer) > 0 {
+		return fmt.Errorf("Cannot change streaming state when not in reset state")
+	}
+	self.streaming = streaming
+	return nil
+}
+
+func (self *EscapeCodeParser) IsStreaming() bool {
+	return self.streaming
+}
+
 func (self *EscapeCodeParser) Parse(data []byte) {
 	prev := UTF8_ACCEPT
 	codep := UTF8_ACCEPT
@@ -79,7 +95,7 @@ func (self *EscapeCodeParser) Reset() {
 }
 
 func (self *EscapeCodeParser) write_ch(ch byte) {
-	if self.Streaming {
+	if self.streaming {
 		if self.current_callback != nil {
 			var data [1]byte = [1]byte{ch}
 			self.current_callback(data[:])
@@ -186,7 +202,7 @@ func (self *EscapeCodeParser) dispatch_char(ch UTF8State) {
 			dispatch()
 		}
 		return
-	} // end bracketed_paste
+	} // end self.state == bracketed_paste
 
 	switch ch {
 	case 0x1b:
@@ -198,7 +214,7 @@ func (self *EscapeCodeParser) dispatch_char(ch UTF8State) {
 		self.state = csi
 		self.current_callback = self.HandleCSI
 	case 0x9d:
-		self.state = st
+		self.state = st_or_bel
 		self.current_callback = self.HandleOSC
 	case 0x98:
 		self.state = st
@@ -226,7 +242,7 @@ func (self *EscapeCodeParser) dispatch_byte(ch byte) {
 			self.csi_state = parameter
 			self.current_callback = self.HandleCSI
 		case ']':
-			self.state = st
+			self.state = st_or_bel
 			self.current_callback = self.HandleOSC
 		case '^':
 			self.state = st
@@ -257,6 +273,12 @@ func (self *EscapeCodeParser) dispatch_byte(ch byte) {
 				self.dispatch_esc_code()
 			}
 		}
+	case st_or_bel:
+		if ch == 0x7 {
+			self.dispatch_esc_code()
+			return
+		}
+		fallthrough
 	case st:
 		if ch == 0x1b {
 			self.state = esc_st
diff --git a/tools/utils/escape_code_parser_test.go b/tools/utils/escape_code_parser_test.go
index 36391ec11..3328a0935 100644
--- a/tools/utils/escape_code_parser_test.go
+++ b/tools/utils/escape_code_parser_test.go
@@ -29,9 +29,9 @@ func TestEscapeCodeParsing(t *testing.T) {
 		d = test_parse_collection{}
 	}
 
-	check_test_result := func() {
+	check_test_result := func(raw string) {
 		if d.actual != d.expected {
-			t.Fatalf("actual != expected: %#v != %#v", string(d.actual), string(d.expected))
+			t.Fatalf("parsing: %#v actual != expected: %#v != %#v", raw, string(d.actual), string(d.expected))
 		}
 	}
 
@@ -39,11 +39,15 @@ func TestEscapeCodeParsing(t *testing.T) {
 		reset_test_parser()
 		d.expected = "\n" + expected
 		test_parser.Parse([]byte(raw))
-		check_test_result()
+		check_test_result(raw)
 	}
 
 	test("\x1b[31m\xc2\x9bm", "CSI: 31m\nCSI: m")
 	test("ab\nc", "CH: a\nCH: b\nCH: \n\nCH: c")
 	test("a\x1b[200m\x1b[mb\x1b[5:3;2;4~", "CH: a\nCSI: 200m\nCSI: m\nCH: b\nCSI: 5:3;2;4~")
+	test("\x1b[200~a\x1b[201m\x1b[201~\x1b[x", "CH: a\nCH: \x1b\nCH: [\nCH: 2\nCH: 0\nCH: 1\nCH: m\nCSI: x")
+	test("a\x1bPb\x1b\x1bc\x1b\\d", "CH: a\nDCS: b\x1bc\nCH: d")
+	test("a\x1b_b\x1b\x1b\x1bc\x1b\\d", "CH: a\nAPC: b\x1b\x1bc\nCH: d")
+	test("\x1b]X\x07\x1b]X\x1b\x07\x1b\\", "OSC: X\nOSC: X\x1b\x07")
 
 }