From 9db9638bc360b0f7d39141aa5adaae9f9291bcad Mon Sep 17 00:00:00 2001 From: Kovid Goyal Date: Sat, 11 Sep 2021 10:14:21 +0530 Subject: [PATCH] Soak up protocol responses after a cancel to avoid outputting garbage to the shell after the kitten exits --- kittens/transfer/main.py | 35 ++++++++++++++++++++++++++++---- kitty/file_transmission.py | 8 +++++--- kitty_tests/file_transmission.py | 2 +- 3 files changed, 37 insertions(+), 8 deletions(-) diff --git a/kittens/transfer/main.py b/kittens/transfer/main.py index de382603c..3a18614ea 100644 --- a/kittens/transfer/main.py +++ b/kittens/transfer/main.py @@ -292,6 +292,7 @@ class SendState(NameReprEnum): waiting_for_permission = auto() permission_granted = auto() permission_denied = auto() + canceled = auto() finished = auto() @@ -385,6 +386,7 @@ class Send(Handler): self.manager = manager self.cli_opts = cli_opts self.transmit_started = False + self.file_metadata_sent = False def send_payload(self, payload: str) -> None: self.write(f'\x1b]{FILE_TRANSFER_CODE};id={self.manager.request_id};') @@ -394,6 +396,11 @@ class Send(Handler): def on_file_transfer_response(self, ftc: FileTransmissionCommand) -> None: if ftc.id != self.manager.request_id: return + if ftc.status == 'CANCELED': + self.quit_loop(1) + return + if self.manager.state in (SendState.finished, SendState.canceled): + return before = self.manager.state self.manager.on_file_transfer_response(ftc) if before == SendState.waiting_for_permission: @@ -405,6 +412,7 @@ class Send(Handler): if self.manager.state == SendState.permission_granted: self.cmd.styled('Permission granted for this transfer', fg='green') self.print() + self.send_file_metadata() self.loop_tick() def check_for_transmit_ok(self) -> None: @@ -435,6 +443,8 @@ class Send(Handler): self.loop_tick() def loop_tick(self) -> None: + if self.manager.state == SendState.waiting_for_permission: + return if self.transmit_started: self.transmit_next_chunk() else: @@ -442,17 +452,34 @@ class Send(Handler): def initialize(self) -> None: self.send_payload(self.manager.start_transfer()) - for payload in self.manager.send_file_metadata(): - self.send_payload(payload) + if self.cli_opts.permissions_password: + # dont wait for permission, not needed with a password and + # avoids a roundtrip + self.send_file_metadata() + + def send_file_metadata(self) -> None: + if not self.file_metadata_sent: + for payload in self.manager.send_file_metadata(): + self.send_payload(payload) + self.file_metadata_sent = True + + def on_term(self) -> None: + self.cmd.styled('Terminate requested, cancelling transfer, transferred files are in undefined state', fg='red') + self.print() + self.abort_transfer(delay=2) def on_interrupt(self) -> None: + if self.manager.state is SendState.canceled: + self.print('Waiting for canceled acknowledgement from terminal, will abort in a few seconds if no response received') + return self.cmd.styled('Interrupt requested, cancelling transfer, transferred files are in undefined state', fg='red') self.print() self.abort_transfer() - def abort_transfer(self) -> None: + def abort_transfer(self, delay: float = 5) -> None: self.send_payload(FileTransmissionCommand(action=Action.cancel).serialize()) - self.quit_loop(1) + self.manager.state = SendState.canceled + self.asyncio_loop.call_later(delay, self.quit_loop, 1) def send_main(cli_opts: TransferCLIOptions, args: List[str]) -> None: diff --git a/kitty/file_transmission.py b/kitty/file_transmission.py index 357fbd296..44416f01e 100644 --- a/kitty/file_transmission.py +++ b/kitty/file_transmission.py @@ -61,7 +61,7 @@ class TransmissionType(NameReprEnum): rsync = auto() -ErrorCode = Enum('ErrorCode', 'OK STARTED EINVAL EPERM EISDIR') +ErrorCode = Enum('ErrorCode', 'OK STARTED CANCELED EINVAL EPERM EISDIR') class TransmissionError(Exception): @@ -428,13 +428,13 @@ class FileTransmission: self.drop_receive(cmd.id) return if not ar.accepted: - log_error(f'File transmission command received for pending id: {cmd.id}, aborting') + log_error(f'File transmission command {cmd.action} received for pending id: {cmd.id}, aborting') self.drop_receive(cmd.id) return ar.last_activity_at = monotonic() else: if cmd.action is not Action.send: - log_error(f'File transmission command received for unknown or rejected id: {cmd.id}, ignoring') + log_error(f'File transmission command {cmd.action} received for unknown or rejected id: {cmd.id}, ignoring') return if len(self.active_receives) >= MAX_ACTIVE_RECEIVES: log_error('New File transmission send with too many active receives, ignoring') @@ -445,6 +445,8 @@ class FileTransmission: if cmd.action is Action.cancel: self.drop_receive(ar.id) + if ar.send_acknowledgements: + self.send_status_response(ErrorCode.CANCELED, request_id=ar.id) elif cmd.action is Action.file: try: df = ar.start_file(cmd) diff --git a/kitty_tests/file_transmission.py b/kitty_tests/file_transmission.py index 4410b40b0..cf6d645eb 100644 --- a/kitty_tests/file_transmission.py +++ b/kitty_tests/file_transmission.py @@ -118,7 +118,7 @@ class TestFileTransmission(BaseTest): ft.handle_serialized_command(serialized_cmd(action='data', data='abcd')) self.assertTrue(os.path.exists(dest)) ft.handle_serialized_command(serialized_cmd(action='cancel')) - self.ae(ft.test_responses, [response(status='OK'), response(status='STARTED', name=dest)]) + self.ae(ft.test_responses, [response(status='OK'), response(status='STARTED', name=dest), response(status='CANCELED')]) self.assertFalse(ft.active_receives) # compress with zlib ft = FileTransmission()