rpc: improve connection check

src/rpcclient/rpcclient/client.py

@@ -14,7 +14,7 @@
 from traitlets.config import Config
 
 from rpcclient.darwin.structs import pid_t, exitcode_t
-from rpcclient.exceptions import ArgumentError, SymbolAbsentError, SpawnError
+from rpcclient.exceptions import ArgumentError, SymbolAbsentError, SpawnError, ServerDiedError
 from rpcclient.fs import Fs
 from rpcclient.network import Network
 from rpcclient.processes import Processes
@@ -366,6 +366,9 @@ def _recvall(self, size: int) -> bytes:
                 chunk = self._sock.recv(size)
             except BlockingIOError:
                 continue
+            if self._sock.gettimeout() == 0 and not chunk:
+                # TODO: replace self._sock.gettimeout() == 0 on -> self._sock.getblocking() on python37+
+                raise ServerDiedError()
             size -= len(chunk)
             buf += chunk
         return buf

src/rpcclient/rpcclient/client_factory.py

@@ -5,9 +5,9 @@
 from rpcclient.darwin.client import DarwinClient
 from rpcclient.ios.client import IosClient
 from rpcclient.linux.client import LinuxClient
-from rpcclient.exceptions import FailedToConnectError
+from rpcclient.exceptions import FailedToConnectError, InvalidServerVersionMagic
 from rpcclient.macos.client import MacosClient
-from rpcclient.protocol import UNAME_VERSION_LEN, DEFAULT_PORT
+from rpcclient.protocol import UNAME_VERSION_LEN, DEFAULT_PORT, SERVER_MAGIC_VERSION
 
 
 def recvall(sock, size: int) -> bytes:
@@ -29,6 +29,10 @@ def create_client(hostname: str, port: int = DEFAULT_PORT):
         # wrap in our own exception
         raise FailedToConnectError() from e
 
+    magic = recvall(sock, len(SERVER_MAGIC_VERSION))
+    if magic != SERVER_MAGIC_VERSION:
+        raise InvalidServerVersionMagic(f'got an invalid server magic: {magic.hex()}')
+
     sysname = recvall(sock, UNAME_VERSION_LEN).split(b'\x00', 1)[0].decode().lower()
     logging.info(f'connection uname.sysname: {sysname}')
 

src/rpcclient/rpcclient/exceptions.py

@@ -2,6 +2,14 @@ class RpcClientException(Exception):
     pass
 
 
+class InvalidServerVersionMagic(RpcClientException):
+    pass
+
+
+class ServerDiedError(RpcClientException):
+    pass
+
+
 class SymbolAbsentError(RpcClientException):
     pass
 

src/rpcclient/rpcclient/protocol.py

@@ -14,6 +14,7 @@
                   CMD_GET_DUMMY_BLOCK=9,
                   )
 DEFAULT_PORT = 5910
+SERVER_MAGIC_VERSION = Int32ul.build(0x88888800)
 MAGIC = 0x12345678
 MAX_PATH_LEN = 1024
 UNAME_VERSION_LEN = 256

src/rpcserver/rpcserver.c

@@ -34,6 +34,7 @@
 #define DEFAULT_PORT ("5910")
 #define DEFAULT_SHELL ("/bin/sh")
 #define USAGE ("Usage: %s [-p port] [-s shell]")
+#define SERVER_MAGIC_VERSION (0x88888800)
 #define MAGIC (0x12345678)
 #define MAX_CONNECTIONS (1024)
 
@@ -565,6 +566,10 @@ void handle_client(int sockfd)
 {
     TRACE("enter. fd: %d", sockfd);
 
+    // send MAGIC
+    u32 magic = SERVER_MAGIC_VERSION;
+    CHECK(sendall(sockfd, (const char *)&magic, sizeof(magic)));
+
     // notify client of the connected target os version
     struct utsname uname_buf;
     CHECK(0 == uname(&uname_buf));