def start_tcp_proxy(src, dst, handler, **kwargs): """Run a simple TCP proxy (tunneling raw connections from src to dst), using a message handler that can be used to intercept messages and return predefined responses for certain requests. Arguments: src -- Source IP address and port string. I.e.: '127.0.0.1:8000' dst -- Destination IP address and port. I.e.: '127.0.0.1:8888' handler -- a handler function to intercept requests (returns tuple (forward_value, response_value)) """ src = "%s:%s" % (BIND_HOST, src) if is_number(src) else src dst = "%s:%s" % (LOCALHOST_IP, dst) if is_number(dst) else dst thread = kwargs.get("_thread") def ip_to_tuple(ip): ip, port = ip.split(":") return ip, int(port) s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) s.bind(ip_to_tuple(src)) s.listen(1) s.settimeout(10) def handle_request(s_src, thread): s_dst = socket.socket(socket.AF_INET, socket.SOCK_STREAM) s_dst.connect(ip_to_tuple(dst)) sockets = [s_src, s_dst] try: while thread.running: s_read, _, _ = select.select(sockets, [], []) for s in s_read: data = s.recv(BUFFER_SIZE) if data in [b"", "", None]: return if s == s_src: forward, response = data, None if handler: forward, response = handler(data) if forward is not None: s_dst.sendall(forward) elif response is not None: s_src.sendall(response) return elif s == s_dst: s_src.sendall(data) finally: run_safe(s_src.close) run_safe(s_dst.close) while thread.running: try: src_socket, _ = s.accept() start_worker_thread( lambda *args, _thread: handle_request(src_socket, _thread)) except socket.timeout: pass
def on_before_start(self): # re-schedule alarms for persistence use-case def restart_alarms(*args): poll_condition(lambda: SERVICE_PLUGINS.is_running("cloudwatch")) self.alarm_scheduler.restart_existing_alarms() start_worker_thread(restart_alarms)
def _start(self): with self._startup_mutex: if self._started: return self._started = True # startup has to run async, otherwise first call to handle() could block a long time. start_worker_thread(self._do_start_retry)
def start_container( self, container_name_or_id: str, stdin=None, interactive: bool = False, attach: bool = False, flags: Optional[str] = None, ) -> Tuple[bytes, bytes]: LOG.debug("Starting container %s", container_name_or_id) try: container = self.client().containers.get(container_name_or_id) stdout = to_bytes(container_name_or_id) stderr = b"" if interactive or attach: params = {"stdout": 1, "stderr": 1, "stream": 1} if interactive: params["stdin"] = 1 sock = container.attach_socket(params=params) sock = sock._sock if hasattr(sock, "_sock") else sock result_queue = queue.Queue() thread_started = threading.Event() start_waiting = threading.Event() # Note: We need to be careful about potential race conditions here - .wait() should happen right # after .start(). Hence starting a thread and asynchronously waiting for the container exit code def wait_for_result(*_): _exit_code = -1 try: thread_started.set() start_waiting.wait() _exit_code = container.wait()["StatusCode"] except APIError as e: _exit_code = 1 raise ContainerException(str(e)) finally: result_queue.put(_exit_code) # start listener thread start_worker_thread(wait_for_result) thread_started.wait() # start container container.start() # start awaiting container result start_waiting.set() # handle container input/output # under windows, the socket has no __enter__ / cannot be used as context manager # therefore try/finally instead of with here try: if stdin: sock.sendall(to_bytes(stdin)) sock.shutdown(socket.SHUT_WR) stdout, stderr = self._read_from_sock(sock, False) except socket.timeout: LOG.debug( f"Socket timeout when talking to the I/O streams of Docker container '{container_name_or_id}'" ) finally: sock.close() # get container exit code exit_code = result_queue.get() if exit_code: raise ContainerException( f"Docker container returned with exit code {exit_code}", stdout=stdout, stderr=stderr, ) else: container.start() return stdout, stderr except NotFound: raise NoSuchContainer(container_name_or_id) except APIError as e: raise ContainerException() from e