Merge pull request #131 from doronz88/feature/socket_timeout

Feature/socket timeout

src/rpcclient/rpcclient/network.py

@@ -2,9 +2,11 @@
 import typing
 from collections import namedtuple
 
-from rpcclient.exceptions import BadReturnValueError
 from rpcclient.allocated import Allocated
-from rpcclient.structs.consts import AF_UNIX, AF_INET, SOCK_STREAM
+from rpcclient.darwin.structs import timeval
+from rpcclient.exceptions import BadReturnValueError
+from rpcclient.structs.consts import AF_UNIX, AF_INET, SOCK_STREAM, SOL_SOCKET, SO_RCVTIMEO, SO_SNDTIMEO, MSG_NOSIGNAL, \
+    EPIPE, EAGAIN, F_GETFL, O_NONBLOCK, F_SETFL
 from rpcclient.structs.generic import sockaddr_in, sockaddr_un, ifaddrs, sockaddr, hostent
 
 Interface = namedtuple('Interface', 'name address netmask broadcast')
@@ -22,6 +24,7 @@ def __init__(self, client, fd: int):
         super().__init__()
         self._client = client
         self.fd = fd
+        self._blocking = self._getblocking()
 
     def _deallocate(self):
         """ close(fd) at remote. read man for more details. """
@@ -39,8 +42,12 @@ def send(self, buf: bytes, size: int = None) -> int:
         """
         if size is None:
             size = len(buf)
-        n = self._client.symbols.send(self.fd, buf, size, 0).c_int64
+        n = self._client.symbols.send(self.fd, buf, size, MSG_NOSIGNAL).c_int64
         if n < 0:
+            if self._client.errno == EPIPE:
+                self.deallocate()
+            elif self._client.errno == EAGAIN:
+                raise pysock.timeout()
             raise BadReturnValueError(f'failed to send on fd: {self.fd}')
         return n
 
@@ -51,24 +58,60 @@ def sendall(self, buf: bytes):
             buf = buf[err:]
 
     def recv(self, size: int = CHUNK_SIZE) -> bytes:
-        """ recv(fd, buf, size, 0) at remote. read man for more details. """
+        """
+        recv(fd, buf, size, 0) at remote. read man for more details.
+
+        :param size: chunk size
+        :return: received bytes
+        """
         with self._client.safe_malloc(size) as chunk:
             err = self._client.symbols.recv(self.fd, chunk, size).c_int64
             if err < 0:
-                raise BadReturnValueError(f'read failed for fd: {self.fd}')
+                if self._client.errno == EAGAIN:
+                    raise TimeoutError()
+                raise BadReturnValueError(f'recv() failed for fd: {self.fd} ({self._client.last_error})')
             return chunk.peek(err)
 
     def recvall(self, size: int) -> bytes:
         """ recv at remote until all buffer is received """
         buf = b''
         with self._client.safe_malloc(size) as chunk:
             while len(buf) < size:
-                err = self._client.symbols.read(self.fd, chunk, size).c_int64
-                if err <= 0:
-                    raise BadReturnValueError(f'read failed for fd: {self.fd}')
+                err = self._client.symbols.recv(self.fd, chunk, size).c_int64
+                if err < 0:
+                    if self._client.errno == EAGAIN:
+                        raise TimeoutError()
+                    raise BadReturnValueError(f'recv() failed for fd: {self.fd} ({self._client.last_error})')
                 buf += chunk.peek(err)
         return buf
 
+    def setsockopt(self, level: int, option_name: int, option_value: bytes):
+        with self._client.safe_malloc(len(option_value)) as option:
+            option.poke(option_value)
+            if 0 != self._client.symbols.setsockopt(self.fd, level, option_name, option, len(option_value)):
+                raise BadReturnValueError(f'setsockopt() failed: {self._client.last_error}')
+
+    def settimeout(self, seconds: int):
+        self.setsockopt(SOL_SOCKET, SO_RCVTIMEO, timeval.build({'tv_sec': seconds, 'tv_usec': 0}))
+        self.setsockopt(SOL_SOCKET, SO_SNDTIMEO, timeval.build({'tv_sec': seconds, 'tv_usec': 0}))
+        self.setblocking(seconds == 0)
+
+    def setblocking(self, blocking: bool):
+        opts = self._client.symbols.fcntl(self.fd, F_GETFL, 0).c_uint64
+        if blocking:
+            opts &= ~O_NONBLOCK
+        else:
+            opts |= ~O_NONBLOCK
+        if 0 != self._client.symbols.fcntl(self.fd, F_SETFL, opts):
+            raise BadReturnValueError(f'fcntl() failed: {self._client.last_error}')
+        self._blocking = blocking
+
+    def getblocking(self) -> bool:
+        return self._blocking
+
+    def _getblocking(self) -> bool:
+        return not bool(self._client.symbols.fcntl(self.fd, F_GETFL, 0) & O_NONBLOCK)
+
     def __repr__(self):
         return f'<{self.__class__.__name__} FD:{self.fd}>'
 

src/rpcclient/rpcclient/structs/consts.py

@@ -22,6 +22,18 @@
 S_ISGID = 0o0002000
 S_ISVTX = 0o0001000
 
+SOL_IP = 0
+SOL_SOCKET = 65535
+SOL_TCP = 6
+SOL_UDP = 17
+
+SO_SNDLOWAT = 0x1003
+SO_RCVLOWAT = 0x1004
+SO_SNDTIMEO = 0x1005
+SO_RCVTIMEO = 0x1006
+
+MSG_NOSIGNAL = 524288
+
 AF_UNIX = 1
 AF_INET = 2
 AF_INET6 = 30
@@ -72,3 +84,11 @@
 SIGWINCH = 28
 SIGXCPU = 24
 SIGXFSZ = 25
+
+EPIPE = 32
+EAGAIN = 35
+
+F_SETFL = 4
+F_GETFL = 3
+
+O_NONBLOCK = 4

src/rpcclient/setup.py

@@ -3,7 +3,7 @@
 from setuptools import setup, find_packages
 
 BASE_DIR = Path(__file__).parent.resolve(strict=True)
-VERSION = '2.10.0'
+VERSION = '2.10.1'
 PACKAGE_NAME = 'rpcclient'
 DATA_FILES_EXTENSIONS = ['*.txt', '*.json', '*.js']
 PACKAGES = [p for p in find_packages() if not p.startswith('tests')]

src/rpcserver/common.c

@@ -115,7 +115,7 @@ bool sendall(int sockfd, const char *buf, size_t len)
 
     while (len > 0)
     {
-        bytes = send(sockfd, buf + total_bytes, len, 0);
+        bytes = send(sockfd, buf + total_bytes, len, MSG_NOSIGNAL);
         CHECK(bytes != -1);
 
         total_bytes += bytes;

src/rpcserver/rpcserver.c

@@ -1034,6 +1034,11 @@ void handle_client(int sockfd)
     close(sockfd);
 }
 
+void signal_handler(int sig)
+{
+    TRACE("entered with signal code: %d", sig);
+}
+
 int main(int argc, const char *argv[])
 {
     int opt;
@@ -1079,6 +1084,8 @@ int main(int argc, const char *argv[])
         }
     }
 
+    signal(SIGPIPE, signal_handler);
+
     int err = 0;
     int server_fd = -1;
     struct addrinfo hints;