async def handle_stream(self, stream: IOStream, address: Tuple[str, str]) -> None: """ :param stream: Client IOStream to read telemetry data from :type stream: IOStream :param address: The IP address and port on which a client connects to the server :type address: Tuple[str,str] :return: None """ try: self.log.info(f"Got Connection from {address[0]}:{address[1]}") while not stream.closed(): header_data: bytes = await stream.read_bytes(self._header_size) ( msg_type, encode_type, msg_version, flags, msg_length, ) = self._header_struct.unpack(header_data) # encoding = {1: "gpb", 2: "json"}[encode_type] # implement json encoding msg_data: bytes = b"" while len(msg_data) < msg_length: packet: bytes = await stream.read_bytes(msg_length - len(msg_data)) msg_data += packet self.data_queue.put_nowait( ("ems", msg_data, None, None, address[0])) except StreamClosedError as error: self.log.error(f'{address[0]}:{address[1]} {error}') stream.close()
async def _pipe_stream(self, _in: IOStream, out: IOStream): while not _in.closed(): try: data = await _in.read_bytes(10240, partial=True) await out.write(data) except StreamClosedError as e: _in.close() # here may call twice out.close()
def _maybe_connect(self, to_pid, callback=None): """Asynchronously establish a connection to the remote pid.""" callback = stack_context.wrap(callback or (lambda stream: None)) def streaming_callback(data): # we are not guaranteed to get an acknowledgment, but log and discard bytes if we do. log.info('Received %d bytes from %s, discarding.' % (len(data), to_pid)) log.debug(' data: %r' % (data, )) def on_connect(exit_cb, stream): log.info('Connection to %s established' % to_pid) with self._connection_callbacks_lock: self._connections[to_pid] = stream self.__dispatch_on_connect_callbacks(to_pid, stream) self.__loop.add_callback(stream.read_until_close, exit_cb, streaming_callback=streaming_callback) create = False with self._connection_callbacks_lock: stream = self._connections.get(to_pid) callbacks = self._connection_callbacks.get(to_pid) if not stream: self._connection_callbacks[to_pid].append(callback) if not callbacks: create = True if stream: self.__loop.add_callback(callback, stream) return if not create: return sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM, 0) if not sock: raise self.SocketError('Failed opening socket') stream = IOStream(sock, io_loop=self.__loop) stream.set_nodelay(True) stream.set_close_callback( partial(self.__on_exit, to_pid, b'reached end of stream')) connect_callback = partial(on_connect, partial(self.__on_exit, to_pid), stream) log.info('Establishing connection to %s' % to_pid) stream.connect((to_pid.ip, to_pid.port), callback=connect_callback) if stream.closed(): raise self.SocketError('Failed to initiate stream connection') log.info('Maybe connected to %s' % to_pid)
def _maybe_connect(self, to_pid, callback=None): """Asynchronously establish a connection to the remote pid.""" callback = stack_context.wrap(callback or (lambda stream: None)) def streaming_callback(data): # we are not guaranteed to get an acknowledgment, but log and discard bytes if we do. log.info('Received %d bytes from %s, discarding.' % (len(data), to_pid)) log.debug(' data: %r' % (data,)) def on_connect(exit_cb, stream): log.info('Connection to %s established' % to_pid) with self._connection_callbacks_lock: self._connections[to_pid] = stream self.__dispatch_on_connect_callbacks(to_pid, stream) self.__loop.add_callback( stream.read_until_close, exit_cb, streaming_callback=streaming_callback) create = False with self._connection_callbacks_lock: stream = self._connections.get(to_pid) callbacks = self._connection_callbacks.get(to_pid) if not stream: self._connection_callbacks[to_pid].append(callback) if not callbacks: create = True if stream: self.__loop.add_callback(callback, stream) return if not create: return sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM, 0) if not sock: raise self.SocketError('Failed opening socket') stream = IOStream(sock, io_loop=self.__loop) stream.set_nodelay(True) stream.set_close_callback(partial(self.__on_exit, to_pid, b'reached end of stream')) connect_callback = partial(on_connect, partial(self.__on_exit, to_pid), stream) log.info('Establishing connection to %s' % to_pid) stream.connect((to_pid.ip, to_pid.port), callback=connect_callback) if stream.closed(): raise self.SocketError('Failed to initiate stream connection') log.info('Maybe connected to %s' % to_pid)
async def _read_stream(self, stream: iostream.IOStream) -> None: while not stream.closed(): try: data = await stream.read_until(b'\x03') except iostream.StreamClosedError as e: return except Exception: logging.exception("Klippy Stream Read Error") continue try: decoded_cmd = json.loads(data[:-1]) self.on_recd(decoded_cmd) except Exception: logging.exception( f"Error processing Klippy Host Response: {data.decode()}")
class ForwardConnection(object): def __init__(self, server, stream, address): self._close_callback = None self.server = server self.stream = stream self.reverse_address = address self.address = stream.socket.getsockname() self.remote_address = server.conf[self.address] sock = socket.socket() self.remote_stream = IOStream(sock) self.remote_stream.connect(self.remote_address, self._on_remote_connected) def close(self): self.remote_stream.close() def set_close_callback(self, callback): self._close_callback = callback def _on_remote_connected(self): ip_from = self.reverse_address[0] fwd_str = get_forwarding_str(self.address[0], self.address[1], self.remote_address[0], self.remote_address[1]) logging.info('Connected ip: %s, forward %s', ip_from, fwd_str) self.remote_stream.read_until_close(self._on_remote_read_close, self.stream.write) self.stream.read_until_close(self._on_read_close, self.remote_stream.write) def _on_remote_read_close(self, data): if self.stream.writing(): self.stream.write(data, self.stream.close) else: if self.stream.closed(): self._on_closed() else: self.stream.close() def _on_read_close(self, data): if self.remote_stream.writing(): self.remote_stream.write(data, self.remote_stream.close) else: if self.remote_stream.closed(): self._on_closed() else: self.remote_stream.close() def _on_closed(self): logging.info('Disconnected ip: %s', self.reverse_address[0]) if self._close_callback: self._close_callback(self)
async def _pipe_stream(self, _in: IOStream, out: IOStream): while not _in.closed(): try: data = await _in.read_bytes(10240, partial=True) if self.__debug: if self.__names.get(_in): print("{} >>> device".format(self.__names[_in])) print(colored.fg("green") + colored.bg("black"), end="") else: print("device >>> {}".format(self.__names[out])) print(colored.fg("orchid") + colored.bg("black"), end="") hexdump(data) print(colored.attr("reset"), end="") await out.write(data) except StreamClosedError as e: _in.close() # here may call twice out.close() self.__names.pop(_in, None)
def maybe_connect(self, to_pid, callback=None): """Synchronously open a connection to to_pid or return a connection if it exists.""" callback = stack_context.wrap(callback or (lambda stream: None)) def streaming_callback(data): # we are not guaranteed to get an acknowledgment, but log and discard bytes if we do. log.info('Received %d bytes from %s, discarding.' % (len(data), to_pid)) def on_connect(exit_cb, stream): log.info('Connection to %s established' % to_pid) self._connections[to_pid] = stream callback(stream) self.loop.add_callback(stream.read_until_close, exit_cb, streaming_callback=streaming_callback) stream = self._connections.get(to_pid) if stream is not None: callback(stream) return sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM, 0) if not sock: raise self.SocketError("Failed opening socket") # Set the socket non-blocking sock.setblocking(0) stream = IOStream(sock, io_loop=self.loop) stream.set_nodelay(True) stream.set_close_callback(partial(self.__on_exit, to_pid, b'closed from maybe_connect')) connect_callback = partial(on_connect, partial(self.__on_exit, to_pid), stream) log.info('Establishing connection to %s' % to_pid) stream.connect((to_pid.ip, to_pid.port), callback=connect_callback) if stream.closed(): raise self.SocketError("Failed to initiate stream connection") log.info('Maybe connected to %s' % to_pid)
class WS2IRCBridge(WebSocketHandler): def open(self, host='irc.freenode.net', port=None): port = int(port or 6667) self.sock = IOStream(socket.socket(socket.AF_INET, socket.SOCK_STREAM, 0)) self.sock.connect((host, port), self.sock_loop) logging.debug("Request received for %s:%d", host, port) def sock_loop(self, data=None): if data: self.write_message(data) if self.sock.closed(): self.close() logging.debug("IRC socket closed. Closing active WebSocket too.") else: self.sock.read_until("\r\n", self.sock_loop) def on_message(self, message): self.sock.write(message.encode('utf-8') + "\r\n") def on_close(self): self.sock.close() logging.debug("Client closed the WebSocket.")
class Trickle(object): """A coroutine-friendly :class:`~tornado.iostream.IOStream` interface. Takes same parameters as ``IOStream``, or takes a single ``IOStream`` as its only parameter. """ def __init__(self, *args, **kwargs): if args and isinstance(args[0], IOStream): if len(args) > 1 or kwargs: raise TypeError('Too many arguments to Trickle()') self.stream = args[0] else: self.stream = IOStream(*args, **kwargs) def connect(self, address, server_hostname=None, timeout=None): """Connects the socket to a remote address without blocking. Like ``IOStream`` :meth:`~tornado.iostream.IOStream.connect`, but returns a :class:`~tornado.concurrent.Future` and takes no callback. """ method = trickle_method('connect', timeout) return method(self, address, server_hostname=server_hostname) def read_until(self, delimiter, timeout=None): """Read up to the given delimiter.. Like ``IOStream`` :meth:`~tornado.iostream.IOStream.read_until`, but returns a :class:`~tornado.concurrent.Future` and takes no callback. """ return trickle_method('read_until', timeout)(self, delimiter) def read_until_regex(self, regex, timeout=None): """Read up to the given regex pattern. Like ``IOStream`` :meth:`~tornado.iostream.IOStream.read_until_regex`, but returns a :class:`~tornado.concurrent.Future` and takes no callback. """ return trickle_method('read_until_regex', timeout)(self, regex) def read_bytes(self, num_bytes, timeout=None): """Read the given number of bytes. Like ``IOStream`` :meth:`~tornado.iostream.IOStream.read_bytes`, but returns a :class:`~tornado.concurrent.Future` and takes no callback or streaming_callback. """ return trickle_method('read_bytes', timeout)(self, num_bytes) @gen.coroutine def read_until_close(self, timeout=None): """Read all remaining data from the socket. Like ``IOStream`` :meth:`~tornado.iostream.IOStream.read_until_close`, but returns a :class:`~tornado.concurrent.Future` and takes no callback or streaming_callback. """ stream = self.stream ioloop_timeout = None if timeout: def on_timeout(): stream.close((socket.timeout, socket.timeout(), None)) ioloop_timeout = stream.io_loop.add_timeout( timeout, callback=on_timeout) stream.read_until_close(callback=(yield gen.Callback(closed))) result = yield Wait(closed) if ioloop_timeout is not None: stream.io_loop.remove_timeout(ioloop_timeout) if stream.error: raise stream.error raise gen.Return(result) def write(self, data, timeout=None, **kwargs): """Write the given data to this stream. Like ``IOStream`` :meth:`~tornado.iostream.IOStream.write`, but returns a :class:`~tornado.concurrent.Future` and takes no callback. yield the returned Future to wait for all data to be written to the stream. """ return trickle_method('write', timeout)(self, data, **kwargs) def closed(self): """Returns true if the stream has been closed.""" return self.stream.closed()
class TornadoClient(Client): """A non-blocking Pomelo client by tornado ioloop Usage : class ClientHandler(object) : def on_recv_data(self, client, proto_type, data) : print "recv_data..." return data def on_connected(self, client, user_data) : print "connect..." client.send_heartbeat() def on_disconnect(self, client) : print "disconnect..." def on_heartbeat(self, client) : print "heartbeat..." send request ... def on_response(self, client, route, request, response) : print "response..." def on_push(self, client, route, push_data) : print "notify..." handler = ClientHandler() client = TornadoClient(handler) client.connect(host, int(port)) client.run() tornado.ioloop.IOLoop.current().start() """ def __init__(self, handler): self.socket = socket(AF_INET, SOCK_STREAM) self.iostream = None self.protocol_package = None super(TornadoClient, self).__init__(handler) def connect(self, host, port): self.iostream = IOStream(self.socket) self.iostream.set_close_callback(self.on_close) self.iostream.connect((host, port), self.on_connect) def on_connect(self): self.send_sync() self.on_data() def on_close(self): if hasattr(self.handler, 'on_disconnect'): self.handler.on_disconnect(self) def send(self, data): assert not self.iostream.closed(), "iostream has closed" if not isinstance(data, bytes): data = bytes(data) self.iostream.write(data) def on_data(self): assert not self.iostream.closed(), "iostream has closed" if None is self.protocol_package or self.protocol_package.completed(): self.iostream.read_bytes(4, self.on_head) def on_head(self, head): self.protocol_package = Protocol.unpack(head) self.iostream.read_bytes(self.protocol_package.length, self.on_body) def on_body(self, body): if hasattr(self.handler, 'on_recv_data'): body = self.handler.on_recv_data(self, self.protocol_package.proto_type, body) self.protocol_package.append(body) self.on_protocol(self.protocol_package) self.on_data() def close(self): if self.iostream: self.iostream.close()
class redisSub: def __init__(self, host="localhost", port=6379, db=0, pw=None): self.socket = socket.socket() self.socket.connect((host, port)) self.fp = self.socket.makefile("rb") self.stream = None if pw: self._command("AUTH", pw) self._get_reply() if db: self._command("SELECT", db) self._get_reply() def _command(self, *args): cmd = [b"*", b(len(args)), b"\r\n"] for val in args: val = b(val) cmd += [b"$", b(len(val)), b"\r\n", val, b"\r\n"] cmd = b"".join(cmd) self.socket.sendall(cmd) def _get_reply(self): data = self.fp.readline() flag, data = data[0], data[1:-2] if flag == 36: # b"$" bulk length = int(data) if length == -1: data = None else: data = self.fp.readline() data = data[:-2].decode() elif flag == 42: # b"*" multi bulk length = int(data) if length == -1: data = None else: data = tuple(self._get_reply() for i in range(length)) elif flag == 58: # b":" integer data = int(data) elif flag == 43: # b"+" status data = data.decode() elif flag == 45: # b"-" error raise Exception(data) return data def subscribe(self, *channels): self._command("SUBSCRIBE", *channels) return tuple(self._get_reply() for i in channels) def unsubscribe(self, *channels): self._command("UNSUBSCRIBE", *channels) return tuple(self._get_reply() for i in channels) def psubscribe(self, *pattenns): self._command("PSUBSCRIBE", *pattenns) return tuple(self._get_reply() for i in pattenns) def punsubscribe(self, *pattenns): self._command("PUNSUBSCRIBE", *pattenns) return tuple(self._get_reply() for i in pattenns) def close(self): if self.stream is not None: self.stream.close() else: self.fp.close() self.socket.close() @gen.coroutine def listen(self, callback): if self.stream is not None: return self.fp.close() self.stream = IOStream(self.socket) self.stream.set_nodelay(True) while not self.stream.closed(): data = yield gen.Task(self.stream.read_until, b"\n") flag, data = data[0], data[1:-2] if flag != 42: raise Exception("unexpected reply flag: " + chr(flag)) length = int(data) if length == -1: reply = None else: reply = [] for i in range(length): data = yield gen.Task(self.stream.read_until, b"\n") flag, data = data[0], data[1:-2] if flag != 36: raise Exception("unexpected reply flag: " + chr(flag)) length = int(data) if length == -1: reply.append(None) else: data = yield gen.Task(self.stream.read_until, b"\n") data = data[:-2].decode() reply.append(data) reply = tuple(reply) callback(reply)
class IPCClient(object): ''' A Tornado IPC client very similar to Tornado's TCPClient class but using either UNIX domain sockets or TCP sockets This was written because Tornado does not have its own IPC server/client implementation. :param IOLoop io_loop: A Tornado ioloop to handle scheduling :param str/int socket_path: A path on the filesystem where a socket belonging to a running IPCServer can be found. It may also be of type 'int', in which case it is used as the port for a tcp localhost connection. ''' # Create singleton map between two sockets instance_map = weakref.WeakKeyDictionary() def __new__(cls, socket_path, io_loop=None): io_loop = io_loop or tornado.ioloop.IOLoop.current() if io_loop not in IPCClient.instance_map: IPCClient.instance_map[io_loop] = weakref.WeakValueDictionary() loop_instance_map = IPCClient.instance_map[io_loop] # FIXME key = str(socket_path) if key not in loop_instance_map: log.debug('Initializing new IPCClient for path: {0}'.format(key)) new_client = object.__new__(cls) # FIXME new_client.__singleton_init__(io_loop=io_loop, socket_path=socket_path) loop_instance_map[key] = new_client else: log.debug('Re-using IPCClient for {0}'.format(key)) return loop_instance_map[key] def __singleton_init__(self, socket_path, io_loop=None): ''' Create a new IPC client IPC clients cannot bind to ports, but must connect to existing IPC servers. Clients can then send messages to the server. ''' self.io_loop = io_loop or tornado.ioloop.IOLoop.current() self.socket_path = socket_path self._closing = False self.stream = None if six.PY2: encoding = None else: encoding = 'utf-8' self.unpacker = msgpack.Unpacker(encoding=encoding) def __init__(self, socket_path, io_loop=None): # Handled by singleton __new__ pass def connected(self): return self.stream is not None and not self.stream.closed() def connect(self, callback=None, timeout=None): ''' Connect to the IPC socket ''' if hasattr(self, '_connecting_future') and not self._connecting_future.done(): # pylint: disable=E0203 future = self._connecting_future # pylint: disable=E0203 else: if hasattr(self, '_connecting_future'): # read previous future result to prevent the "unhandled future exception" error self._connecting_future.exc_info() # pylint: disable=E0203 future = tornado.concurrent.Future() self._connecting_future = future self._connect(timeout=timeout) if callback is not None: def handle_future(future): response = future.result() self.io_loop.add_callback(callback, response) future.add_done_callback(handle_future) return future @tornado.gen.coroutine def _connect(self, timeout=None): ''' Connect to a running IPCServer ''' if isinstance(self.socket_path, int): sock_type = socket.AF_INET sock_addr = ('127.0.0.1', self.socket_path) else: sock_type = socket.AF_UNIX sock_addr = self.socket_path self.stream = None if timeout is not None: timeout_at = time.time() + timeout while True: if self._closing: break if self.stream is None: self.stream = IOStream( socket.socket(sock_type, socket.SOCK_STREAM), io_loop=self.io_loop, ) try: log.trace('IPCClient: Connecting to socket: {0}'.format(self.socket_path)) yield self.stream.connect(sock_addr) self._connecting_future.set_result(True) break except Exception as e: if self.stream.closed(): self.stream = None if timeout is None or time.time() > timeout_at: if self.stream is not None: self.stream.close() self.stream = None self._connecting_future.set_exception(e) break yield tornado.gen.sleep(1) def __del__(self): self.close() def close(self): ''' Routines to handle any cleanup before the instance shuts down. Sockets and filehandles should be closed explicitly, to prevent leaks. ''' if self._closing: return self._closing = True if self.stream is not None and not self.stream.closed(): self.stream.close() # Remove the entry from the instance map so # that a closed entry may not be reused. # This forces this operation even if the reference # count of the entry has not yet gone to zero. if self.io_loop in IPCClient.instance_map: loop_instance_map = IPCClient.instance_map[self.io_loop] key = str(self.socket_path) if key in loop_instance_map: del loop_instance_map[key]
class WebSocketProxyHandler(WebSocketHandler): """ Proxy a websocket connection to a service listening on a given (host, port) pair """ def initialize(self, **kwargs): self.remote_address = kwargs.get("address") self.io_stream = IOStream(socket.socket(kwargs.get("family", socket.AF_INET), kwargs.get("type", socket.SOCK_STREAM), 0)) self.filters = kwargs.get("filters", []) self.io_stream.set_close_callback(self.on_close) def open(self): """ Open the connection to the service when the WebSocket connection has been established """ logger.info("Forwarding connection to server %s" % tuple_to_address(self.remote_address)) self.io_stream.connect(self.remote_address, self.on_connect) def on_message(self, message): """ On message received from WebSocket, forward data to the service """ try: data = None if message is None else bytes(message) for filtr in self.filters: data = filtr.ws_to_socket(data=data) if data: self.io_stream.write(data) except Exception as e: logger.exception(e) self.close() def on_close(self, *args, **kwargs): """ When web socket gets closed, close the connection to the service too """ logger.info("Closing connection with peer at %s" % tuple_to_address(self.remote_address)) logger.debug("Received args %s and %s", args, kwargs) #if not self.io_stream._closed: for message in args: self.on_peer_message(message) if not self.io_stream.closed(): self.io_stream.close() self.close() def on_connect(self): """ Callback invoked on connection with mapped service """ logger.info("Connection established with peer at %s" % tuple_to_address(self.remote_address)) self.io_stream.read_until_close(self.on_close, self.on_peer_message) def on_peer_message(self, message): """ On message received from peer service, send back to client through WebSocket """ try: data = None if message is None else bytes(message) for filtr in self.filters: data = filtr.socket_to_ws(data=data) if data: self.write_message(data, binary=True) except FilterException as e: logger.exception(e) self.on_close()
class Connection(object): def __init__(self, host='localhost', port=11211, pool=None): self._host = host self._port = port self._pool = pool self._socket = None self._stream = None self._ioloop = IOLoop.instance() self.connect() def connect(self): try: self._socket = socket(AF_INET, SOCK_STREAM, 0) self._socket.connect((self._host, self._port)) self._stream = IOStream(self._socket, io_loop=self._ioloop) self._stream.set_close_callback(self.on_disconnect) except error as e: raise ConnectionError(e) def disconect(self): callback = self._final_callback self._final_callback = None try: if callback: callback(None) finally: self._stream._close_callback = None self._stream.close() def on_disconnect(self): callback = self._final_callback self._final_callback = None try: if callback: callback(None) finally: logging.debug('asyncmemcached closing connection') self._pool.release(self) def closed(self): return self._stream.closed() def send_command(self, fullcmd, expect_str, callback): self._final_callback = callback if self._stream.closed(): self.connect() with stack_context.StackContext(self.cleanup): if fullcmd[0:3] == 'get' or \ fullcmd[0:4] == 'incr' or \ fullcmd[0:4] == 'decr': self._stream.write(fullcmd, self.read_value) else: self._stream.write( fullcmd, functools.partial(self.read_response, expect_str)) def read_response(self, expect_str): self._stream.read_until( '\r\n', functools.partial(self._expect_callback, expect_str)) def read_value(self): self._stream.read_until('\r\n', self._expect_value_header_callback) def _expect_value_header_callback(self, response): response = response[:-2] if response[:5] == 'VALUE': resp, key, flag, length = response.split() flag = int(flag) length = int(length) self._stream.read_bytes(length + 2, self._expect_value_callback) elif response.isdigit(): try: callback = self._final_callback self._final_callback = None if callback: callback(int(response)) finally: self._pool.release(self) else: try: callback = self._final_callback self._final_callback = None if callback: callback(None) finally: self._pool.release(self) def _expect_value_callback(self, value): value = value[:-2] self._stream.read_until( '\r\n', functools.partial(self._end_value_callback, value)) def _end_value_callback(self, value, response): response = response.rstrip('\r\n') if response == 'END': try: callback = self._final_callback self._final_callback = None if callback: callback(value) finally: self._pool.release(self) else: raise RedisError('error %s' % response) def _expect_callback(self, expect_str, response): response = response.rstrip('\r\n') if response == expect_str: try: callback = self._final_callback self._final_callback = None if callback: callback(None) finally: self._pool.release(self) else: raise RedisError('error %s' % response) @contextlib.contextmanager def cleanup(self): try: yield except Exception as e: logging.warning("uncaught exception", exc_info=True) try: callback = self._final_callback self._final_callback = None if callback: callback(None) finally: self._pool.release(self)
class Trickle(object): """A coroutine-friendly :class:`~tornado.iostream.IOStream` interface. Takes same parameters as ``IOStream``, or takes a single ``IOStream`` as its only parameter. """ def __init__(self, *args, **kwargs): if args and isinstance(args[0], IOStream): if len(args) > 1 or kwargs: raise TypeError('Too many arguments to Trickle()') self.stream = args[0] else: self.stream = IOStream(*args, **kwargs) def connect(self, address, server_hostname=None, timeout=None): """Connects the socket to a remote address without blocking. Like ``IOStream`` :meth:`~tornado.iostream.IOStream.connect`, but returns a :class:`~tornado.concurrent.Future` and takes no callback. """ method = trickle_method('connect', timeout) return method(self, address, server_hostname=server_hostname) def read_until(self, delimiter, timeout=None): """Read up to the given delimiter.. Like ``IOStream`` :meth:`~tornado.iostream.IOStream.read_until`, but returns a :class:`~tornado.concurrent.Future` and takes no callback. """ return trickle_method('read_until', timeout)(self, delimiter) def read_until_regex(self, regex, timeout=None): """Read up to the given regex pattern. Like ``IOStream`` :meth:`~tornado.iostream.IOStream.read_until_regex`, but returns a :class:`~tornado.concurrent.Future` and takes no callback. """ return trickle_method('read_until_regex', timeout)(self, regex) def read_bytes(self, num_bytes, timeout=None): """Read the given number of bytes. Like ``IOStream`` :meth:`~tornado.iostream.IOStream.read_bytes`, but returns a :class:`~tornado.concurrent.Future` and takes no callback or streaming_callback. """ return trickle_method('read_bytes', timeout)(self, num_bytes) @gen.coroutine def read_until_close(self, timeout=None): """Read all remaining data from the socket. Like ``IOStream`` :meth:`~tornado.iostream.IOStream.read_until_close`, but returns a :class:`~tornado.concurrent.Future` and takes no callback or streaming_callback. """ stream = self.stream ioloop_timeout = None if timeout: def on_timeout(): stream.close((socket.timeout, socket.timeout(), None)) ioloop_timeout = stream.io_loop.add_timeout(timeout, callback=on_timeout) stream.read_until_close(callback=(yield gen.Callback(closed))) result = yield Wait(closed) if ioloop_timeout is not None: stream.io_loop.remove_timeout(ioloop_timeout) if stream.error: raise stream.error raise gen.Return(result) def write(self, data, timeout=None, **kwargs): """Write the given data to this stream. Like ``IOStream`` :meth:`~tornado.iostream.IOStream.write`, but returns a :class:`~tornado.concurrent.Future` and takes no callback. yield the returned Future to wait for all data to be written to the stream. """ return trickle_method('write', timeout)(self, data, **kwargs) def closed(self): """Returns true if the stream has been closed.""" return self.stream.closed()
class Connection(Connection): def __init__(self, pool=None, *args, **kwargs): super(Connection, self).__init__(*args, **kwargs) self._pool = pool self._stream = None self._callbacks = [] self._ready = False def __enter__(self): return self def __exit__(self, exc_type, exc_val, exc_tb): self._pool.release(self) def _add_callback(self, func): self._callbacks.append(func) def _do_callbacks(self): self._ready = True while 1: try: func = self._callbacks.pop() func() except IndexError: # all done break except: # other error continue def connect(self): s = socket.socket(socket.AF_INET, socket.SOCK_STREAM, 0) self._sock = IOStream(s) # tornado iostream self._sock.connect((self._host, self._port), self._do_callbacks) def send(self, payload, correlation_id=-1, callback=None): """ :param payload: an encoded kafka packet :param correlation_id: for now, just for debug logging :return: """ if not self._ready: def _callback(*args, **kwargs): self.send(payload, correlation_id, callback) self._add_callback(_callback) return log.debug("About to send %d bytes to Kafka, request %d" % (len(payload), correlation_id)) if payload: _bytes = struct.pack('>i%ds' % len(payload), len(payload), payload) else: _bytes = struct.pack('>i', -1) try: self._sock.write(_bytes, callback) # simply using sendall except: self.close() callback(None) self._log_and_raise('Unable to send payload to Kafka') def _recv(self, size, callback): try: self._sock.read_bytes(min(size, 4096), callback) except: self.close() callback(None) # if error, set None self._log_and_raise('Unable to receive data from Kafka') def recv(self, correlation_id=-1, callback=None): """ :param correlation_id: for now, just for debug logging :return: kafka response packet """ log.debug("Reading response %d from Kafka" % correlation_id) if not self._ready: def _callback(): self.recv(correlation_id, callback) self._add_callback(_callback) return def get_size(resp): if resp == None: callback(None) size, = struct.unpack('>i', resp) self._recv(size, callback) self._recv(4, get_size) # read the response length def close(self): self._callbacks = [] log.debug("Closing socket connection" + self._log_tail) if self._sock: self._sock.close() self._sock = None else: log.debug("Socket connection not exists" + self._log_tail) def closed(self): return self._sock.closed()
class IPCClient(object): ''' A Tornado IPC client very similar to Tornado's TCPClient class but using either UNIX domain sockets or TCP sockets This was written because Tornado does not have its own IPC server/client implementation. :param IOLoop io_loop: A Tornado ioloop to handle scheduling :param str/int socket_path: A path on the filesystem where a socket belonging to a running IPCServer can be found. It may also be of type 'int', in which case it is used as the port for a tcp localhost connection. ''' def __init__(self, socket_path, io_loop=None): ''' Create a new IPC client IPC clients cannot bind to ports, but must connect to existing IPC servers. Clients can then send messages to the server. ''' self.io_loop = io_loop or tornado.ioloop.IOLoop.current() self.socket_path = socket_path self._closing = False self.stream = None if six.PY2: encoding = None else: encoding = 'utf-8' self.unpacker = msgpack.Unpacker(encoding=encoding) def connected(self): return self.stream is not None and not self.stream.closed() def connect(self, callback=None, timeout=None): ''' Connect to the IPC socket ''' if hasattr( self, '_connecting_future') and not self._connecting_future.done(): # pylint: disable=E0203 future = self._connecting_future # pylint: disable=E0203 else: if hasattr(self, '_connecting_future'): # read previous future result to prevent the "unhandled future exception" error self._connecting_future.exception() # pylint: disable=E0203 future = tornado.concurrent.Future() self._connecting_future = future self._connect(timeout=timeout) if callback is not None: def handle_future(future): response = future.result() self.io_loop.add_callback(callback, response) future.add_done_callback(handle_future) return future @tornado.gen.coroutine def _connect(self, timeout=None): ''' Connect to a running IPCServer ''' if isinstance(self.socket_path, int): sock_type = socket.AF_INET sock_addr = ('127.0.0.1', self.socket_path) else: sock_type = socket.AF_UNIX sock_addr = self.socket_path self.stream = None if timeout is not None: timeout_at = time.time() + timeout while True: if self._closing: break if self.stream is None: with salt.utils.asynchronous.current_ioloop(self.io_loop): self.stream = IOStream( socket.socket(sock_type, socket.SOCK_STREAM)) try: log.trace('IPCClient: Connecting to socket: %s', self.socket_path) yield self.stream.connect(sock_addr) self._connecting_future.set_result(True) break except Exception as e: if self.stream.closed(): self.stream = None if timeout is None or time.time() > timeout_at: if self.stream is not None: self.stream.close() self.stream = None self._connecting_future.set_exception(e) break yield tornado.gen.sleep(1) def __del__(self): try: self.close() except socket.error as exc: if exc.errno != errno.EBADF: # If its not a bad file descriptor error, raise raise except TypeError: # This is raised when Python's GC has collected objects which # would be needed when calling self.close() pass def close(self): ''' Routines to handle any cleanup before the instance shuts down. Sockets and filehandles should be closed explicitly, to prevent leaks. ''' if self._closing: return self._closing = True log.debug('Closing %s instance', self.__class__.__name__) if self.stream is not None and not self.stream.closed(): self.stream.close()
class Flashlight(object): """ TCP клиент фонарика принимает TLV комманды от :host :port и выполняет метод Словарь TLV выполняет функцию роутера команда -> метод для вызова """ def __init__(self, status='OFF', host='127.0.0.1', port=9999, timeout=2): self.status = status self.host = host self.port = port self.timeout = datetime.timedelta(seconds=timeout) self._init_time = datetime.datetime.now() self.color = '#ffffff' self._clear_command() def connect(self): logging.info('Connecting to the server: {} {}'.format(self.host, self.port)) s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) self.stream = IOStream(s) self.stream.set_close_callback(self._on_close) self.stream.connect((self.host, self.port), self._on_connect) IOLoop.instance().start() if self.stream.closed(): logging.error("Can't connect to {} {}".format(self.host, self.port)) def close_connection(self): self.stream.close() IOLoop.instance().stop() IOLoop.instance().close() def on(self): self.status = 'ON' logging.info('Status changed to {status}'.format(status=self.status)) self.send_status() def off(self): self.status = 'OFF' logging.info('Status changed to {status}'.format(status=self.status)) self.send_status() def ch_color(self): value = char_sec_to_int(self.command[3:]) rgb = '#{}'.format(hex(value)[2:]) self.color = rgb logging.info('Switched to {color}'.format(color=self.color)) self.send_status() def send_status(self, callback=None): if not self.stream.reading(): self.stream.write('{:<4}{}\n'.format(self.status, self.color), callback=callback) def _on_close(self): self.close_connection() def _on_connect(self): self._callback() def _callback(self): self.stream.read_bytes(1, self._collect_command) def _clear_command(self): self.command = [] self.length = 0 def _collect_command(self, data): self.command.append(data) if len(self.command) == 3: self.length = char_sec_to_int(self.command[1:3]) elif self.length: self.length -= 1 if not self.length and len(self.command) >= 3: self._run_command() self._callback() def _run_command(self): TLV[ord(self.command[0])]['callback'](self) self._clear_command()
class RconConnection(object): def __init__(self, ip, port, rcon_password): self.ip = ip self.port = port self.rcon_password = rcon_password self._socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM, 0) self._stream = IOStream(self._socket) self.authed = False self.request_id = 0 self.error = None # use a deque for the queue because it supports popleft(), which is # what we want (FIFO) self._queue = deque() self._busy = False # async connect & call _auth when connected self._stream.connect((ip, port), self._auth) def _construct_packet(self, code, body): #packets are in the form dictated at # https://developer.valvesoftware.com/wiki/Source_RCON_Protocol #<packet size><request id><request code><body string><empty string> #strings must be null terminated, even the empty one if body is None: body = r'' self.request_id += 1 # use little endian packing of ints, terminate body with a null, # and add a null at the end for the empty string packet = struct.pack('<l', self.request_id) packet += struct.pack('<l', code) packet += body.encode('ascii') + '\x00\x00' # add packet length to the front of the packet packet = struct.pack('<l', len(packet)) + packet return packet def _send_packet(self, packet, callback=None): logging.debug("Sending packet %s", repr(packet)) self._stream.write(packet, callback) def _read_single_packet(self, callback=None): # reads a single packet from the stream and pushes the processed # response through the given callback if provided #logging.debug("Reading single packet from stream") def process_packet(packed_packet): #logging.debug("Processing packet: %s", repr(packed_packet)) #packet <id (packed int)><response code (packed int)><body>\x00\x00 curr_packet_id = struct.unpack('<l', packed_packet[0:4])[0] response_code = struct.unpack('<l', packed_packet[4:8])[0] message = packed_packet[8:].strip('\x00') #strip the terminators """logging.debug("ID: %d CODE: %d MESSAGE: %s", curr_packet_id, response_code, message)""" # we now have the packet message, response code, and response id, # so push it through the given callback if callback is not None: callback((curr_packet_id, response_code, message)) def process_packet_len(packed_packet_len): packet_len = struct.unpack('<l', packed_packet_len)[0] """logging.debug("Processing packet length: %s, length: %s", repr(packed_packet_len), packet_len)""" # read the entire packet #logging.debug("Reading the rest of the packet") self._stream.read_bytes(packet_len, process_packet) self._stream.read_bytes(4, process_packet_len) def _auth(self, data=None, auth_sent=False, junked=False): """ Called when connect is successful. First thing we always do after connecting is attempt to authenticate. Authentication is asynchronous. We'll utilise partials to do this """ if self.error: raise self.error if not auth_sent: auth_packet = self._construct_packet(SERVERDATA_AUTH, self.rcon_password) # _auth will be called again once the auth packet has been sent f = partial(self._auth, auth_sent=True) self._send_packet(auth_packet, f) elif not junked: # read the junk packet out, and call again with junked = True f = partial(self._auth, auth_sent=True, junked=True) self._read_single_packet(f) else: # now read the real response. call _auth_response once it has been # processed self._read_single_packet(self._auth_response) def _auth_response(self, response): logging.debug("Auth response: %s", repr(response)) if response[1] == SERVERDATA_AUTH_RESPONSE: if response[0] == -1: self.error = RconAuthError("Invalid RCON password specified") elif response[0] == self.request_id: self.authed = True logging.debug("Successfully authed") self._process_queue() else: self.error = RconAuthError("Expected packet id %d, got %d" % (response[0], self.request_id)) else: self.error = RconAuthError("Expected auth response, got %d" % (response[1])) def _exec(self, command, callback=None): """ Send a command packet to the server if we are authenticated. When doing this, we send a packet with SERVERDATA_EXEC_COMMAND and the command we want to execute, along with an empty SERVERDATA_COMMAND_RESPONSE, which the server will mirror back at us once it has sent the response to our command (in order, because TCP is ordered). Doing this lets us know exactly when we've received the full response to our command. """ if self.authed: self._busy = True # send command packet with no callback, it'll just execute and # and do nothing packet = self._construct_packet(SERVERDATA_EXEC_COMMAND, command) self._send_packet(packet) # add empty packet, with callback. this packet is just appended # to the stream's internal write queue packet = self._construct_packet(SERVERDATA_COMMAND_RESPONSE, r'') f = partial(self._command_sent_callback, handle_callback=callback) self._send_packet(packet, f) def _command_sent_callback(self, handle_callback=None): """ Called when a complete command has been sent (command + mirror packets) This will let us know when we should start reading. @param handle_callback The callback to be used once we've read all data from a response """ # begin reading f = partial(self._handle_multi_packet_read, complete_callback=handle_callback) self._read_single_packet(f) def _handle_multi_packet_read(self, data, previous=None, complete_callback=None): # data is a tuple in the form (id, code, message) if previous is None: previous = [data] else: previous.append(data) length = len(previous) if length > 0: """ To signify the end of a multi-line response, we'll receive an empty packet (which is the mirror of our empty packet), along with an additional empty packet whose body consists of solely \x01. So we need to check the last two packets. This means we'll receive a MINIMUM of THREE packets for ANY command response """ response_complete = False got_mirror_packet = False empty = previous[length - 2:] # last 2 packets in previous list for packet in empty: packet_id = packet[0] response_code = packet[1] message = packet[2] if (response_code == SERVERDATA_COMMAND_RESPONSE and packet_id == self.request_id): if got_mirror_packet and message == '\x01': """ Have already gotten mirror packet, so this packet is the expected packet after the mirror, with body of \x01 Therefore, our response is complete! """ response_complete = True else: got_mirror_packet = True if response_complete: self._busy = False if complete_callback is not None: complete = self._compile_multi_packet(previous) complete_callback(complete) self._process_queue() elif not response_complete: # response not complete, get the next packet f = partial(self._handle_multi_packet_read, previous=previous, complete_callback=complete_callback) self._read_single_packet(f) def _compile_multi_packet(self, packets): """ Compiles a list of packets into a single packet, chopping off the empty packets. @param packets A list of packet tuples @return Complete packet """ first = packets[0] response_id = first[0] response_code = first[1] message = "".join([x[2] for x in packets[:-2]]) return (response_id, response_code, message) def send_cmd(self, command, callback=None): if self.error: raise self.error if self.busy() or not self.authed: # we're already reading/writing from the socket. this command # should be queued. the queue is processed at the end of a # read cycle logging.debug("Stream is busy. Adding command to queue") self._add_to_queue((command, callback)) else: # execute the command! self._exec(command, callback) def _add_to_queue(self, qtuple): self._queue.append(qtuple) def _process_queue(self): """ Process any RCON commands in the queue. Called as soon as we are authed and after a command has been completed. Commands are added to the queue if `send_cmd` is called when we are busy (i.e authing or reading/writing) """ logging.debug("Processing RCON command queue") try: command, callback = self._queue.popleft() logging.debug("QUEUE - command: %s, callback: %s", command, callback) self.send_cmd(command, callback) except IndexError: pass except: logging.exception("Exception processing queue") def busy(self): return (self._stream.reading() or self._stream.writing() or self._busy) @property def closed(self): return self._stream.closed()
class _HTTPConnection(object): _SUPPORTED_METHODS = set(["GET", "HEAD", "POST", "PUT", "DELETE", "PATCH", "OPTIONS"]) # map of (host, port, af) ==> {'addrinfo': addrinfo, # 'last_update': timestamp} _DNS_CACHE = {} def __init__(self, io_loop, client, request, callback, max_buffer_size): self.start_time = time.time() self.io_loop = io_loop self.client = client self.request = request self.callback = callback self.code = None self.headers = None self.chunks = None self._decompressor = None # Timeout handle returned by IOLoop.add_timeout self._timeout = None self._connect_timeout = None if self.request.headers.get('Connection') == 'close': self.keep_alive = False else: self.keep_alive = True with stack_context.StackContext(self.cleanup): parsed = urlparse.urlsplit(_unicode(self.request.url)) if ssl is None and parsed.scheme == "https": raise ValueError("HTTPS requires either python2.6+ or " "curl_httpclient") if parsed.scheme not in ("http", "https"): raise ValueError("Unsupported url scheme: %s" % self.request.url) # urlsplit results have hostname and port results, but they # didn't support ipv6 literals until python 2.7. netloc = parsed.netloc if "@" in netloc: userpass, _, netloc = netloc.rpartition("@") match = re.match(r'^(.+):(\d+)$', netloc) if match: host = match.group(1) port = int(match.group(2)) else: host = netloc port = 443 if parsed.scheme == "https" else 80 if re.match(r'^\[.*\]$', host): # raw ipv6 addresses in urls are enclosed in brackets host = host[1:-1] parsed_hostname = host # save final parsed host for _on_connect if self.client.hostname_mapping is not None: host = self.client.hostname_mapping.get(host, host) if request.allow_ipv6: af = socket.AF_UNSPEC else: # We only try the first IP we get from getaddrinfo, # so restrict to ipv4 by default. af = socket.AF_INET # Ignore keep_alive logic if explicitly requesting non-presistent connections if self.keep_alive: self.stream_key = (host, port, parsed.scheme) if self.client.stream_map.has_key(self.stream_key): self._set_stream(parsed, parsed_hostname) if self.stream: return else: self.client.stream_map[self.stream_key] = collections.deque() af, socktype, proto, canonname, sockaddr = self._getaddrinfo(host, port, af) if parsed.scheme == "https": ssl_options = {} if request.validate_cert: ssl_options["cert_reqs"] = ssl.CERT_REQUIRED if request.ca_certs is not None: ssl_options["ca_certs"] = request.ca_certs else: ssl_options["ca_certs"] = _DEFAULT_CA_CERTS if request.client_key is not None: ssl_options["keyfile"] = request.client_key if request.client_cert is not None: ssl_options["certfile"] = request.client_cert # SSL interoperability is tricky. We want to disable # SSLv2 for security reasons; it wasn't disabled by default # until openssl 1.0. The best way to do this is to use # the SSL_OP_NO_SSLv2, but that wasn't exposed to python # until 3.2. Python 2.7 adds the ciphers argument, which # can also be used to disable SSLv2. As a last resort # on python 2.6, we set ssl_version to SSLv3. This is # more narrow than we'd like since it also breaks # compatibility with servers configured for TLSv1 only, # but nearly all servers support SSLv3: # http://blog.ivanristic.com/2011/09/ssl-survey-protocol-support.html if sys.version_info >= (2, 7): ssl_options["ciphers"] = "DEFAULT:!SSLv2" else: # This is really only necessary for pre-1.0 versions # of openssl, but python 2.6 doesn't expose version # information. ssl_options["ssl_version"] = ssl.PROTOCOL_SSLv3 self.stream = SSLIOStream(socket.socket(af, socktype, proto), io_loop=self.io_loop, ssl_options=ssl_options, max_buffer_size=max_buffer_size) self.stream.last_used = time.time() else: self.stream = IOStream(socket.socket(af, socktype, proto), io_loop=self.io_loop, max_buffer_size=max_buffer_size) self.stream.last_used = time.time() timeout = min(request.connect_timeout, request.request_timeout) if timeout: self._connect_timeout = self.io_loop.add_timeout( self.start_time + timeout, self._on_timeout) self.stream.set_close_callback(self._on_close) self.stream.connect(sockaddr, functools.partial(self._on_connect, parsed, parsed_hostname)) def _getaddrinfo(self, host, port, af): """ Custom getaddrinfo method caches the response for `DNS_TTL` seconds """ key = (host, port, af) # cache based on this tuple now = datetime.datetime.now() # if bool(DNS_TTL) == False (e.g. 0.0, None, etc), always # ignore cache if not DNS_TTL or not key in self._DNS_CACHE: addrinfo = socket.getaddrinfo(host, port, af, socket.SOCK_STREAM, 0, 0) #Assign TTL for each IP addrinfo_dict = dict([(i,now) for i in addrinfo]) self._DNS_CACHE[key] = {'addrinfo': addrinfo_dict,"last_changed":now,"last_addrinfo":addrinfo,"last_updated":now} #If past DNS_TTL or it was changed recently elif random.randint(0,(now-self._DNS_CACHE[key]['last_changed']).seconds+1000) < 10 or (now-self._DNS_CACHE[key]['last_updated']).seconds>DNS_TTL: addrinfo = socket.getaddrinfo(host, port, af, socket.SOCK_STREAM, 0, 0) self._DNS_CACHE[key]['last_updated'] = now if addrinfo != self._DNS_CACHE[key]['last_addrinfo']: self._DNS_CACHE[key]['last_changed'] = now self._DNS_CACHE[key]['last_addrinfo'] = addrinfo for info in addrinfo: self._DNS_CACHE[key]['addrinfo'].update({info:now}) random_addr = None while not random_addr and self._DNS_CACHE[key]['addrinfo']: random_addr = random.choice(self._DNS_CACHE[key]['addrinfo'].keys()) if (now-self._DNS_CACHE[key]['addrinfo'][random_addr]).seconds > DNS_TTL: del self._DNS_CACHE[key]['addrinfo'][random_addr] random_addr = None if not random_addr: raise Exception("Error. No DNS could be resolved. We should never get here.") return random_addr @classmethod def clear_dns_cache(cls): cls._DNS_CACHE = {} def _set_stream(self, parsed, parsed_hostname): while self.client.stream_map[self.stream_key]: self.stream = self.client.stream_map[self.stream_key].pop() # Ditch closed streams and get a new one if self.stream.closed(): continue # Double check the stream isn't in use # Don't put back in the queue because if it's in use whoever's using it will, # or if it's closed it shouldn't be there if not (self.stream.reading() or self.stream.writing()): self.stream.set_close_callback(self._on_close) self._on_connect(parsed, parsed_hostname) self.stream.last_used = time.time() self.client._garbage_collect(self.stream_key) return else: self.stream.close() logging.warning("Stream is writing, this should be fixed now though") self.stream = None def _on_timeout(self): self._timeout = None self._run_callback(HTTPResponse(self.request, 599, request_time=time.time() - self.start_time, error=HTTPError(599, "Timeout"))) self.stream.close() def _on_connect(self, parsed, parsed_hostname): if self._connect_timeout is not None: self.io_loop.remove_timeout(self._connect_timeout) self._connect_timeout = None if self._timeout is not None: self.io_loop.remove_timeout(self._timeout) self._timeout = None if self.request.request_timeout: self._timeout = self.io_loop.add_timeout( self.start_time + self.request.request_timeout, self._on_timeout) if (self.request.validate_cert and isinstance(self.stream, SSLIOStream)): match_hostname(self.stream.socket.getpeercert(), # ipv6 addresses are broken (in # parsed.hostname) until 2.7, here is # correctly parsed value calculated in # __init__ parsed_hostname) if (self.request.method not in self._SUPPORTED_METHODS and not self.request.allow_nonstandard_methods): raise KeyError("unknown method %s" % self.request.method) for key in ('network_interface', 'proxy_host', 'proxy_port', 'proxy_username', 'proxy_password'): if getattr(self.request, key, None): raise NotImplementedError('%s not supported' % key) if "Connection" not in self.request.headers: self.request.headers["Connection"] = "close" if "Host" not in self.request.headers: if '@' in parsed.netloc: self.request.headers["Host"] = parsed.netloc.rpartition('@')[-1] else: self.request.headers["Host"] = parsed.netloc username, password = None, None if parsed.username is not None: username, password = parsed.username, parsed.password elif self.request.auth_username is not None: username = self.request.auth_username password = self.request.auth_password or '' if username is not None: auth = utf8(username) + b(":") + utf8(password) self.request.headers["Authorization"] = (b("Basic ") + base64.b64encode(auth)) if self.request.user_agent: self.request.headers["User-Agent"] = self.request.user_agent if not self.request.allow_nonstandard_methods: if self.request.method in ("POST", "PATCH", "PUT"): assert self.request.body is not None else: assert self.request.body is None if self.request.body is not None: self.request.headers["Content-Length"] = str(len( self.request.body)) if (self.request.method == "POST" and "Content-Type" not in self.request.headers): self.request.headers["Content-Type"] = "application/x-www-form-urlencoded" if self.request.use_gzip: self.request.headers["Accept-Encoding"] = "gzip" req_path = ((parsed.path or '/') + (('?' + parsed.query) if parsed.query else '')) request_lines = [utf8("%s %s HTTP/1.1" % (self.request.method, req_path))] for k, v in self.request.headers.get_all(): line = utf8(k) + b(": ") + utf8(v) if b('\n') in line: raise ValueError('Newline in header: ' + repr(line)) request_lines.append(line) self.stream.write(b("\r\n").join(request_lines) + b("\r\n\r\n")) if self.request.body is not None: self.stream.write(self.request.body) self.stream.read_until(b("\r\n\r\n"), self._on_headers) def _run_callback(self, response): if self.callback is not None: callback = self.callback self.callback = None callback(response) @contextlib.contextmanager def cleanup(self): try: yield except Exception, e: logging.warning("uncaught exception: %s" % self.request.url, exc_info=True) self._run_callback(HTTPResponse(self.request, 599, error=e, request_time=time.time() - self.start_time)) if hasattr(self, 'stream'): self._close_stream()
class IPCClient(object): ''' A Tornado IPC client very similar to Tornado's TCPClient class but using either UNIX domain sockets or TCP sockets This was written because Tornado does not have its own IPC server/client implementation. :param IOLoop io_loop: A Tornado ioloop to handle scheduling :param str/int socket_path: A path on the filesystem where a socket belonging to a running IPCServer can be found. It may also be of type 'int', in which case it is used as the port for a tcp localhost connection. ''' # Create singleton map between two sockets instance_map = weakref.WeakKeyDictionary() def __new__(cls, socket_path, io_loop=None): io_loop = io_loop or tornado.ioloop.IOLoop.current() if io_loop not in IPCClient.instance_map: IPCClient.instance_map[io_loop] = weakref.WeakValueDictionary() loop_instance_map = IPCClient.instance_map[io_loop] # FIXME key = str(socket_path) if key not in loop_instance_map: log.debug('Initializing new IPCClient for path: {0}'.format(key)) new_client = object.__new__(cls) # FIXME new_client.__singleton_init__(io_loop=io_loop, socket_path=socket_path) loop_instance_map[key] = new_client else: log.debug('Re-using IPCClient for {0}'.format(key)) return loop_instance_map[key] def __singleton_init__(self, socket_path, io_loop=None): ''' Create a new IPC client IPC clients cannot bind to ports, but must connect to existing IPC servers. Clients can then send messages to the server. ''' self.io_loop = io_loop or tornado.ioloop.IOLoop.current() self.socket_path = socket_path self._closing = False self.stream = None if six.PY2: encoding = None else: encoding = 'utf-8' self.unpacker = msgpack.Unpacker(encoding=encoding) def __init__(self, socket_path, io_loop=None): # Handled by singleton __new__ pass def connected(self): return self.stream is not None and not self.stream.closed() def connect(self, callback=None, timeout=None): ''' Connect to the IPC socket ''' if hasattr( self, '_connecting_future') and not self._connecting_future.done(): # pylint: disable=E0203 future = self._connecting_future # pylint: disable=E0203 else: future = tornado.concurrent.Future() self._connecting_future = future self.io_loop.add_callback(self._connect, timeout=timeout) if callback is not None: def handle_future(future): response = future.result() self.io_loop.add_callback(callback, response) future.add_done_callback(handle_future) return future @tornado.gen.coroutine def _connect(self, timeout=None): ''' Connect to a running IPCServer ''' if isinstance(self.socket_path, int): sock_type = socket.AF_INET sock_addr = ('127.0.0.1', self.socket_path) else: sock_type = socket.AF_UNIX sock_addr = self.socket_path self.stream = None if timeout is not None: timeout_at = time.time() + timeout while True: if self._closing: break if self.stream is None: self.stream = IOStream( socket.socket(sock_type, socket.SOCK_STREAM), io_loop=self.io_loop, ) try: log.trace('IPCClient: Connecting to socket: {0}'.format( self.socket_path)) yield self.stream.connect(sock_addr) self._connecting_future.set_result(True) break except Exception as e: if self.stream.closed(): self.stream = None if timeout is None or time.time() > timeout_at: if self.stream is not None: self.stream.close() self.stream = None self._connecting_future.set_exception(e) break yield tornado.gen.sleep(1) def __del__(self): self.close() def close(self): ''' Routines to handle any cleanup before the instance shuts down. Sockets and filehandles should be closed explicitly, to prevent leaks. ''' if self._closing: return self._closing = True if self.stream is not None and not self.stream.closed(): self.stream.close()
class NTcpConnector(object): def __init__(self, host, port): self.routes = {} self.host = host self.port = port self._s = socket.socket(socket.AF_INET, socket.SOCK_STREAM, 0) self.stream = IOStream(self._s) self.stream.connect((self.host, self.port), self._start_recv) def unregister(self, client): self.routes = dict( filter(lambda x: x[1] != client, self.routes.items())) def __lt__(self, other): return id(self) < id(other) def sendMsg(self, client, content): sn = client.application.proxys.getSN() self.routes[sn] = client data = struct.pack('<i6I%dsI' % len(content), int(-1), 10020, 20 + len(content), sn, 0, int(time.time()), 1, content.encode('utf-8'), int((20 + len(content)) ^ 0xaaaaaaaa)) self.stream.write(data) def is_connected(self): return not self.stream.closed() def invalidate(self): self.stream.close_fd() def _start_recv(self): self.stream.read_bytes(12, self._on_frame) def _on_frame(self, data): nLen = struct.unpack('<i2I', data)[2] self.stream.read_bytes(nLen, self._on_msg) def _on_msg(self, data): nLen = len(data) sn, nTag, nTime, nCmdId, dataS = struct.unpack('<4I%dsI' % (nLen - 20), data)[0:-1] if sn == 0: self.stream.write( struct.pack('<i7I', int(-1), 10000, 20, 0, 0, int(time.time()), 0, int(20 ^ 0xaaaaaaaa))) elif sn > 0 and (sn in self.routes): fs, strField = {}, '' if nCmdId == 110 and nLen == 292: #十档报价 ds = struct.unpack( '<2iIqIqIqIqIqIqIqIqIqIqIqIqIqIqIqIqIqIqIqIqIqIq', dataS) strField = 'nSecurityID,nTime,nPxBid1,llVolumeBid1,nPxBid2,llVolumeBid2,nPxBid3,llVolumeBid3,nPxBid4,llVolumeBid4,nPxBid5,llVolumeBid5,nPxBid6,llVolumeBid6,nPxBid7,llVolumeBid7,nPxBid8,llVolumeBid8,nPxBid9,llVolumeBid9,nPxBid10,llVolumeBid10,nWeightedAvgBidPx,llTotalBidVolume,nPxOffer1,llVolumeOffer1,nPxOffer2,llVolumeOffer2,nPxOffer3,llVolumeOffer3,nPxOffer4,llVolumeOffer4,nPxOffer5,llVolumeOffer5,nPxOffer6,llVolumeOffer6,nPxOffer7,llVolumeOffer7,nPxOffer8,llVolumeOffer8,nPxOffer9,llVolumeOffer9,nPxOffer10,llVolumeOffer10,nWeightedAvgOfferPx,llTotalOfferVolume' elif nCmdId == 165 and nLen == 644: #委托明细 ds = struct.unpack('<2iI3i150i', dataS) strField = 'nSecurityID,nTime,nPx,nLevel,nOrderCount,nRevealCount,nStatus1,nVolume1,nChangeVolume1,nStatus2,nVolume2,nChangeVolume2,nStatus3,nVolume3,nChangeVolume3,nStatus4,nVolume4,nChangeVolume4,nStatus5,nVolume5,nChangeVolume5,nStatus6,nVolume6,nChangeVolume6,nStatus7,nVolume7,nChangeVolume7,nStatus8,nVolume8,nChangeVolume8,nStatus9,nVolume9,nChangeVolume9,nStatus10,nVolume10,nChangeVolume10,nStatus11,nVolume11,nChangeVolume11,nStatus12,nVolume12,nChangeVolume12,nStatus13,nVolume13,nChangeVolume13,nStatus14,nVolume14,nChangeVolume14,nStatus15,nVolume15,nChangeVolume15,nStatus16,nVolume16,nChangeVolume16,nStatus17,nVolume17,nChangeVolume17,nStatus18,nVolume18,nChangeVolume18,nStatus19,nVolume19,nChangeVolume19,nStatus20,nVolume20,nChangeVolume20,nStatus21,nVolume21,nChangeVolume21,nStatus22,nVolume22,nChangeVolume22,nStatus23,nVolume23,nChangeVolume23,nStatus24,nVolume24,nChangeVolume24,nStatus25,nVolume25,nChangeVolume25,nStatus26,nVolume26,nChangeVolume26,nStatus27,nVolume27,nChangeVolume27,nStatus28,nVolume28,nChangeVolume28,nStatus29,nVolume29,nChangeVolume29,nStatus30,nVolume30,nChangeVolume30,nStatus31,nVolume31,nChangeVolume31,nStatus32,nVolume32,nChangeVolume32,nStatus33,nVolume33,nChangeVolume33,nStatus34,nVolume34,nChangeVolume34,nStatus35,nVolume35,nChangeVolume35,nStatus36,nVolume36,nChangeVolume36,nStatus37,nVolume37,nChangeVolume37,nStatus38,nVolume38,nChangeVolume38,nStatus39,nVolume39,nChangeVolume39,nStatus40,nVolume40,nChangeVolume40,nStatus41,nVolume41,nChangeVolume41,nStatus42,nVolume42,nChangeVolume42,nStatus43,nVolume43,nChangeVolume43,nStatus44,nVolume44,nChangeVolume44,nStatus45,nVolume45,nChangeVolume45,nStatus46,nVolume46,nChangeVolume46,nStatus47,nVolume47,nChangeVolume47,nStatus48,nVolume48,nChangeVolume48,nStatus49,nVolume49,nChangeVolume49,nStatus50,nVolume50,nChangeVolume50' else: pass if strField: fields = strField.split(',') for i in range(0, len(fields)): fs[fields[i]] = ds[i] fs['nCmdId'] = nCmdId self.routes[sn].callback(fs) self._start_recv()
class Client(object): def __init__(self, host='localhost', port=11300, connect_timeout=socket.getdefaulttimeout(), io_loop=None): self._connect_timeout = connect_timeout self.host = host self.port = port self.io_loop = io_loop or IOLoop.instance() self._stream = None self._using = 'default' # current tube self._watching = set(['default']) # set of watched tubes self._queue = deque() self._talking = False self._reconnect_cb = None def _reconnect(self): # wait some time before trying to re-connect self.io_loop.add_timeout(time.time() + RECONNECT_TIMEOUT, lambda: self.connect(self._reconnected)) def _reconnected(self): # re-establish the used tube and tubes being watched watch = self._watching.difference(['default']) # ignore "default", if it is not in the client's watch list ignore = set(['default']).difference(self._watching) def do_next(_=None): try: if watch: self.watch(watch.pop(), do_next) elif ignore: self.ignore(ignore.pop(), do_next) elif self._using != 'default': # change the tube used, and callback to user self.use(self._using, self._reconnect_cb) elif self._reconnect_cb: # callback to user self._reconnect_cb() except: # ignored, as next re-connect will retry the operation pass do_next() @coroutine def connect(self): """Connect to beanstalkd server.""" if not self.closed(): return self._talking = False self._socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM, socket.IPPROTO_TCP) self._stream = IOStream(self._socket, io_loop=self.io_loop) self._stream.set_close_callback(self._reconnect) yield Task(self._stream.connect, (self.host, self.port)) def set_reconnect_callback(self, callback): """Set callback to be called if connection has been lost and re-established again. If the connection is closed unexpectedly, the client will automatically attempt to re-connect with 1 second intervals. After re-connecting, the client will attempt to re-establish the used tube and watched tubes. """ self._reconnect_cb = callback @coroutine def close(self): """Close connection to server.""" key = object() if self._stream: self._stream.set_close_callback((yield Callback(key))) if not self.closed(): yield Task(self._stream.write, b'quit\r\n') self._stream.close() yield Wait(key) def closed(self): """"Returns True if the connection is closed.""" return not self._stream or self._stream.closed() def _interact(self, request, callback): # put the interaction request into the FIFO queue cb = stack_context.wrap(callback) self._queue.append((request, cb)) self._process_queue() def _process_queue(self): if self._talking or not self._queue: return # pop a request of the queue and perform the send-receive interaction self._talking = True with stack_context.NullContext(): req, cb = self._queue.popleft() command = req.cmd + b'\r\n' if req.body: command += req.body + b'\r\n' # write command and body to socket stream self._stream.write(command, # when command is written: read line from socket stream lambda: self._stream.read_until(b'\r\n', # when a line has been read: return status and results lambda data: self._recv(req, data, cb))) def _recv(self, req, data, cb): # parse the data received as server response spl = data.decode('utf8').split() status, values = spl[0], spl[1:] error = None err_args = ObjectDict(request=req, status=status, values=values) if req.ok and status in req.ok: # avoid raising a Buried exception when using the bury command pass elif status == 'BURIED': error = Buried(**err_args) elif status == 'TIMED_OUT': error = TimedOut(**err_args) elif status == 'DEADLINE_SOON': error = DeadlineSoon(**err_args) elif req.err and status in req.err: error = CommandFailed(**err_args) else: error = UnexpectedResponse(**err_args) resp = Bunch(req=req, status=status, values=values, error=error) if error or not req.read_body: # end the request and callback with results self._do_callback(cb, resp) else: # read the body including the terminating two bytes of crlf if len(values) == 2: job_id, size = int(values[0]), int(values[1]) resp.job_id = int(job_id) else: size = int(values[0]) self._stream.read_bytes(size + 2, lambda data: self._recv_body(data[:-2], resp, cb)) def _recv_body(self, data, resp, cb): if resp.req.parse_yaml: # parse the yaml encoded body self._parse_yaml(data, resp, cb) else: # don't parse body, it is a job! # end the request and callback with results resp.body = ObjectDict(id=resp.job_id, body=data) self._do_callback(cb, resp) def _parse_yaml(self, data, resp, cb): # dirty parsing of yaml data # (assumes that data is a yaml encoded list or dict) spl = data.decode('utf8').split('\n')[1:-1] if spl[0].startswith('- '): # it is a list resp.body = [s[2:] for s in spl] else: # it is a dict conv = lambda v: ((float(v) if '.' in v else int(v)) if v.replace('.', '', 1).isdigit() else v) resp.body = ObjectDict((k, conv(v.strip())) for k, v in (s.split(':') for s in spl)) self._do_callback(cb, resp) def _do_callback(self, cb, resp): # end the request and process next item in the queue # and callback with results self._talking = False self.io_loop.add_callback(self._process_queue) if not cb: return # default is to callback with error state (None or exception) obj = None req = resp.req if resp.error: obj = resp.error elif req.read_value: # callback with an integer value or a string if resp.values[0].isdigit(): obj = int(resp.values[0]) else: obj = resp.values[0] elif req.read_body: # callback with the body (job or parsed yaml) obj = resp.body self.io_loop.add_callback(lambda: cb(obj)) # # Producer commands # @coroutine def put(self, body, priority=DEFAULT_PRIORITY, delay=0, ttr=120): """Put a job body (a byte string) into the current tube. The job can be delayed a number of seconds, before it is put in the ready queue, default is no delay. The job is assigned a Time To Run (ttr, in seconds), the mininum is 1 sec., default is ttr=120 sec. Calls back with id when job is inserted. If an error occured, the callback gets a Buried or CommandFailed exception. The job is buried when either the body is too big, so server ran out of memory, or when the server is in draining mode. """ cmd = 'put {} {} {} {}'.format(priority, delay, ttr, len(body)).encode('utf8') assert isinstance(body, bytes) request = Bunch(cmd=cmd, ok=['INSERTED'], err=['BURIED', 'JOB_TOO_BIG', 'DRAINING'], body=body, read_value=True) resp = yield Task(self._interact, request) raise Return(resp) @coroutine def use(self, name): """Use the tube with given name. Calls back with the name of the tube now being used. """ cmd = 'use {}'.format(name).encode('utf8') request = Bunch(cmd=cmd, ok=['USING'], read_value=True) resp = yield Task(self._interact, request) if not isinstance(resp, Exception): self._using = resp raise Return(resp) # # Worker commands # @coroutine def reserve(self, timeout=None): """Reserve a job from one of the watched tubes, with optional timeout in seconds. Not specifying a timeout (timeout=None, the default) will make the client put the communication with beanstalkd on hold, until either a job is reserved, or a already reserved job is approaching it's TTR deadline. Commands issued while waiting for the "reserve" callback will be queued and sent in FIFO order, when communication is resumed. A timeout value of 0 will cause the server to immediately return either a response or TIMED_OUT. A positive value of timeout will limit the amount of time the client will will hold communication until a job becomes available. Calls back with a job dict (keys id and body). If the request timed out, the callback gets a TimedOut exception. If a reserved job has deadline within the next second, the callback gets a DeadlineSoon exception. """ if timeout is not None: cmd = 'reserve-with-timeout {}'.format(timeout).encode('utf8') else: cmd = b'reserve' request = Bunch(cmd=cmd, ok=['RESERVED'], err=['DEADLINE_SOON', 'TIMED_OUT'], read_body=True) resp = yield Task(self._interact, request) raise Return(resp) @coroutine def delete(self, job_id): """Delete job with given id. Calls back when job is deleted. If the job does not exist, or it is not neither reserved by the client, ready or buried; the callback gets a CommandFailed exception. """ cmd = 'delete {}'.format(job_id).encode('utf8') request = Bunch(cmd=cmd, ok=['DELETED'], err=['NOT_FOUND']) resp = yield Task(self._interact, request) raise Return(resp) @coroutine def release(self, job_id, priority=DEFAULT_PRIORITY, delay=0): """Release a reserved job back into the ready queue. A new priority can be assigned to the job. It is also possible to specify a delay (in seconds) to wait before putting the job in the ready queue. The job will be in the "delayed" state during this time. Calls back when job is released. If the job was buried, the callback gets a Buried exception. If the job does not exist, or it is not reserved by the client, the callback gets a CommandFailed exception. """ cmd = 'release {} {} {}'.format(job_id, priority, delay).encode('utf8') request = Bunch(cmd=cmd, ok=['RELEASED'], err=['BURIED', 'NOT_FOUND']) resp = yield Task(self._interact, request) raise Return(resp) @coroutine def bury(self, job_id, priority=DEFAULT_PRIORITY): """Bury job with given id. A new priority can be assigned to the job. Calls back when job is burried. If the job does not exist, or it is not reserved by the client, the callback gets a CommandFailed exception. """ cmd = 'bury {} {}'.format(job_id, priority).encode('utf8') request = Bunch(cmd=cmd, ok=['BURIED'], err=['NOT_FOUND']) resp = yield Task(self._interact, request) raise Return(resp) @coroutine def touch(self, job_id): """Touch job with given id. This is for requesting more time to work on a reserved job before it expires. Calls back when job is touched. If the job does not exist, or it is not reserved by the client, the callback gets a CommandFailed exception. """ cmd = 'touch {}'.format(job_id).encode('utf8') request = Bunch(cmd=cmd, ok=['TOUCHED'], err=['NOT_FOUND']) resp = yield Task(self._interact, request) raise Return(resp) @coroutine def watch(self, name): """Watch tube with given name. Calls back with number of tubes currently in the watch list. """ cmd = 'watch {}'.format(name).encode('utf8') request = Bunch(cmd=cmd, ok=['WATCHING'], read_value=True) resp = yield Task(self._interact, request) # add to the client's watch list self._watching.add(name) raise Return(resp) @coroutine def ignore(self, name): """Stop watching tube with given name. Calls back with the number of tubes currently in the watch list. On an attempt to ignore the only tube in the watch list, the callback gets a CommandFailed exception. """ cmd = 'ignore {}'.format(name).encode('utf8') request = Bunch(cmd=cmd, ok=['WATCHING'], err=['NOT_IGNORED'], read_value=True) resp = yield Task(self._interact, request) if name in self._watching: # remove from the client's watch list self._watching.remove(name) raise Return(resp) # # Other commands # def _peek(self, variant, callback): # a shared gateway for the peek* commands cmd = 'peek{}'.format(variant).encode('utf8') request = Bunch(cmd=cmd, ok=['FOUND'], err=['NOT_FOUND'], read_body=True) self._interact(request, callback) @coroutine def peek(self, job_id): """Peek at job with given id. Calls back with a job dict (keys id and body). If no job exists with that id, the callback gets a CommandFailed exception. """ resp = yield Task(self._peek, ' {}'.format(job_id)) raise Return(resp) @coroutine def peek_ready(self): """Peek at next ready job in the current tube. Calls back with a job dict (keys id and body). If no ready jobs exist, the callback gets a CommandFailed exception. """ resp = yield Task(self._peek, '-ready') raise Return(resp) @coroutine def peek_delayed(self): """Peek at next delayed job in the current tube. Calls back with a job dict (keys id and body). If no delayed jobs exist, the callback gets a CommandFailed exception. """ resp = yield Task(self._peek, '-delayed') raise Return(resp) @coroutine def peek_buried(self): """Peek at next buried job in the current tube. Calls back with a job dict (keys id and body). If no buried jobs exist, the callback gets a CommandFailed exception. """ resp = yield Task(self._peek, '-buried') raise Return(resp) @coroutine def kick(self, bound=1): """Kick at most `bound` jobs into the ready queue from the current tube. Calls back with the number of jobs actually kicked. """ cmd = 'kick {}'.format(bound).encode('utf8') request = Bunch(cmd=cmd, ok=['KICKED'], read_value=True) resp = yield Task(self._interact, request) raise Return(resp) @coroutine def kick_job(self, job_id): """Kick job with given id into the ready queue. (Requires Beanstalkd version >= 1.8) Calls back when job is kicked. If no job exists with that id, or if job is not in a kickable state, the callback gets a CommandFailed exception. """ cmd = 'kick-job {}'.format(job_id).encode('utf8') request = Bunch(cmd=cmd, ok=['KICKED'], err=['NOT_FOUND']) resp = yield Task(self._interact, request) raise Return(resp) @coroutine def stats_job(self, job_id): """A dict of stats about the job with given id. If no job exists with that id, the callback gets a CommandFailed exception. """ cmd = 'stats-job {}'.format(job_id).encode('utf8') request = Bunch(cmd=cmd, ok=['OK'], err=['NOT_FOUND'], read_body=True, parse_yaml=True) resp = yield Task(self._interact, request) raise Return(resp) @coroutine def stats_tube(self, name): """A dict of stats about the tube with given name. If no tube exists with that name, the callback gets a CommandFailed exception. """ cmd = 'stats-tube {}'.format(name).encode('utf8') request = Bunch(cmd=cmd, ok=['OK'], err=['NOT_FOUND'], read_body=True, parse_yaml=True) resp = yield Task(self._interact, request) raise Return(resp) @coroutine def stats(self): """A dict of beanstalkd statistics.""" request = Bunch(cmd=b'stats', ok=['OK'], read_body=True, parse_yaml=True) resp = yield Task(self._interact, request) raise Return(resp) @coroutine def list_tubes(self): """List of all existing tubes.""" request = Bunch(cmd=b'list-tubes', ok=['OK'], read_body=True, parse_yaml=True) resp = yield Task(self._interact, request) raise Return(resp) @coroutine def list_tube_used(self): """Name of the tube currently being used.""" request = Bunch(cmd=b'list-tube-used', ok=['USING'], read_value=True) resp = yield Task(self._interact, request) raise Return(resp) @coroutine def list_tubes_watched(self): """List of tubes currently being watched.""" request = Bunch(cmd=b'list-tubes-watched', ok=['OK'], read_body=True, parse_yaml=True) resp = yield Task(self._interact, request) raise Return(resp) @coroutine def pause_tube(self, name, delay): """Delay any new job being reserved from the tube for a given time. The delay is an integer number of seconds to wait before reserving any more jobs from the queue. Calls back when tube is paused. If tube does not exists, the callback will get a CommandFailed exception. """ cmd = 'pause-tube {} {}'.format(name, delay).encode('utf8') request = Bunch(cmd=cmd, ok=['PAUSED'], err=['NOT_FOUND']) resp = yield Task(self._interact, request) raise Return(resp)
class TornadoClient(BaseClient): def __init__(self, jid, password, hostname='localhost', port=5222, tls=False, registercls=None): BaseClient.__init__(self, jid, password, tls, registercls) s = socket.socket(socket.AF_INET, socket.SOCK_STREAM, 0) s.connect((hostname, port)) self.io = IOStream(s) self.io.set_close_callback(self.socket_error) self._read() def socket_error(self): BaseClient.socket_error(self) self.stop() def send_raw_stanza(self, stanza): self.log(stanza, 'OUTGOING') self.io.write(stanza) def start_tls(self): assert ssl, "Python 2.6+ and OpenSSL required for SSL" self.io.io_loop.remove_handler(self.io.socket.fileno()) self.io.socket = ssl.wrap_socket(self.io.socket, server_side=False) self.io.io_loop.add_handler(self.io.socket.fileno(), self.io._handle_events, self.io._state) self.tls_ok() def start(self): self.running = True self.run(start_loop=False) def stop(self, stop_loop=False): self.stopping() if not self.io.closed(): BaseClient.stop(self) self.io._close_callback = None self.io.close() self.cleanup() if stop_loop: ioloop.IOLoop.instance().stop() self.terminated() def run(self, start_loop=True): """ Starts the client. ``start_loop`` True - flag indicating if the :func:`ioloop.IOLoop.start` method should be called too. """ header = self.stream.stream_header() self.send_raw_stanza(header) if start_loop: ioloop.IOLoop.instance().start() def handle_read(self, data): try: self.parser.feed(data) except SAXParseException, exc: self.log(traceback=True) self._read()
class IcecastSourceClient(object): BUFFER_TIME = 3.0 icecast_source_client_set = set([]) def __init__(self, stream_id, kbps, audiostream_connection): self.stream_id = stream_id self.connection = audiostream_connection self.didStart = False self.isFinishing = False self.kbps = kbps self.queue = Queue.Queue() IcecastSourceClient.icecast_source_client_set.add(self) s = socket.socket(socket.AF_INET, socket.SOCK_STREAM, 0) self.stream = IOStream(s) self.stream.set_close_callback(self._on_close) self.stream.connect(("localhost", 8000), self.connect) self.curr_queue_time = 0.0 self.periodic = PeriodicCallback(self.manage_audio, self.bytes2time(1000*AudioStreamConnection.BYTES_PER_READ), IOLoop.instance()) @staticmethod def getIcecastSourceClient(stream_id, kbps, audiostream_connection): for source_client in IcecastSourceClient.icecast_source_client_set: if source_client.stream_id == stream_id: source_client.isFinishing = False return source_client return IcecastSourceClient(stream_id, kbps, audiostream_connection) def connect(self): logging.info('%s-icecast port connected', self.stream_id) self.stream.write(("SOURCE /%s HTTP/1.0\n" "Authorization: Basic c291cmNlOnRlc3RpbmcjIyNzcGFjZWJhcg==\n" "User-Agent: libshout/2.3.1\n" "Content-Type: audio/mpeg\n" "ice-description: HLS Test\n\n") % self.stream_id) self.stream.read_until("\n", self.on_response) self.periodic.start() def add_audio(self, data): self.queue.put(data) self.curr_queue_time += self.bytes2time(len(data)) logging.info('%s-received,BUFFER:%f', self.stream_id, self.curr_queue_time) def manage_audio(self): if self.stream.closed(): self.connection.stream.close() self.periodic.stop() return if not self.didStart: if self.curr_queue_time > IcecastSourceClient.BUFFER_TIME: self.didStart = True else: return if not self.isFinishing and self.curr_queue_time < 1.0: logging.info('%s - Halt sending to refill buffer', self.stream_id) return if self.queue.empty() and self.isFinishing: logging.info('%s - sent all data' % self.stream_id) self.periodic.stop() self.stream.close() else: data = self.queue.get() self.curr_queue_time -= self.bytes2time(len(data)) logging.info('%s-sending, BUFFER:%f', self.stream_id, self.curr_queue_time) self.stream.write(data) def bytes2time(self, num_bytes): return 8*num_bytes/(self.kbps*1024.0) def on_response(self, data): logging.info('%s-Received from icecast: %s', self.stream_id, data) def _on_close(self): logging.info('%s-closed icecast stream', self.stream_id) IcecastSourceClient.icecast_source_client_set.remove(self)
class Connection(object): def __init__(self, host='localhost', port=11211, pool=None): self._host = host self._port = port self._pool = pool self._socket = None self._stream = None self._ioloop = IOLoop.instance() self.connect() def connect(self): try: self._socket = socket(AF_INET, SOCK_STREAM, 0) self._socket.connect((self._host, self._port)) self._stream = IOStream(self._socket, io_loop=self._ioloop) self._stream.set_close_callback(self.on_disconnect) except error as e: raise ConnectionError(e) def disconect(self): callback = self._final_callback self._final_callback = None try: if callback: callback(None) finally: self._stream._close_callback = None self._stream.close() def on_disconnect(self): callback = self._final_callback self._final_callback = None try: if callback: callback(None) finally: logging.debug('asyncmemcached closing connection') self._pool.release(self) def closed(self): return self._stream.closed() def send_command(self, fullcmd, expect_str, callback): self._final_callback = callback if self._stream.closed(): self.connect() with stack_context.StackContext(self.cleanup): if fullcmd[0:3] == 'get' or \ fullcmd[0:4] == 'incr' or \ fullcmd[0:4] == 'decr': self._stream.write(fullcmd, self.read_value) else: self._stream.write(fullcmd, functools.partial(self.read_response, expect_str)) def read_response(self, expect_str): self._stream.read_until('\r\n', functools.partial(self._expect_callback, expect_str)) def read_value(self): self._stream.read_until('\r\n', self._expect_value_header_callback) def _expect_value_header_callback(self, response): response = response[:-2] if response[:5] == 'VALUE': resp, key, flag, length = response.split() flag = int(flag) length = int(length) self._stream.read_bytes(length+2, self._expect_value_callback) elif response.isdigit(): try: callback = self._final_callback self._final_callback = None if callback: callback(int(response)) finally: self._pool.release(self) else: try: callback = self._final_callback self._final_callback = None if callback: callback(None) finally: self._pool.release(self) def _expect_value_callback(self, value): value = value[:-2] self._stream.read_until('\r\n', functools.partial(self._end_value_callback, value)) def _end_value_callback(self, value, response): response = response.rstrip('\r\n') if response == 'END': try: callback = self._final_callback self._final_callback = None if callback: callback(value) finally: self._pool.release(self) else: raise RedisError('error %s' % response) def _expect_callback(self, expect_str, response): response = response.rstrip('\r\n') if response == expect_str: try: callback = self._final_callback self._final_callback = None if callback: callback(None) finally: self._pool.release(self) else: raise RedisError('error %s' % response) @contextlib.contextmanager def cleanup(self): try: yield except Exception as e: logging.warning("uncaught exception", exc_info=True) try: callback = self._final_callback self._final_callback = None if callback: callback(None) finally: self._pool.release(self)
class IPCClient(object): ''' A Tornado IPC client very similar to Tornado's TCPClient class but using either UNIX domain sockets or TCP sockets This was written because Tornado does not have its own IPC server/client implementation. :param IOLoop io_loop: A Tornado ioloop to handle scheduling :param str/int socket_path: A path on the filesystem where a socket belonging to a running IPCServer can be found. It may also be of type 'int', in which case it is used as the port for a tcp localhost connection. ''' # Create singleton map between two sockets instance_map = weakref.WeakKeyDictionary() def __new__(cls, socket_path, io_loop=None): io_loop = io_loop or tornado.ioloop.IOLoop.current() if io_loop not in IPCClient.instance_map: IPCClient.instance_map[io_loop] = weakref.WeakValueDictionary() loop_instance_map = IPCClient.instance_map[io_loop] # FIXME key = six.text_type(socket_path) client = loop_instance_map.get(key) if client is None: log.debug('Initializing new IPCClient for path: %s', key) client = object.__new__(cls) # FIXME client.__singleton_init__(io_loop=io_loop, socket_path=socket_path) client._instance_key = key loop_instance_map[key] = client client._refcount = 1 client._refcount_lock = threading.RLock() else: log.debug('Re-using IPCClient for %s', key) with client._refcount_lock: client._refcount += 1 return client def __singleton_init__(self, socket_path, io_loop=None): ''' Create a new IPC client IPC clients cannot bind to ports, but must connect to existing IPC servers. Clients can then send messages to the server. ''' self.io_loop = io_loop or tornado.ioloop.IOLoop.current() self.socket_path = socket_path self._closing = False self.stream = None if six.PY2: encoding = None else: encoding = 'utf-8' self.unpacker = msgpack.Unpacker(encoding=encoding) def __init__(self, socket_path, io_loop=None): # Handled by singleton __new__ pass def connected(self): return self.stream is not None and not self.stream.closed() def connect(self, callback=None, timeout=None): ''' Connect to the IPC socket ''' if hasattr( self, '_connecting_future') and not self._connecting_future.done(): # pylint: disable=E0203 future = self._connecting_future # pylint: disable=E0203 else: if hasattr(self, '_connecting_future'): # read previous future result to prevent the "unhandled future exception" error self._connecting_future.exception() # pylint: disable=E0203 future = tornado.concurrent.Future() self._connecting_future = future self._connect(timeout=timeout) if callback is not None: def handle_future(future): response = future.result() self.io_loop.add_callback(callback, response) future.add_done_callback(handle_future) return future @tornado.gen.coroutine def _connect(self, timeout=None): ''' Connect to a running IPCServer ''' if isinstance(self.socket_path, int): sock_type = socket.AF_INET sock_addr = ('127.0.0.1', self.socket_path) else: sock_type = socket.AF_UNIX sock_addr = self.socket_path self.stream = None if timeout is not None: timeout_at = time.time() + timeout while True: if self._closing: break if self.stream is None: with salt.utils.asynchronous.current_ioloop(self.io_loop): self.stream = IOStream( socket.socket(sock_type, socket.SOCK_STREAM), ) try: log.trace('IPCClient: Connecting to socket: %s', self.socket_path) yield self.stream.connect(sock_addr) self._connecting_future.set_result(True) break except Exception as e: if self.stream.closed(): self.stream = None if timeout is None or time.time() > timeout_at: if self.stream is not None: self.stream.close() self.stream = None self._connecting_future.set_exception(e) break yield tornado.gen.sleep(1) def __del__(self): try: with self._refcount_lock: # Make sure we actually close no matter if something # went wrong with our ref counting self._refcount = 1 try: self.close() except socket.error as exc: if exc.errno != errno.EBADF: # If its not a bad file descriptor error, raise raise except TypeError: # This is raised when Python's GC has collected objects which # would be needed when calling self.close() pass def close(self): ''' Routines to handle any cleanup before the instance shuts down. Sockets and filehandles should be closed explicitly, to prevent leaks. ''' if self._closing: return if self._refcount > 1: # Decrease refcount with self._refcount_lock: self._refcount -= 1 log.debug('This is not the last %s instance. Not closing yet.', self.__class__.__name__) return self._closing = True log.debug('Closing %s instance', self.__class__.__name__) if self.stream is not None and not self.stream.closed(): self.stream.close() # Remove the entry from the instance map so # that a closed entry may not be reused. # This forces this operation even if the reference # count of the entry has not yet gone to zero. if self.io_loop in self.__class__.instance_map: loop_instance_map = self.__class__.instance_map[self.io_loop] if self._instance_key in loop_instance_map: del loop_instance_map[self._instance_key] if not loop_instance_map: del self.__class__.instance_map[self.io_loop]
class TornadoClient(Client) : """A non-blocking Pomelo client by tornado ioloop Usage : class ClientHandler(object) : def on_recv_data(self, client, proto_type, data) : print "recv_data..." return data def on_connected(self, client, user_data) : print "connect..." client.send_heartbeat() def on_disconnect(self, client) : print "disconnect..." def on_heartbeat(self, client) : print "heartbeat..." send request ... def on_response(self, client, route, request, response) : print "response..." def on_push(self, client, route, push_data) : print "notify..." handler = ClientHandler() client = TornadoClient(handler) client.connect(host, int(port)) client.run() tornado.ioloop.IOLoop.current().start() """ def __init__(self, handler) : self.socket = socket(AF_INET, SOCK_STREAM) self.iostream = None self.protocol_package = None super(TornadoClient, self).__init__(handler) def connect(self, host, port) : self.iostream = IOStream(self.socket) self.iostream.set_close_callback(self.on_close) self.iostream.connect((host, port), self.on_connect) def on_connect(self) : self.send_sync() self.on_data() def on_close(self) : if hasattr(self.handler, 'on_disconnect') : self.handler.on_disconnect(self) def send(self, data) : assert not self.iostream.closed(), "iostream has closed" if not isinstance(data, bytes) : data = bytes(data) self.iostream.write(data) def on_data(self) : assert not self.iostream.closed(), "iostream has closed" if None is self.protocol_package or self.protocol_package.completed() : self.iostream.read_bytes(4, self.on_head) def on_head(self, head) : self.protocol_package = Protocol.unpack(head) self.iostream.read_bytes(self.protocol_package.length, self.on_body) def on_body(self, body) : if hasattr(self.handler, 'on_recv_data') : body = self.handler.on_recv_data(self, self.protocol_package.proto_type, body) self.protocol_package.append(body) self.on_protocol(self.protocol_package) self.on_data() def close(self) : if self.iostream : self.iostream.close()