From 3a198833dae2ca697d1cd34df63774eeec087b2d Mon Sep 17 00:00:00 2001 From: Kovid Goyal Date: Thu, 1 Sep 2022 15:29:20 +0530 Subject: [PATCH] Simplify escape code wrapping --- tools/cmd/at/main.go | 72 ++++++++++++--------------------------- tools/cmd/at/socket_io.go | 21 +++++++++--- tools/cmd/at/tty_io.go | 16 ++++++--- 3 files changed, 49 insertions(+), 60 deletions(-) diff --git a/tools/cmd/at/main.go b/tools/cmd/at/main.go index b6fb4b05e..5fb9b5908 100644 --- a/tools/cmd/at/main.go +++ b/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) { - one_escape_code_done = self.serializer.state == 2 - block, err := self.serializer.next(self) - if err != nil && !errors.Is(err, io.EOF) { - return +func (self *rc_io_data) next_chunk() (chunk []byte, err error) { + if self.chunks_done { + return make([]byte, 0), nil } - err = nil - chunk = block - return + if self.multiple_payload_generator != nil { + is_last, err := self.multiple_payload_generator(self) + 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") } diff --git a/tools/cmd/at/socket_io.go b/tools/cmd/at/socket_io.go index a6209a668..e6b11e989 100644 --- a/tools/cmd/at/socket_io.go +++ b/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, one_escape_code_done, err = io_data.next_chunk() + chunk, 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) diff --git a/tools/cmd/at/tty_io.go b/tools/cmd/at/tty_io.go index 588340d46..374eca198 100644 --- a/tools/cmd/at/tty_io.go +++ b/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()