From d1745d6749c3ec72a313b6e99c151b351dede37e Mon Sep 17 00:00:00 2001 From: kgrofelnik Date: Wed, 29 Jan 2025 17:22:19 +0100 Subject: [PATCH 1/2] WIP --- verified-inference/host/dns_forwarder.py | 56 ++++++++++++ verified-inference/host/traffic_forwarder.py | 95 ++++++++++++++++++++ 2 files changed, 151 insertions(+) create mode 100644 verified-inference/host/dns_forwarder.py create mode 100644 verified-inference/host/traffic_forwarder.py diff --git a/verified-inference/host/dns_forwarder.py b/verified-inference/host/dns_forwarder.py new file mode 100644 index 0000000..3e7e812 --- /dev/null +++ b/verified-inference/host/dns_forwarder.py @@ -0,0 +1,56 @@ +import socket + +VSOCK_PORT = 5053 + + +def get_system_dns(): + """Retrieve the first nameserver from /etc/resolv.conf.""" + try: + with open("/etc/resolv.conf", "r") as f: + for line in f: + if line.startswith("nameserver"): + return line.split()[1], 53 # Extract the IP and set port 53 + except Exception as e: + print(f"Failed to get system DNS: {e}") + + return "8.8.8.8", 53 # Fallback to Google DNS if none found + + +def forward_dns_to_public(dns_request, dns_ip, dns_port): + """Forward the DNS request to a public DNS resolver.""" + with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as dns_sock: + dns_sock.sendto( + dns_request, (dns_ip, dns_port) + ) # Send request to the DNS resolver + response, _ = dns_sock.recvfrom(512) # Receive response from the DNS resolver + return response + + +def vsock_dns_proxy(): + """Listen on a vsock port and forward DNS requests to a public DNS server.""" + server_socket = socket.socket(socket.AF_VSOCK, socket.SOCK_STREAM) + server_socket.bind((socket.VMADDR_CID_ANY, VSOCK_PORT)) + server_socket.listen(5) + print(f"Listening for DNS requests on vsock port {VSOCK_PORT}...") + dns_ip, dns_port = get_system_dns() + print(f"Using DNS server {dns_ip}:{dns_port}") + + while True: + client_socket, _ = server_socket.accept() + print("Received DNS request from enclave.") + + try: + # Receive the DNS request + dns_request = client_socket.recv(512) + dns_response = forward_dns_to_public(dns_request, dns_ip, dns_port) + + # Send response back to enclave + client_socket.sendall(dns_response) + except Exception as e: + print(f"Error: {e}") + finally: + client_socket.close() + + +if __name__ == "__main__": + vsock_dns_proxy() diff --git a/verified-inference/host/traffic_forwarder.py b/verified-inference/host/traffic_forwarder.py new file mode 100644 index 0000000..1be0e6d --- /dev/null +++ b/verified-inference/host/traffic_forwarder.py @@ -0,0 +1,95 @@ +import socket +import struct +import threading +import logging + +# Configure logging +logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(message)s") +logger = logging.getLogger() + + +class HostProxy: + BUFFER_SIZE = 1024 + CID = 3 # The CID of the enclave + PORT = 8001 # The port where this proxy listens for traffic from the enclave + + def forward(self, source, destination): + """Forward data between two sockets.""" + try: + while True: + data = source.recv(self.BUFFER_SIZE) + if not data: + break + destination.sendall(data) + except Exception as exc: + logger.error(f"Exception in forward: {exc}") + finally: + source.close() + destination.close() + + def handle_vsock_connection(self, vsock_client): + """ + Handles a new connection from the enclave over vsock, + extracts the original destination, and forwards traffic. + """ + try: + # Read first 6 bytes which contain the prepended IP and port + header = vsock_client.recv(6) + if len(header) < 6: + logger.error("Error: Incomplete header received") + vsock_client.close() + return + + # Extract original IP and port + original_ip_bin, original_port_bin = header[:4], header[4:6] + original_ip = socket.inet_ntoa( + original_ip_bin + ) # Convert binary IP to string + original_port = struct.unpack("!H", original_port_bin)[ + 0 + ] # Convert binary port to integer + + logger.info(f"Forwarding traffic to {original_ip}:{original_port}") + + # Connect to the real destination + destination_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + destination_socket.connect((original_ip, original_port)) + + # Start forwarding traffic in both directions + outgoing_thread = threading.Thread( + target=self.forward, args=(vsock_client, destination_socket) + ) + incoming_thread = threading.Thread( + target=self.forward, args=(destination_socket, vsock_client) + ) + + outgoing_thread.start() + incoming_thread.start() + except Exception as exc: + logger.error(f"Error handling vsock connection: {exc}") + vsock_client.close() + + def start(self): + """Starts the host-side proxy listening on vsock.""" + try: + vsock_server = socket.socket(socket.AF_VSOCK, socket.SOCK_STREAM) + vsock_server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + vsock_server.bind((self.CID, self.PORT)) + vsock_server.listen(10) + + logger.info( + f"Host proxy listening on vsock {self.vsock_cid}:{self.vsock_port}" + ) + + while True: + client_socket, _ = vsock_server.accept() + threading.Thread( + target=self.handle_vsock_connection, args=(client_socket,) + ).start() + except Exception as exc: + logger.error(f"HostProxy exception: {exc}") + + +if __name__ == "__main__": + proxy = HostProxy() + proxy.start() From f8f5fb3071c430b43705d56d1d89ee1d897ae2a8 Mon Sep 17 00:00:00 2001 From: kgrofelnik Date: Tue, 11 Feb 2025 09:26:00 +0100 Subject: [PATCH 2/2] DNS and traffic forwarders --- verified-inference/host/dns_forwarder.py | 15 ++++++++++----- verified-inference/host/run_forwarders.sh | 10 ++++++++++ verified-inference/host/stop_forwarders.sh | 19 +++++++++++++++++++ verified-inference/host/traffic_forwarder.py | 4 +--- 4 files changed, 40 insertions(+), 8 deletions(-) create mode 100755 verified-inference/host/run_forwarders.sh create mode 100755 verified-inference/host/stop_forwarders.sh diff --git a/verified-inference/host/dns_forwarder.py b/verified-inference/host/dns_forwarder.py index 3e7e812..6b98916 100644 --- a/verified-inference/host/dns_forwarder.py +++ b/verified-inference/host/dns_forwarder.py @@ -1,7 +1,12 @@ import socket +import logging VSOCK_PORT = 5053 +# Configure logging +logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(message)s") +logger = logging.getLogger() + def get_system_dns(): """Retrieve the first nameserver from /etc/resolv.conf.""" @@ -11,7 +16,7 @@ def get_system_dns(): if line.startswith("nameserver"): return line.split()[1], 53 # Extract the IP and set port 53 except Exception as e: - print(f"Failed to get system DNS: {e}") + logger.error(f"Failed to get system DNS: {e}") return "8.8.8.8", 53 # Fallback to Google DNS if none found @@ -31,13 +36,13 @@ def vsock_dns_proxy(): server_socket = socket.socket(socket.AF_VSOCK, socket.SOCK_STREAM) server_socket.bind((socket.VMADDR_CID_ANY, VSOCK_PORT)) server_socket.listen(5) - print(f"Listening for DNS requests on vsock port {VSOCK_PORT}...") + logger.info(f"Listening for DNS requests on vsock port {VSOCK_PORT}...") dns_ip, dns_port = get_system_dns() - print(f"Using DNS server {dns_ip}:{dns_port}") + logger.info(f"Using DNS server {dns_ip}:{dns_port}") while True: client_socket, _ = server_socket.accept() - print("Received DNS request from enclave.") + logger.info("Received DNS request from enclave.") try: # Receive the DNS request @@ -47,7 +52,7 @@ def vsock_dns_proxy(): # Send response back to enclave client_socket.sendall(dns_response) except Exception as e: - print(f"Error: {e}") + logger.error(f"Error: {e}") finally: client_socket.close() diff --git a/verified-inference/host/run_forwarders.sh b/verified-inference/host/run_forwarders.sh new file mode 100755 index 0000000..70fdb1d --- /dev/null +++ b/verified-inference/host/run_forwarders.sh @@ -0,0 +1,10 @@ +#!/usr/bin/env bash +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +echo "Starting DNS forwarder..." +nohup python "$SCRIPT_DIR/dns_forwarder.py" > "$SCRIPT_DIR/dns_forwarder.log" 2>&1 & +echo "Started DNS forwarder with PID $!" + +echo "Starting traffic forwarder..." +nohup python "$SCRIPT_DIR/traffic_forwarder.py" > "$SCRIPT_DIR/traffic_forwarder.log" 2>&1 & +echo "Started traffic forwarder with PID $!" \ No newline at end of file diff --git a/verified-inference/host/stop_forwarders.sh b/verified-inference/host/stop_forwarders.sh new file mode 100755 index 0000000..caa873a --- /dev/null +++ b/verified-inference/host/stop_forwarders.sh @@ -0,0 +1,19 @@ +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +PIDS=$(pgrep -f "$SCRIPT_DIR/dns_forwarder.py") + +if [ -n "$PIDS" ]; then + echo "Killing existing instance(s) of DNS forwarder..." + kill $PIDS + sleep 1 +fi + +PIDS=$(pgrep -f "$SCRIPT_DIR/traffic_forwarder.py") + +if [ -n "$PIDS" ]; then + echo "Killing existing instance(s) of traffic forwarder..." + kill $PIDS + sleep 1 +fi + +echo "All forwarders have been stopped." diff --git a/verified-inference/host/traffic_forwarder.py b/verified-inference/host/traffic_forwarder.py index 1be0e6d..4ad1b6b 100644 --- a/verified-inference/host/traffic_forwarder.py +++ b/verified-inference/host/traffic_forwarder.py @@ -77,9 +77,7 @@ def start(self): vsock_server.bind((self.CID, self.PORT)) vsock_server.listen(10) - logger.info( - f"Host proxy listening on vsock {self.vsock_cid}:{self.vsock_port}" - ) + logger.info(f"Host proxy listening on vsock {self.CID}:{self.PORT}") while True: client_socket, _ = vsock_server.accept()