Work on supporting streaming remote commands with passwords

This commit is contained in:
Kovid Goyal
2022-08-31 20:03:51 +05:30
parent d7985689c9
commit 364533b1ed
14 changed files with 206 additions and 85 deletions

View File

@@ -9,6 +9,7 @@ import (
"fmt"
"io"
"os"
"reflect"
"strings"
"time"
@@ -41,6 +42,14 @@ type GlobalOptions struct {
var global_options GlobalOptions
func set_payload_string_field(io_data *rc_io_data, field, data string) {
payload_interface := reflect.ValueOf(&io_data.rc.Payload).Elem()
struct_in_interface := reflect.New(payload_interface.Elem().Type()).Elem()
struct_in_interface.Set(payload_interface.Elem()) // copies the payload to struct_in_interface
struct_in_interface.FieldByName(field).SetString(data)
payload_interface.Set(struct_in_interface) // copies struct_in_interface back to payload
}
func get_pubkey(encoded_key string) (encryption_version string, pubkey []byte, err error) {
if encoded_key == "" {
encoded_key = os.Getenv("KITTY_PUBLIC_KEY")
@@ -171,7 +180,8 @@ type rc_io_data struct {
multiple_payload_generator func(io_data *rc_io_data) (bool, error)
}
func (self *rc_io_data) next_chunk() (chunk []byte, err error) {
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

View File

@@ -18,9 +18,7 @@ func parse_send_text(io_data *rc_io_data, args []string) error {
if len(args) > 0 {
text := strings.Join(args, " ")
text_gen := func(io_data *rc_io_data) (bool, error) {
payload := io_data.rc.Payload.(send_text_json_type)
payload.Data = "text:" + text[:2048]
io_data.rc.Payload = payload
set_payload_data(io_data, "text:"+text[:2048])
text = text[2048:]
return len(text) == 0, nil
}
@@ -38,9 +36,7 @@ func parse_send_text(io_data *rc_io_data, args []string) error {
if err != nil && !errors.Is(err, io.EOF) {
return false, err
}
payload := io_data.rc.Payload.(send_text_json_type)
payload.Data = "base64:" + base64.StdEncoding.EncodeToString(chunk[:n])
io_data.rc.Payload = payload
set_payload_data(io_data, "base64:"+base64.StdEncoding.EncodeToString(chunk[:n]))
return n == 0, nil
}
generators = append(generators, file_gen)
@@ -48,9 +44,7 @@ func parse_send_text(io_data *rc_io_data, args []string) error {
io_data.multiple_payload_generator = func(io_data *rc_io_data) (bool, error) {
if len(generators) == 0 {
payload := io_data.rc.Payload.(send_text_json_type)
payload.Data = "text:"
io_data.rc.Payload = payload
set_payload_data(io_data, "text:")
return true, nil
}
finished, err := generators[0](io_data)

View File

@@ -16,5 +16,6 @@ func parse_set_font_size(arg string, io_data *rc_io_data) error {
return err
}
payload.Size = val
io_data.rc.Payload = payload
return nil
}

View File

@@ -11,10 +11,18 @@ import (
"strings"
)
type struct_with_data interface {
SetData(data string)
}
func set_payload_data(io_data *rc_io_data, data string) {
set_payload_string_field(io_data, "Data", data)
}
func read_window_logo(path string) (func(io_data *rc_io_data) (bool, error), error) {
if strings.ToLower(path) == "none" {
return func(io_data *rc_io_data) (bool, error) {
io_data.rc.Payload = "-"
set_payload_data(io_data, "-")
return true, nil
}, nil
}
@@ -35,16 +43,20 @@ func read_window_logo(path string) (func(io_data *rc_io_data) (bool, error), err
f.Close()
return nil, fmt.Errorf("%s is not a PNG image", path)
}
is_first_call := true
return func(io_data *rc_io_data) (bool, error) {
payload := io_data.rc.Payload.(set_window_logo_json_type)
if is_first_call {
is_first_call = false
} else {
io_data.rc.Stream = false
}
if len(buf) == 0 {
payload.Data = ""
io_data.rc.Payload = payload
set_payload_data(io_data, "")
io_data.rc.Stream = false
return true, nil
}
payload.Data = base64.StdEncoding.EncodeToString(buf)
io_data.rc.Payload = payload
set_payload_data(io_data, base64.StdEncoding.EncodeToString(buf))
buf = buf[:cap(buf)]
n, err := f.Read(buf)
if err != nil && err != io.EOF {

View File

@@ -57,7 +57,7 @@ func read_response_from_conn(conn *net.Conn, timeout time.Duration) (serialized_
func simple_socket_io(conn *net.Conn, io_data *rc_io_data) (serialized_response []byte, err error) {
for {
var chunk []byte
chunk, err = io_data.next_chunk()
chunk, _, err = io_data.next_chunk()
if err != nil {
return
}

View File

@@ -43,6 +43,14 @@ func create_rc_CMD_NAME(args []string) (*utils.RemoteControlCmd, error) {
Cmd: "CLI_NAME",
Version: ProtocolVersion,
NoResponse: NO_RESPONSE_BASE,
Stream: STREAM_WANTED,
}
if rc.Stream {
stream_id, err := utils.HumanRandomId(128)
if err != nil {
return nil, err
}
rc.StreamId = stream_id
}
if IS_ASYNC {
async_id, err := utils.HumanRandomId(128)

View File

@@ -3,12 +3,27 @@
package at
import (
"encoding/json"
"os"
"time"
"kitty/tools/tui/loop"
)
type stream_response struct {
Ok bool `json:"ok"`
Stream bool `json:"stream"`
}
func is_stream_response(serialized_response []byte) bool {
var response stream_response
if len(serialized_response) > 32 {
return false
}
err := json.Unmarshal(serialized_response, &response)
return err == nil && response.Stream
}
func do_chunked_io(io_data *rc_io_data) (serialized_response []byte, err error) {
serialized_response = make([]byte, 0)
lp, err := loop.New()
@@ -17,11 +32,21 @@ func do_chunked_io(io_data *rc_io_data) (serialized_response []byte, err error)
return
}
const (
BEFORE_FIRST_ESCAPE_CODE_SENT = iota
WAITING_FOR_STREAMING_RESPONSE
SENDING
WAITING_FOR_RESPONSE
)
state := BEFORE_FIRST_ESCAPE_CODE_SENT
var last_received_data_at time.Time
var final_write_id loop.IdType
var check_for_timeout func(timer_id loop.IdType) error
wants_streaming := false
check_for_timeout = func(timer_id loop.IdType) error {
if state != WAITING_FOR_RESPONSE && state != WAITING_FOR_STREAMING_RESPONSE {
return nil
}
time_since_last_received_data := time.Now().Sub(last_received_data_at)
if time_since_last_received_data >= io_data.timeout {
return os.ErrDeadlineExceeded
@@ -31,7 +56,7 @@ func do_chunked_io(io_data *rc_io_data) (serialized_response []byte, err error)
}
transition_to_read := func() {
if io_data.rc.NoResponse {
if state == WAITING_FOR_RESPONSE && io_data.rc.NoResponse {
lp.Quit(0)
}
last_received_data_at = time.Now()
@@ -44,36 +69,48 @@ func do_chunked_io(io_data *rc_io_data) (serialized_response []byte, err error)
}
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
}
write_id := lp.QueueWriteBytesDangerous(chunk)
lp.QueueWriteBytesDangerous(chunk)
if len(chunk) == 0 {
final_write_id = write_id
state = WAITING_FOR_RESPONSE
transition_to_read()
}
return "", nil
}
lp.OnWriteComplete = func(completed_write_id loop.IdType) error {
if final_write_id > 0 {
if completed_write_id == final_write_id {
transition_to_read()
}
if state == WAITING_FOR_STREAMING_RESPONSE || state == WAITING_FOR_RESPONSE {
return nil
}
chunk, err := io_data.next_chunk()
chunk, one_escape_code_done, err := io_data.next_chunk()
if err != nil {
return err
}
write_id := lp.QueueWriteBytesDangerous(chunk)
lp.QueueWriteBytesDangerous(chunk)
if len(chunk) == 0 {
final_write_id = write_id
state = WAITING_FOR_RESPONSE
transition_to_read()
}
if one_escape_code_done && state == BEFORE_FIRST_ESCAPE_CODE_SENT {
if wants_streaming {
state = WAITING_FOR_STREAMING_RESPONSE
transition_to_read()
} else {
state = SENDING
}
}
return nil
}
lp.OnRCResponse = func(raw []byte) error {
if state == WAITING_FOR_STREAMING_RESPONSE && is_stream_response(raw) {
state = SENDING
return lp.OnWriteComplete(0)
}
serialized_response = raw
lp.Quit(0)
return nil