diff --git a/tools/cmd/at/main.go b/tools/cmd/at/main.go index 54a39845a..63709782b 100644 --- a/tools/cmd/at/main.go +++ b/tools/cmd/at/main.go @@ -66,47 +66,48 @@ func get_pubkey(encoded_key string) (encryption_version string, pubkey []byte, e return } -func wrap_in_escape_code(data []byte) []byte { - const prefix = "\x1bP@kitty-cmd" - const suffix = "\x1b\\" - ans := make([]byte, len(prefix)+len(data)+len(suffix)) - n := copy(ans, prefix) - n += copy(ans[n:], data) - copy(ans[n:], suffix) - return ans -} - func simple_serializer(rc *utils.RemoteControlCmd) (ans []byte, err error) { - ans, err = json.Marshal(rc) - if err != nil { - return - } - ans = wrap_in_escape_code(ans) - return + return json.Marshal(rc) } type serializer_func func(rc *utils.RemoteControlCmd) ([]byte, error) +type wrapped_serializer struct { + state int + serializer serializer_func +} + +func (self *wrapped_serializer) next(rc *utils.RemoteControlCmd) ([]byte, error) { + const prefix = "\x1bP@kitty-cmd" + const suffix = "\x1b\\" + defer func() { self.state++ }() + switch self.state { + case 0: + return []byte(prefix), nil + case 1: + return self.serializer(rc) + case 2: + 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 = simple_serializer + io_data.serializer.serializer = simple_serializer if password != "" { encryption_version, pubkey, err := get_pubkey(encoded_pubkey) if err != nil { return err } - io_data.serializer = func(rc *utils.RemoteControlCmd) (ans []byte, err error) { + io_data.serializer.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 } - ans, err = json.Marshal(ec) - if err != nil { - return - } - ans = wrap_in_escape_code(ans) - return + return json.Marshal(ec) } if io_data.timeout < 120*time.Second { io_data.timeout = 120 * time.Second @@ -145,8 +146,7 @@ type Response struct { type rc_io_data struct { cmd *cobra.Command rc *utils.RemoteControlCmd - serializer serializer_func - next_block func(rc *utils.RemoteControlCmd, serializer serializer_func) (b []byte, err error) + serializer wrapped_serializer send_keypresses bool string_response_is_err bool timeout time.Duration @@ -161,7 +161,7 @@ func (self *rc_io_data) next_chunk(limit_size bool) (chunk []byte, err error) { self.pending_chunks = self.pending_chunks[:len(self.pending_chunks)-1] return } - block, err := self.next_block(self.rc, self.serializer) + block, err := self.serializer.next(self.rc) if err != nil && !errors.Is(err, io.EOF) { return } @@ -180,14 +180,6 @@ func (self *rc_io_data) next_chunk(limit_size bool) (chunk []byte, err error) { return } -func single_rc_sender(rc *utils.RemoteControlCmd, serializer serializer_func) ([]byte, error) { - if rc.SingleSent() { - return make([]byte, 0), nil - } - rc.SetSingleSent() - return serializer(rc) -} - func get_response(do_io func(io_data *rc_io_data) ([]byte, error), io_data *rc_io_data) (ans *Response, err error) { serialized_response, err := do_io(io_data) if err != nil { @@ -195,7 +187,7 @@ func get_response(do_io func(io_data *rc_io_data) ([]byte, error), io_data *rc_i io_data.rc.Payload = nil io_data.rc.CancelAsync = true io_data.rc.NoResponse = true - io_data.rc.ResetSingleSent() + io_data.serializer.state = 0 do_io(io_data) err = fmt.Errorf("Timed out waiting for a response from kitty") } diff --git a/tools/cmd/at/main_test.go b/tools/cmd/at/main_test.go index 767e1654f..18c6a90ef 100644 --- a/tools/cmd/at/main_test.go +++ b/tools/cmd/at/main_test.go @@ -47,7 +47,7 @@ func TestRCSerialization(t *testing.T) { Cmd: "test", Version: ver, } simple := func(expected string) { - actual, err := io_data.serializer(&rc) + actual, err := io_data.serializer.serializer(&rc) if err != nil { t.Fatal(err) } @@ -56,7 +56,7 @@ func TestRCSerialization(t *testing.T) { t.Fatalf("Incorrect serialization: %s != %s", expected, as) } } - simple(string(wrap_in_escape_code([]byte(`{"cmd":"test","version":[1,2,3]}`)))) + simple(string(`{"cmd":"test","version":[1,2,3]}`)) pubkey_b, _, err := crypto.KeyPair("1") if err != nil { t.Fatal(err) @@ -69,8 +69,7 @@ func TestRCSerialization(t *testing.T) { if err != nil { t.Fatal(err) } - raw, err := io_data.serializer(&rc) - raw = raw[len("\x1bP@kitty-cmd") : len(raw)-2] + raw, err := io_data.serializer.serializer(&rc) var ec utils.EncryptedRemoteControlCmd err = json.Unmarshal([]byte(raw), &ec) if err != nil { diff --git a/tools/cmd/at/template.go b/tools/cmd/at/template.go index ef4bc7e86..280787413 100644 --- a/tools/cmd/at/template.go +++ b/tools/cmd/at/template.go @@ -69,7 +69,6 @@ func run_CMD_NAME(cmd *cobra.Command, args []string) (err error) { rc: rc, timeout: time.Duration(timeout * float64(time.Second)), string_response_is_err: STRING_RESPONSE_IS_ERROR, - next_block: single_rc_sender, } err = create_payload_CMD_NAME(&io_data, args) if err != nil { diff --git a/tools/utils/types.go b/tools/utils/types.go index 6078a5414..fea896d4f 100644 --- a/tools/utils/types.go +++ b/tools/utils/types.go @@ -11,14 +11,8 @@ type RemoteControlCmd struct { Password string `json:"password,omitempty"` Async string `json:"async,omitempty"` CancelAsync bool `json:"cancel_async,omitempty"` - - single_sent bool } -func (self *RemoteControlCmd) SingleSent() bool { return self.single_sent } -func (self *RemoteControlCmd) SetSingleSent() { self.single_sent = true } -func (self *RemoteControlCmd) ResetSingleSent() { self.single_sent = false } - type EncryptedRemoteControlCmd struct { Version [3]int `json:"version"` IV string `json:"iv"`