Simplify serialization code and avoid extra copy

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)
+	return json.Marshal(rc)
-	if err != nil {
-		return
-	}
-	ans = wrap_in_escape_code(ans)
-	return
 }
 
 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)
+			return json.Marshal(ec)
-			if err != nil {
-				return
-			}
-			ans = wrap_in_escape_code(ans)
-			return
 		}
 		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
+	serializer             wrapped_serializer
-	next_block             func(rc *utils.RemoteControlCmd, serializer serializer_func) (b []byte, err error)
 	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")
 		}

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, err := io_data.serializer.serializer(&rc)
-	raw = raw[len("\x1bP@kitty-cmd") : len(raw)-2]
 	var ec utils.EncryptedRemoteControlCmd
 	err = json.Unmarshal([]byte(raw), &ec)
 	if err != nil {

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 {

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"`