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