Merge pull request #30 from kmaork/fixing_todos

Fixing todos

Author: kmaork

.github/workflows/tests.yml

@@ -18,7 +18,7 @@ jobs:
         with:
           python-version: ${{ matrix.python-version }}
       - name: Install dependencies
-        run: pip install tox tox-wheel
+        run: pip install tox tox-gh-actions tox-wheel
       - name: Test sdist with tox
         run: python -m tox
       - name: Upload sdist

madbg/client.py

@@ -7,7 +7,7 @@
 import termios
 from contextlib import contextmanager
 
-from .communication import pipe_until_closed, send_message
+from .communication import Piping, send_message
 from .consts import DEFAULT_IP, DEFAULT_PORT, STDIN_FILENO, STDOUT_FILENO, DEFAULT_CONNECT_TIMEOUT
 
 
@@ -60,7 +60,8 @@ def connect_to_server(ip, port, timeout):
         s.close()
 
 
-def connect_to_debugger(ip=DEFAULT_IP, port=DEFAULT_PORT, timeout=DEFAULT_CONNECT_TIMEOUT):
+def connect_to_debugger(ip=DEFAULT_IP, port=DEFAULT_PORT, timeout=DEFAULT_CONNECT_TIMEOUT,
+                        in_fd=STDIN_FILENO, out_fd=STDOUT_FILENO):
     with connect_to_server(ip, port, timeout) as socket:
         tty_handle = get_tty_handle()
         term_size = os.get_terminal_size(tty_handle)
