network: support setsockopt() and settimeout()

src/rpcclient/rpcclient/network.py

@@ -3,9 +3,10 @@
 from collections import namedtuple
 
 from rpcclient.allocated import Allocated
+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
+    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')
@@ -23,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. """
@@ -44,6 +46,8 @@ def send(self, buf: bytes, size: int = None) -> int:
         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
 
@@ -54,11 +58,18 @@ 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:
@@ -67,11 +78,40 @@ def recvall(self, size: int) -> bytes:
         with self._client.safe_malloc(size) as chunk:
             while len(buf) < size:
                 err = self._client.symbols.recv(self.fd, chunk, size).c_int64
-                if err <= 0:
-                    raise BadReturnValueError(f'read failed for fd: {self.fd}')
+                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

@@ -86,3 +86,9 @@
 SIGXFSZ = 25
 
 EPIPE = 32
+EAGAIN = 35
+
+F_SETFL = 4
+F_GETFL = 3
+
+O_NONBLOCK = 4