commit 7b0de170a79d6da1935598c17eeb65a1e4db9a05 Author: Tom Brennan Date: Thu Apr 3 08:05:44 2025 -0400 initial implementation diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..0e5ac79 --- /dev/null +++ b/.gitignore @@ -0,0 +1,2 @@ +.venv +__pycache__ \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..b2ee5ba --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,21 @@ +[build-system] +requires = ["hatchling >= 1.26"] +build-backend = "hatchling.build" + +[project] +name = "rproxy" +dynamic = ["version"] + +dependencies = [ + "pyyaml", + "pydantic", +] + +[project.scripts] +rproxy = "rproxy.__main__:start_reverse_proxy" + +[tool.hatch.version] +path = "rproxy/__about__.py" + +# [tool.hatch.build.targets.wheel] +# only-include = ["rproxy"] diff --git a/rproxy/__about__.py b/rproxy/__about__.py new file mode 100644 index 0000000..3dc1f76 --- /dev/null +++ b/rproxy/__about__.py @@ -0,0 +1 @@ +__version__ = "0.1.0" diff --git a/rproxy/__init__.py b/rproxy/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/rproxy/__main__.py b/rproxy/__main__.py new file mode 100755 index 0000000..82c8e9d --- /dev/null +++ b/rproxy/__main__.py @@ -0,0 +1,42 @@ +#!/usr/bin/env python +import importlib +import importlib.util +import os +from pathlib import Path +import sys + +import yaml + +from rproxy.sockets import start_reverse_proxy +from rproxy.config import Config + + +def usage() -> str: + example = ( + Path(importlib.util.find_spec("rproxy").origin).parent + / "resources" + / "example.config.yaml" + ) + return ( + f"Reverse proxy server.\n\nusage: {prog} config\n" + f"example: {prog} ./my-config.yaml\n\n" + f"See example config file at {example}" + ) + + +try: + with open(sys.argv[1]) as f: + config = Config(**yaml.safe_load(f)) + print(config) +except (FileNotFoundError, IndexError) as e: + prog = os.path.basename(sys.argv[0]) + print( + f"Error: missing config file argument. {e}.\n{usage()}", + file=sys.stderr, + ) + sys.exit(1) + +try: + start_reverse_proxy(config) +except KeyboardInterrupt: + sys.exit(0) diff --git a/rproxy/config.py b/rproxy/config.py new file mode 100644 index 0000000..34b2bc0 --- /dev/null +++ b/rproxy/config.py @@ -0,0 +1,19 @@ +import pydantic + + +class SocketAddr(pydantic.BaseModel): + host: str + port: int = 8000 + + +class LocationMatch(pydantic.BaseModel): + request_line: str = "" + + +class Location(SocketAddr): + match: LocationMatch = LocationMatch() + + +class Config(pydantic.BaseModel): + listen: SocketAddr + locations: list[Location] diff --git a/rproxy/resources/example.config.yaml b/rproxy/resources/example.config.yaml new file mode 100644 index 0000000..833ee68 --- /dev/null +++ b/rproxy/resources/example.config.yaml @@ -0,0 +1,14 @@ +listen: + host: &local localhost + port: 8001 + +locations: +- match: + request_line: ^GET /api.*\r\n + # headers: + # foo: bar + host: *local + port: 8081 + +- host: *local + port: 8089 diff --git a/rproxy/sockets.py b/rproxy/sockets.py new file mode 100644 index 0000000..baff876 --- /dev/null +++ b/rproxy/sockets.py @@ -0,0 +1,96 @@ +import re +import socket +import threading + +from rproxy.config import Config + + +# Function to forward data from one socket to another +def forward_data(source: socket.socket, destination: socket.socket): + while True: + data = source.recv(4 << 10) + if not data: + break + destination.sendall(data) + + +def handle_client( + request_line: bytes, + client_socket: socket.socket, + target_host: str, + target_port: int, +): + try: + # Connect to the target server + server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + server_socket.connect((target_host, target_port)) + + # We do this just once, and then assume any other requests to this + # connection will be for the selected backend. + server_socket.sendall(request_line) + + # Create threads to forward data in both directions + client_to_server = threading.Thread( + target=forward_data, args=(client_socket, server_socket) + ) + server_to_client = threading.Thread( + target=forward_data, args=(server_socket, client_socket) + ) + + # Start the threads + client_to_server.start() + server_to_client.start() + + # Wait for the threads to complete + client_to_server.join() + server_to_client.join() + + except Exception as e: + print(f"Error: {e}") + finally: + client_socket.close() + server_socket.close() + + +def read_request_line(socket: socket.socket): + data = b"" + while not data.endswith(b"\r\n"): + data += socket.recv(1) + return data + + +def start_reverse_proxy(config: Config): + # Create a socket + server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + + # Bind the socket to the specified address and port + server_socket.bind((config.listen.host, config.listen.port)) + + # Listen for incoming connections + server_socket.listen(5) + print(f"Reverse proxy listening on {config.listen.host}:{config.listen.port}") + + while True: + # Accept a connection + client_socket, client_address = server_socket.accept() + print(f"Accepted connection from {client_address[0]}:{client_address[1]}") + + # Handle the client connection in a new thread + request_line = read_request_line(client_socket) + target_host, target_port = None, None + for location in config.locations: + if re.match(location.match.request_line, request_line.decode()): + target_host = location.host + target_port = location.port + break + + if not (target_host and target_port): + print(f'Missing location match for request line "{request_line}"') + client_socket.close() + continue + + client_thread = threading.Thread( + target=handle_client, + args=(request_line, client_socket, target_host, target_port), + ) + client_thread.start()