diff --git a/src/rpcclient/rpcclient/network.py b/src/rpcclient/rpcclient/network.py index ce538033..afc7b158 100644 --- a/src/rpcclient/rpcclient/network.py +++ b/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,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: @@ -63,12 +77,41 @@ def recvall(self, size: int) -> bytes: 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}>' diff --git a/src/rpcclient/rpcclient/structs/consts.py b/src/rpcclient/rpcclient/structs/consts.py index 8fbd4138..a7229ecc 100644 --- a/src/rpcclient/rpcclient/structs/consts.py +++ b/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 diff --git a/src/rpcclient/setup.py b/src/rpcclient/setup.py index dbd8adcb..880d3963 100644 --- a/src/rpcclient/setup.py +++ b/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')] diff --git a/src/rpcserver/common.c b/src/rpcserver/common.c index 23bc3483..c34362ef 100644 --- a/src/rpcserver/common.c +++ b/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; diff --git a/src/rpcserver/rpcserver.c b/src/rpcserver/rpcserver.c index 675db068..1defe1b0 100644 --- a/src/rpcserver/rpcserver.c +++ b/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;