Simplify serialization code and avoid extra copy
This commit is contained in:
parent
0cda5d43a6
commit
a960937095
@ -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")
|
||||
}
|
||||
|
||||
@ -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 {
|
||||
|
||||
@ -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 {
|
||||
|
||||
@ -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"`
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user