@@ -71,4 +72,4 @@ def connect_to_debugger(ip=DEFAULT_IP, port=DEFAULT_PORT, timeout=DEFAULT_CONNEC
         send_message(socket, term_data)
         with prepare_terminal():
             socket_fd = socket.fileno()
-            pipe_until_closed({STDIN_FILENO: socket_fd, socket_fd: STDOUT_FILENO})
+            Piping({in_fd: socket_fd, socket_fd: out_fd}).run()

madbg/communication.py

@@ -1,12 +1,14 @@
 import pickle
-import select
 import fcntl
 import os
 import struct
+from collections import defaultdict
 from contextlib import contextmanager
+from functools import partial
+from asyncio import new_event_loop
 from io import BytesIO
 
-from madbg.utils import loop_in_thread, opposite_dict
+from .utils import run_thread, opposite_dict
 
 MESSAGE_LENGTH_FMT = 'I'
 
@@ -28,41 +30,51 @@ def blocking_read(fd, n):
     return io.getvalue()
 
 
-def pipe_once(pipe_dict):
-    # If read fails, write will fail. But if write fail, we might be still able to read.
-    # TODO: use wakeup fd instead of 0 timeout polling (os.pipe? eventfd?)
-    # TODO: can use splice or ebpf
-    if not pipe_dict:
-        return
-    reverse_pipe_dict = opposite_dict(pipe_dict)
-    for read_fd in select.select(list(pipe_dict), [], [], 0)[0]:
-        write_fd = pipe_dict[read_fd]
+class Piping:
+    def __init__(self, pipe_dict):
+        self.buffers = defaultdict(bytes)
+        self.loop = new_event_loop()
+        for src_fd, dest_fd in pipe_dict.items():
+            self.loop.add_reader(src_fd, partial(self._read, src_fd, dest_fd))
+            self.loop.add_writer(dest_fd, partial(self._write, dest_fd))
+        self.readers_to_writers = dict(pipe_dict)
+        self.writers_to_readers = opposite_dict(pipe_dict)
+
+    def _remove_writer(self, writer_fd):
+        self.loop.remove_writer(writer_fd)
+        for reader_fd in self.writers_to_readers.pop(writer_fd):
+            self.readers_to_writers.pop(reader_fd)
+
+    def _remove_reader(self, reader_fd):
+        # remove all writers that im the last to write to, remove all that write to me, if nothing left stop loop
+        self.loop.remove_reader(reader_fd)
+        writer_fd = self.readers_to_writers.pop(reader_fd)
+        writer_readers = self.writers_to_readers[writer_fd]
+        writer_readers.remove(reader_fd)
+        if not writer_fd:
+            self._remove_writer(writer_fd)
+
+    def _read(self, src_fd, dest_fd):
         try:
-            data = os.read(read_fd, 1024)
-            if not data:
-                raise OSError('EOF')
+            data = os.read(src_fd, 1024)
         except OSError:
-            pipe_dict.pop(read_fd, None)
-            for writing_fd in reverse_pipe_dict[read_fd]:
-                pipe_dict.pop(writing_fd, None)
+            data = ''
+        if data:
+            self.buffers[dest_fd] += data
         else:
-            try:
-                os.write(write_fd, data)
-            except OSError:
-                pipe_dict.pop(read_fd, None)
-
-
-@contextmanager
-def pipe_in_background(pipe_dict):
-    pipe_dict = dict(pipe_dict)
-    with loop_in_thread(pipe_once, pipe_dict):
-        yield
-
-
-def pipe_until_closed(pipe_dict):
-    pipe_dict = dict(pipe_dict)
-    while pipe_dict:
-        pipe_once(pipe_dict)
+            self._remove_reader(src_fd)
+            if src_fd in self.writers_to_readers:
+                self._remove_writer(src_fd)
+            if not self.readers_to_writers:
+                self.loop.stop()
+
+    def _write(self, dest_fd):
+        buffer = self.buffers[dest_fd]
+        if buffer:
+            self.buffers[dest_fd] = buffer[os.write(dest_fd, buffer):]
+
+    def run(self):
+        self.loop.run_forever()
 
 
 def send_message(sock, obj):

madbg/debugger.py

@@ -8,10 +8,11 @@
 from IPython.terminal.interactiveshell import TerminalInteractiveShell
 from prompt_toolkit.input.vt100 import Vt100Input
 from prompt_toolkit.output.vt100 import Vt100_Output
+from inspect import currentframe
 
-from .utils import preserve_sys_state, get_client_connection, use_context
-from .tty_utils import print_to_ctty, open_pty, resize_terminal, modify_terminal, set_ctty
-from .communication import receive_message, pipe_in_background
+from .utils import preserve_sys_state, get_client_connection, use_context, run_thread
+from .tty_utils import print_to_ctty, PTY
+from .communication import receive_message, Piping
 
 
 class RemoteIPythonDebugger(TerminalPdb):
@@ -58,7 +59,7 @@ def set_trace(self, frame=None, done_callback=None):
         """ Overriding super to add the done_callback argument, allowing cleanup after a debug session """
         td = lambda *args: self.trace_dispatch(*args, done_callback=done_callback)
         if frame is None:
-            frame = sys._getframe().f_back
+            frame = currentframe().f_back
         self.reset()
         while frame:
             frame.f_trace = td
@@ -116,24 +117,22 @@ def connect_and_set_trace(cls, ip, port, frame=None):
     def start(cls, sock_fd):
         term_data = receive_message(sock_fd)
         term_attrs, term_type, term_size = term_data['term_attrs'], term_data['term_type'], term_data['term_size']
-        with open_pty() as (master_fd, slave_fd):
-            resize_terminal(slave_fd, term_size[0], term_size[1])
-            modify_terminal(slave_fd, term_attrs)
-            set_ctty(slave_fd)
-            with pipe_in_background({sock_fd: master_fd, master_fd: sock_fd}):
-                slave_reader = os.fdopen(slave_fd, 'r')
-                slave_writer = os.fdopen(slave_fd, 'w')
+        with PTY.open() as pty:
+            pty.resize(term_size[0], term_size[1])
+            pty.set_tty_attrs(term_attrs)
+            pty.make_ctty()
+            piping = Piping({sock_fd: pty.master_fd, pty.master_fd: sock_fd})
+            with run_thread(piping.run):
+                slave_reader = os.fdopen(pty.slave_fd, 'r')
+                slave_writer = os.fdopen(pty.slave_fd, 'w')
                 try:
                     yield cls(slave_reader, slave_writer, term_type)
                 except Exception:
                     print(traceback.format_exc(), file=slave_writer)
                     raise
                 finally:
-                    # We flush to make sure the message is written before the last pipe iteration
                     print('Closing connection', file=slave_writer, flush=True)
-                    # TODO: solve race
-                    import time
-                    time.sleep(0.05)
+                    slave_writer.close()
 
     @classmethod
     @contextmanager

madbg/tty_utils.py

@@ -3,6 +3,7 @@
 import pty
 import struct
 from contextlib import contextmanager
+from dataclasses import dataclass
 from multiprocessing.pool import Pool
 import signal
 import fcntl
@@ -77,40 +78,49 @@ def detach_current_ctty():
         os.close(ctty_fd)
 
 
-def set_ctty(fd):
-    detach_current_ctty()
-    attach_ctty(fd)
-
-
-def resize_terminal(fd, rows, cols):
-    winsize = struct.pack("HHHH", rows, cols, 0, 0)
-    fcntl.ioctl(fd, termios.TIOCSWINSZ, winsize)
-
-
-def modify_terminal(fd, tc_attrs, when=termios.TCSANOW):
-    return
-    IFLAG = 0
-    OFLAG = 1
-    CFLAG = 2
-    LFLAG = 3
-    ISPEED = 4
-    OSPEED = 5
-    CC = 6
-    tc_attrs = tc_attrs[:]
-    current_attrs = termios.tcgetattr(fd)
-    tc_attrs[CC] = current_attrs[termios.CC]
-    termios.tcsetattr(fd, when, tc_attrs)
-
-
-@contextmanager
-def open_pty():
-    master_fd, slave_fd = pty.openpty()
-    try:
-        yield master_fd, slave_fd
-    finally:
-        with ignore_signal(signal.SIGHUP):
-            os.close(master_fd)
-            os.close(slave_fd)
+@dataclass
+class PTY:
+    master_fd: int
+    slave_fd: int
+    _closed: bool = False
+
+    def resize(self, rows, cols):
+        winsize = struct.pack("HHHH", rows, cols, 0, 0)
+        fcntl.ioctl(self.slave_fd, termios.TIOCSWINSZ, winsize)
+
+    def close(self):
+        if not self._closed:
+            with ignore_signal(signal.SIGHUP):
+                os.close(self.master_fd)
+            self._closed = True
+
+    def set_tty_attrs(self, tc_attrs, when=termios.TCSANOW):
+        return
+        IFLAG = 0
+        OFLAG = 1
+        CFLAG = 2
+        LFLAG = 3
+        ISPEED = 4
+        OSPEED = 5
+        CC = 6
+        tc_attrs = tc_attrs[:]
+        current_attrs = termios.tcgetattr(fd)
+        tc_attrs[CC] = current_attrs[termios.CC]
+        termios.tcsetattr(self.slave_fd, when, tc_attrs)
+
+    def make_ctty(self):
+        detach_current_ctty()
+        attach_ctty(self.slave_fd)
+
+    @classmethod
+    @contextmanager
+    def open(cls):
+        master_fd, slave_fd = pty.openpty()
+        self = cls(master_fd, slave_fd)
+        try:
+            yield self
+        finally:
+            self.close()
 
 
 def print_to_ctty(string):

madbg/utils.py

@@ -6,6 +6,8 @@
 from concurrent.futures.thread import ThreadPoolExecutor
 from contextlib import contextmanager, ExitStack
 
+from madbg.tty_utils import print_to_ctty
+
 
 @contextmanager
 def preserve_sys_state():
@@ -24,7 +26,8 @@ def get_client_connection(ip, port):
     server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, True)
     server_socket.bind((ip, port))
     server_socket.listen(1)
