Simplify escape code wrapping

tools/cmd/at/main.go

@@ -81,12 +81,6 @@ func simple_serializer(rc *utils.RemoteControlCmd) (ans []byte, err error) {
 
 type serializer_func func(rc *utils.RemoteControlCmd) ([]byte, error)
 
-type wrapped_serializer struct {
-	state             int
-	serializer        serializer_func
-	all_payloads_done bool
-}
-
 func debug_to_log(args ...interface{}) {
 	f, err := os.OpenFile("/tmp/kdlog", os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0666)
 	if err == nil {
@@ -95,49 +89,16 @@ func debug_to_log(args ...interface{}) {
 	}
 }
 
-func (self *wrapped_serializer) next(io_data *rc_io_data) ([]byte, error) {
-	const prefix = "\x1bP@kitty-cmd"
-	const suffix = "\x1b\\"
-	switch self.state {
-	case 0:
-		self.state++
-		return []byte(prefix), nil
-	case 1:
-		if io_data.multiple_payload_generator != nil {
-			is_last, err := io_data.multiple_payload_generator(io_data)
-			if err != nil {
-				return nil, err
-			}
-			if is_last {
-				self.all_payloads_done = true
-			}
-		} else {
-			self.all_payloads_done = true
-		}
-		self.state++
-		return self.serializer(io_data.rc)
-	case 2:
-		if self.all_payloads_done {
-			self.state++
-		} else {
-			self.state = 0
-		}
-		return []byte(suffix), nil
-	default:
-		return make([]byte, 0), nil
-	}
-}
-
 var serializer serializer_func = simple_serializer
 
 func create_serializer(password string, encoded_pubkey string, io_data *rc_io_data) (err error) {
-	io_data.serializer.serializer = simple_serializer
+	io_data.serializer = simple_serializer
 	if password != "" {
 		encryption_version, pubkey, err := get_pubkey(encoded_pubkey)
 		if err != nil {
 			return err
 		}
-		io_data.serializer.serializer = func(rc *utils.RemoteControlCmd) (ans []byte, err error) {
+		io_data.serializer = func(rc *utils.RemoteControlCmd) (ans []byte, err error) {
 			ec, err := crypto.Encrypt_cmd(rc, global_options.password, pubkey, encryption_version)
 			if err != nil {
 				return
@@ -181,22 +142,31 @@ type Response struct {
 type rc_io_data struct {
 	cmd                        *cobra.Command
 	rc                         *utils.RemoteControlCmd
-	serializer                 wrapped_serializer
+	serializer                 serializer_func
 	send_keypresses            bool
 	string_response_is_err     bool
 	timeout                    time.Duration
 	multiple_payload_generator func(io_data *rc_io_data) (bool, error)
+
+	chunks_done bool
 }
 
-func (self *rc_io_data) next_chunk() (chunk []byte, one_escape_code_done bool, err error) {
+func (self *rc_io_data) next_chunk() (chunk []byte, err error) {
-	one_escape_code_done = self.serializer.state == 2
+	if self.chunks_done {
-	block, err := self.serializer.next(self)
+		return make([]byte, 0), nil
-	if err != nil && !errors.Is(err, io.EOF) {
-		return
 	}
-	err = nil
+	if self.multiple_payload_generator != nil {
-	chunk = block
+		is_last, err := self.multiple_payload_generator(self)
-	return
+		if err != nil {
+			return nil, err
+		}
+		if is_last {
+			self.chunks_done = true
+		}
+		return self.serializer(self.rc)
+	}
+	self.chunks_done = true
+	return self.serializer(self.rc)
 }
 
 func get_response(do_io func(io_data *rc_io_data) ([]byte, error), io_data *rc_io_data) (ans *Response, err error) {
@@ -207,7 +177,7 @@ func get_response(do_io func(io_data *rc_io_data) ([]byte, error), io_data *rc_i
 			io_data.rc.CancelAsync = true
 			io_data.multiple_payload_generator = nil
 			io_data.rc.NoResponse = true
-			io_data.serializer.state = 0
+			io_data.chunks_done = false
 			do_io(io_data)
 			err = fmt.Errorf("Timed out waiting for a response from kitty")
 		}

tools/cmd/at/socket_io.go

@@ -30,6 +30,17 @@ func write_all_to_conn(conn *net.Conn, data []byte) error {
 	return nil
 }
 
+func write_many_to_conn(conn *net.Conn, datums ...[]byte) error {
+	for len(datums) > 0 {
+		err := write_all_to_conn(conn, datums[0])
+		if err != nil {
+			return err
+		}
+		datums = datums[1:]
+	}
+	return nil
+}
+
 func read_response_from_conn(conn *net.Conn, timeout time.Duration) (serialized_response []byte, err error) {
 	p := wcswidth.EscapeCodeParser{}
 	keep_going := true
@@ -54,6 +65,9 @@ func read_response_from_conn(conn *net.Conn, timeout time.Duration) (serialized_
 	return
 }
 
+const cmd_escape_code_prefix = "\x1bP@kitty-cmd"
+const cmd_escape_code_suffix = "\x1b\\"
+
 func simple_socket_io(conn *net.Conn, io_data *rc_io_data) (serialized_response []byte, err error) {
 	const (
 		BEFORE_FIRST_ESCAPE_CODE_SENT = iota
@@ -65,8 +79,7 @@ func simple_socket_io(conn *net.Conn, io_data *rc_io_data) (serialized_response
 	wants_streaming := io_data.rc.Stream
 	for {
 		var chunk []byte
-		var one_escape_code_done bool
+		chunk, err = io_data.next_chunk()
-		chunk, one_escape_code_done, err = io_data.next_chunk()
 		if err != nil {
 			return
 		}
@@ -74,11 +87,11 @@ func simple_socket_io(conn *net.Conn, io_data *rc_io_data) (serialized_response
 			state = WAITING_FOR_RESPONSE
 			break
 		}
-		err = write_all_to_conn(conn, chunk)
+		err = write_many_to_conn(conn, []byte(cmd_escape_code_prefix), chunk, []byte(cmd_escape_code_suffix))
 		if err != nil {
 			return
 		}
-		if state == BEFORE_FIRST_ESCAPE_CODE_SENT && one_escape_code_done {
+		if state == BEFORE_FIRST_ESCAPE_CODE_SENT {
 			if wants_streaming {
 				var streaming_response []byte
 				streaming_response, err = read_response_from_conn(conn, io_data.timeout)

tools/cmd/at/tty_io.go

@@ -67,13 +67,19 @@ func do_chunked_io(io_data *rc_io_data) (serialized_response []byte, err error)
 		return nil
 	}
 
+	queue_escape_code := func(data []byte) {
+		lp.QueueWriteString(cmd_escape_code_prefix)
+		lp.QueueWriteBytesDangerous(data)
+		lp.QueueWriteString(cmd_escape_code_suffix)
+	}
+
 	lp.OnInitialize = func() (string, error) {
-		chunk, _, err := io_data.next_chunk()
+		chunk, err := io_data.next_chunk()
 		wants_streaming = io_data.rc.Stream
 		if err != nil {
 			return "", err
 		}
-		lp.QueueWriteBytesDangerous(chunk)
+		queue_escape_code(chunk)
 		if len(chunk) == 0 {
 			state = WAITING_FOR_RESPONSE
 			transition_to_read()
@@ -85,16 +91,16 @@ func do_chunked_io(io_data *rc_io_data) (serialized_response []byte, err error)
 		if state == WAITING_FOR_STREAMING_RESPONSE || state == WAITING_FOR_RESPONSE {
 			return nil
 		}
-		chunk, one_escape_code_done, err := io_data.next_chunk()
+		chunk, err := io_data.next_chunk()
 		if err != nil {
 			return err
 		}
-		lp.QueueWriteBytesDangerous(chunk)
+		queue_escape_code(chunk)
 		if len(chunk) == 0 {
 			state = WAITING_FOR_RESPONSE
 			transition_to_read()
 		}
-		if one_escape_code_done && state == BEFORE_FIRST_ESCAPE_CODE_SENT {
+		if state == BEFORE_FIRST_ESCAPE_CODE_SENT {
 			if wants_streaming {
 				state = WAITING_FOR_STREAMING_RESPONSE
 				transition_to_read()