adbproxy 0.1.0__py3-none-any.whl

adbproxy/device/connection.py

@@ -0,0 +1,505 @@
+import socket
+import threading
+
+from adbproxy.observers.stream import StreamBuffer
+from adbproxy.protocols.adb import AdbProtocol
+from adbproxy.protocols.transport import (
+    A_AUTH,
+    A_CLSE,
+    A_CNXN,
+    A_OKAY,
+    A_OPEN,
+    A_STLS,
+    A_WRTE,
+    MAX_SEND_QUEUE,
+    MAX_PAYLOAD_V1,
+    TransportCodec,
+    TransportProtocolError,
+)
+from adbproxy.runtime.logging import Logger
+
+from .backend import AdbClientBackend
+from .session import (
+    AcceptOpen,
+    AttachObserver,
+    BackendFailed,
+    BackendReady,
+    CloseStream,
+    ForwardToBackend,
+    Ignore,
+    LogEvent,
+    MarkTransportTerminal,
+    NegotiateCodec,
+    OpenBackend,
+    QueueFromBackend,
+    RejectOpen,
+    ResetStreamsForCnxn,
+    SendPacket,
+    SetCanSend,
+    AdbdSession,
+)
+from .stream_table import StreamEntry, StreamTable, _StreamEntry
+
+# re-export for tests: from adbproxy.device.connection import _StreamEntry
+__all__ = ["TransportConnection", "_StreamEntry", "StreamEntry"]
+
+
+_CONN_ID_LOCK = threading.Lock()
+_CONN_ID_SEQ = 0
+
+
+def _next_conn_id():
+    global _CONN_ID_SEQ
+    with _CONN_ID_LOCK:
+        _CONN_ID_SEQ += 1
+        return _CONN_ID_SEQ
+
+
+class TransportConnection:
+    """per-conn front 传输状态机：CNXN 握手 + OPEN/OKAY/WRTE/CLSE 流表 + id 重映射
+    + v1 停等流控 + CLSE 幂等 + service 派发。
+
+    Decision layer: ``AdbdSession`` (Track 2b). This class executes actions and
+    owns I/O / threads / ``_try_flush`` (algorithm unchanged).
+    """
+
+    def __init__(
+        self,
+        conn_id,
+        front_sock,
+        file_logger,
+        capture_enabled,
+        debug_logger,
+        serial,
+        adb_server_port,
+        banner_features,
+        banner_type,
+        banner_product,
+        banner_model,
+        banner_device,
+    ):
+        self.conn_id = conn_id
+        self.front = front_sock
+        self.logger = Logger(conn_id, file_logger)
+        self.file_logger = file_logger
+        self.capture_enabled = capture_enabled
+        self.codec = TransportCodec()
+        self.debug_logger = debug_logger
+        self.serial = serial
+        self.adb_server_port = adb_server_port
+        self.banner_features = banner_features
+        self.banner_type = banner_type
+        self.banner_product = banner_product
+        self.banner_model = banner_model
+        self.banner_device = banner_device
+        self.stream_table = StreamTable()
+        # Compat views used by tests / existing methods:
+        self.streams = self.stream_table.streams
+        self._streams_lock = self.stream_table.lock
+        self.session = AdbdSession(
+            self.stream_table,
+            banner_features=banner_features,
+            banner_type=banner_type,
+            banner_product=banner_product,
+            banner_model=banner_model,
+            banner_device=banner_device,
+            peer_maxdata=MAX_PAYLOAD_V1,
+        )
+        self._front_send_lock = threading.Lock()
+        self.peer_maxdata = MAX_PAYLOAD_V1
+        self._finished = threading.Event()
+        self._finish_lock = threading.Lock()
+        self._backend_threads = set()
+        self._backend_threads_lock = threading.Lock()
+        self._stop = False
+
+    @property
+    def _cnxn_done(self):
+        return self.session.cnxn_done
+
+    @_cnxn_done.setter
+    def _cnxn_done(self, value):
+        self.session._cnxn_done = bool(value)
+
+    def run(self):
+        try:
+            while not self._stop:
+                try:
+                    data = self.front.recv(65536)
+                except OSError:
+                    break
+                if not data:
+                    break
+                if self.debug_logger is not None:
+                    try:
+                        self.debug_logger.write_hex(self.conn_id, "c", data)
+                    except Exception:
+                        pass
+                for pkt in self.codec.feed(data):
+                    self._dispatch(pkt)
+        except TransportProtocolError as exc:
+            self.logger.event(f"[!] transport 协议错误，终止连接: {exc}")
+        except OSError:
+            pass
+        finally:
+            self._finish()
+
+    def _dispatch(self, apacket):
+        actions = self.session.receive(apacket)
+        self._execute_actions(actions)
+
+    def _execute_actions(self, actions):
+        for action in actions:
+            self._execute_action(action)
+
+    def _execute_action(self, action):
+        if isinstance(action, Ignore):
+            return
+        if isinstance(action, LogEvent):
+            self.logger.event(action.message)
+            return
+        if isinstance(action, MarkTransportTerminal):
+            raise TransportProtocolError(action.reason)
+        if isinstance(action, ResetStreamsForCnxn):
+            for entry in self.stream_table.snapshot():
+                self._close_entry(entry, send_clse=False)
+            return
+        if isinstance(action, NegotiateCodec):
+            self.peer_maxdata = action.peer_maxdata
+            self.session.peer_maxdata = action.peer_maxdata
+            self.codec.negotiate(action.version, action.peer_maxdata)
+            return
+        if isinstance(action, SendPacket):
+            try:
+                self._front_send(
+                    self.codec.encode(
+                        action.cmd, action.arg0, action.arg1, action.payload
+                    )
+                )
+            except Exception:
+                if action.cmd == A_OKAY:
+                    # backend ready path: fail closed without front OKAY already sent
+                    pass
+                raise
+            return
+        if isinstance(action, RejectOpen):
+            self._front_send(self.codec.encode(A_CLSE, 0, action.peer_local))
+            return
+        if isinstance(action, AcceptOpen):
+            # allocation already done in session; nothing extra
+            return
+        if isinstance(action, AttachObserver):
+            entry = action.entry
+            service = action.service
+            try:
+                entry.sb = StreamBuffer(
+                    _next_conn_id(), self.file_logger, self.capture_enabled
+                )
+                entry.sb.logger.event(f"📥 [User Command]: {service}")
+                entry.sb.shell.set_mode_from_service(service)
+                if AdbProtocol.is_sync_service(service):
+                    entry.sb.sync.activate()
+                if AdbProtocol.is_install_stream_service(service):
+                    size = AdbProtocol.parse_install_size(service)
+                    if size is not None:
+                        entry.sb.install.activate(size, service)
+                        cap = (
+                            entry.sb.install.capture_path
+                            if (
+                                entry.sb.install.capture_enabled
+                                and entry.sb.install.capture_path
+                            )
+                            else "<disabled>"
+                        )
+                        entry.sb.logger.event(
+                            f"📤 [install push]: remote={service} size={size} local_capture={cap}"
+                        )
+            except Exception:
+                if entry.sb is not None:
+                    try:
+                        entry.sb.close()
+                    except Exception:
+                        pass
+                entry.sb = None
+            return
+        if isinstance(action, OpenBackend):
+            entry = action.entry
+            observer_logger = entry.sb.logger if entry.sb is not None else self.logger
+            back = AdbClientBackend(
+                self, entry, self.serial, self.adb_server_port, observer_logger
+            )
+            entry.back = back
+            worker = threading.Thread(
+                target=self._backend_worker,
+                args=(entry, back, action.generation),
+                daemon=True,
+            )
+            with self._backend_threads_lock:
+                self._backend_threads.add(worker)
+            worker.start()
+            return
+        if isinstance(action, SetCanSend):
+            entry = action.entry
+            with entry.lock:
+                entry.can_send = action.value
+            if action.flush:
+                self._try_flush(entry)
+            return
+        if isinstance(action, ForwardToBackend):
+            entry = action.entry
+            if action.also_okay:
+                self._front_send(
+                    self.codec.encode(A_OKAY, entry.local_id, entry.remote_id)
+                )
+            if entry.sb is not None:
+                try:
+                    entry.sb.add_client(action.payload)
+                except Exception:
+                    pass
+            if entry.back is not None and entry.back.sock is not None:
+                try:
+                    entry.back.sock.sendall(action.payload)
+                except Exception as e:
+                    self.logger.event(f"[!] back sendall(wrte) 失败: {e}")
+                    self._front_clse(entry)
+            return
+        if isinstance(action, CloseStream):
+            self._close_entry(
+                action.entry,
+                send_clse=action.send_clse,
+                close_arg0=action.close_arg0,
+                close_arg1=action.close_arg1,
+            )
+            return
+        if isinstance(action, QueueFromBackend):
+            self._execute_queue_from_backend(action.entry, action.data)
+            return
+        self.logger.event(f"[!] unknown TransportAction {type(action)!r}")
+
+    def _backend_worker(self, entry, back, generation):
+        """在独立 worker 中完成后端握手，front dispatch 不等待网络 I/O。"""
+        try:
+            if not back.start():
+                self._backend_failed(entry, generation)
+                return
+            if not self._backend_ready(entry, generation):
+                back.stop()
+                return
+            back.run()
+        finally:
+            current = threading.current_thread()
+            with self._backend_threads_lock:
+                self._backend_threads.discard(current)
+
+    def _backend_ready(self, entry, generation):
+        actions = self.session.on_session_event(BackendReady(entry, generation))
+        # Special-case: if only Ignore, return False (stale generation)
+        if len(actions) == 1 and isinstance(actions[0], Ignore):
+            return False
+        try:
+            # Execute without double SetCanSend flush before OKAY send order:
+            # plan: SendPacket(OKAY) then flush. Session emits SendPacket + SetCanSend.
+            for action in actions:
+                if isinstance(action, SendPacket):
+                    try:
+                        self._front_send(
+                            self.codec.encode(
+                                action.cmd, action.arg0, action.arg1, action.payload
+                            )
+                        )
+                    except Exception:
+                        self._close_entry(entry, send_clse=False)
+                        return False
+                elif isinstance(action, SetCanSend):
+                    # can_send already set under lock in session; just flush
+                    if action.flush:
+                        self._try_flush(action.entry)
+                elif isinstance(action, Ignore):
+                    continue
+                else:
+                    self._execute_action(action)
+        except Exception:
+            self._close_entry(entry, send_clse=False)
+            return False
+        return True
+
+    def _backend_failed(self, entry, generation):
+        actions = self.session.on_session_event(BackendFailed(entry, generation))
+        self._execute_actions(actions)
+
+    # Compat wrappers for tests that call handlers directly
+    def _handle_cnxn(self, arg0, arg1, payload):
+        self._execute_actions(self.session.receive((A_CNXN, arg0, arg1, payload)))
+
+    def _handle_open(self, arg0, arg1, payload):
+        self._execute_actions(self.session.receive((A_OPEN, arg0, arg1, payload)))
+
+    def _handle_okay(self, arg0, arg1, payload):
+        self._execute_actions(self.session.receive((A_OKAY, arg0, arg1, payload)))
+
+    def _handle_wrte(self, arg0, arg1, payload):
+        self._execute_actions(self.session.receive((A_WRTE, arg0, arg1, payload)))
+
+    def _handle_clse(self, arg0, arg1, payload):
+        self._execute_actions(self.session.receive((A_CLSE, arg0, arg1, payload)))
+
+    def _handle_stls(self, arg0, arg1, payload):
+        self._execute_actions(self.session.receive((A_STLS, arg0, arg1, payload)))
+
+    def _handle_auth(self, arg0, arg1, payload):
+        self._execute_actions(self.session.receive((A_AUTH, arg0, arg1, payload)))
+
+    def _front_wrte_from_back(self, entry, data):
+        """back 收到设备流字节：经 QueueFromBackend 专用执行路径入队（旁路观察 + 切块 + 背压）。"""
+        self._execute_action(QueueFromBackend(entry=entry, data=data))
+
+    def _execute_queue_from_backend(self, entry, data):
+        """QueueFromBackend 专用路径。字面序与历史 _front_wrte_from_back 一致；不改 _try_flush。"""
+        if not data:
+            return
+        if entry.sb is not None:
+            try:
+                entry.sb.add_server(data)
+            except Exception:
+                pass
+        maxdata = self.codec.max_payload
+        with entry.lock:
+            if entry.closed:
+                return
+            pos = 0
+            dlen = len(data)
+            while pos < dlen:
+                entry.send_queue.append(data[pos : pos + maxdata])
+                pos += maxdata
+            entry.queued_bytes += dlen
+            while entry.queued_bytes > MAX_SEND_QUEUE and not entry.closed:
+                entry.flow_cv.wait()
+            if entry.closed:
+                return
+        self._try_flush(entry)
+
+    def _try_flush(self, entry):
+        """v1 停等：can_send 且队列非空时取一条发 WRTE，置 can_send=False；
+        排出 chunk 后唤醒背压等待的 back 线程；后端 EOF 时等队列和在途 WRTE 排空再 CLSE。"""
+        chunk = None
+        should_close = False
+        with entry.lock:
+            if entry.closed or not entry.can_send:
+                return
+            if entry.send_queue:
+                chunk = entry.send_queue.popleft()
+                entry.queued_bytes -= len(chunk)
+                entry.can_send = False
+                entry.flow_cv.notify_all()
+            elif entry.back_eof:
+                should_close = True
+        if should_close:
+            self._front_clse(entry)
+            return
+        if chunk is None:
+            return
+        try:
+            self._front_send(
+                self.codec.encode(A_WRTE, entry.local_id, entry.remote_id, chunk)
+            )
+        except Exception as e:
+            self.logger.event(f"[!] front sendall(wrte) 失败: {e}")
+            self._front_clse(entry)
+
+    def _back_eof(self, entry):
+        """记录后端 EOF；若还有在途或排队 WRTE，延迟到最后一次 OKAY 后再关闭。"""
+        with entry.lock:
+            if entry.closed:
+                return
+            entry.back_eof = True
+        self._try_flush(entry)
+
+    def _front_clse(self, entry):
+        if entry is None:
+            return
+        self._close_entry(entry)
+
+    def _close_entry(self, entry, send_clse=True, close_arg0=None, close_arg1=None):
+        """一次性完成 stream 状态、等待者、后端、观察器和 registry 的收尾。"""
+        with entry.lock:
+            if entry.closed:
+                return False
+            entry.closed = True
+            entry.state = "closed"
+            entry.flow_cv.notify_all()
+        if send_clse:
+            arg0 = entry.local_id if close_arg0 is None else close_arg0
+            arg1 = entry.remote_id if close_arg1 is None else close_arg1
+            try:
+                self._front_send(self.codec.encode(A_CLSE, arg0, arg1))
+            except Exception:
+                pass
+        if entry.back is not None:
+            try:
+                entry.back.stop()
+            except Exception:
+                pass
+        if entry.sb is not None:
+            try:
+                entry.sb.close()
+            except Exception:
+                pass
+        self.stream_table.remove_if_same(entry)
+        return True
+
+    def _find_by_local(self, local_id):
+        return self.stream_table.find_by_local(local_id)
+
+    def _find_by_remote(self, remote_id):
+        return self.stream_table.find_by_remote(remote_id)
+
+    def _find_by_pair(self, local_id, remote_id):
+        return self.stream_table.find_by_pair(local_id, remote_id)
+
+    def _front_send(self, raw):
+        """front sendall 串行化 + --debug per-apacket 落盘（outgoing = device→program，标 "s"）。"""
+        with self._front_send_lock:
+            self.front.sendall(raw)
+        if self.debug_logger is not None:
+            try:
+                self.debug_logger.write_hex(self.conn_id, "s", raw)
+            except Exception:
+                pass
+
+    def stop(self):
+        self._stop = True
+        try:
+            self.front.shutdown(socket.SHUT_RDWR)
+        except Exception:
+            pass
+        self._finish()
+
+    def wait(self, timeout=None):
+        done = self._finished.wait(timeout)
+        if not done:
+            return False
+        current = threading.current_thread()
+        with self._backend_threads_lock:
+            workers = list(self._backend_threads)
+        for worker in workers:
+            if worker is not current:
+                worker.join(timeout)
+        return True
+
+    def _finish(self):
+        with self._finish_lock:
+            if self._finished.is_set():
+                return
+            self._stop = True
+            entries = self.stream_table.snapshot()
+            for e in entries:
+                self._close_entry(e, send_clse=False)
+            try:
+                self.front.shutdown(socket.SHUT_RDWR)
+            except Exception:
+                pass
+            try:
+                self.front.close()
+            except Exception:
+                pass
+            self._finished.set()