-    sock, _ = server_socket.accept()
+    sock, (client_ip, client_port) = server_socket.accept()
+    print_to_ctty(f'Client connected from {client_ip}:{client_port}')
     server_socket.close()
     try:
         yield sock.fileno()
@@ -60,25 +63,8 @@ def run_thread(func, *args, **kwargs):
             future.result()
 
 
-@contextmanager
-def loop_in_thread(func, *args, iteration_after_exit=True, **kwargs):
-    cont = True
-
-    def loop():
-        while cont:
-            func(*args, **kwargs)
-
-    with run_thread(loop) as future:
-        try:
-            yield future
-        finally:
-            cont = False
-    if iteration_after_exit:
-        func(*args, **kwargs)
-
-
-def opposite_dict(dikt):
-    opposite = defaultdict(list)
-    for key, value in dikt.items():
-        opposite[value].append(key)
+def opposite_dict(dict_):
+    opposite = defaultdict(set)
+    for key, value in dict_.items():
+        opposite[value].add(key)
     return opposite

tests/system/test_set_trace.py

@@ -25,7 +25,7 @@ def test_set_trace(port, start_debugger_with_ctty):
     client_future = run_in_process(run_client, port, b'value_to_change += 1\nc\n')
     assert debugger_future.result(JOIN_TIMEOUT)
     client_output = client_future.result(JOIN_TIMEOUT)
-    # TODO: why does this assert fail? Problem in piping?
+    # TODO: waiting for fix in run_client
     # assert b'Closing connection' in client_output