67 lines
2.7 KiB
Python
67 lines
2.7 KiB
Python
|
#!/usr/bin/env python3
|
||
|
|
"""TCP forwarding with trusted PROXY v1 metadata, before any TLS bytes."""
|
||
|
|
import argparse
|
||
|
|
import asyncio
|
||
|
|
import ipaddress
|
||
|
|
import os
|
||
|
|
|
||
|
|
from env_file import load, env_int
|
||
|
|
|
||
|
|
|
||
|
|
async def forward(reader, writer, target_host, target_port):
|
||
|
|
upstream = None
|
||
|
|
try:
|
||
|
|
peer = writer.get_extra_info('peername')
|
||
|
|
local = writer.get_extra_info('sockname')
|
||
|
|
family = 'TCP6' if ipaddress.ip_address(peer[0]).version == 6 else 'TCP4'
|
||
|
|
remote, upstream = await asyncio.wait_for(asyncio.open_connection(target_host, target_port), 5)
|
||
|
|
upstream.write(f'PROXY {family} {peer[0]} {local[0]} {peer[1]} {local[1]}\r\n'.encode('ascii'))
|
||
|
|
await upstream.drain()
|
||
|
|
async def copy(source, destination):
|
||
|
|
while data := await asyncio.wait_for(source.read(65536), 60):
|
||
|
|
destination.write(data)
|
||
|
|
await asyncio.wait_for(destination.drain(), 10)
|
||
|
|
if destination.can_write_eof():
|
||
|
|
destination.write_eof()
|
||
|
|
await asyncio.wait_for(asyncio.gather(copy(reader, upstream), copy(remote, writer)), 120)
|
||
|
|
except (OSError, asyncio.TimeoutError):
|
||
|
|
pass
|
||
|
|
finally:
|
||
|
|
if upstream:
|
||
|
|
upstream.close()
|
||
|
|
writer.close()
|
||
|
|
|
||
|
|
|
||
|
|
async def serve(args):
|
||
|
|
active = set()
|
||
|
|
allowed = [ipaddress.ip_network(item.strip(), strict=False)
|
||
|
|
for item in (args.allow_from or '').split(',') if item.strip()]
|
||
|
|
|
||
|
|
def permitted(peer):
|
||
|
|
return not allowed or any(ipaddress.ip_address(peer[0]) in network for network in allowed)
|
||
|
|
|
||
|
|
async def accept(reader, writer):
|
||
|
|
if len(active) >= 64 or not permitted(writer.get_extra_info('peername')):
|
||
|
|
writer.close()
|
||
|
|
return
|
||
|
|
task = asyncio.current_task()
|
||
|
|
active.add(task)
|
||
|
|
try:
|
||
|
|
await forward(reader, writer, args.target, args.target_port)
|
||
|
|
finally:
|
||
|
|
active.discard(task)
|
||
|
|
server = await asyncio.start_server(accept, args.bind, args.port, limit=65536)
|
||
|
|
async with server:
|
||
|
|
await server.serve_forever()
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
load()
|
||
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
||
|
|
parser.add_argument('--bind', default=os.environ.get('INGRESS_BIND', '127.0.0.1'))
|
||
|
|
parser.add_argument('--port', type=int, default=env_int('INGRESS_PORT', 18443))
|
||
|
|
parser.add_argument('--target', default=os.environ.get('INGRESS_TARGET', '127.0.0.1'))
|
||
|
|
parser.add_argument('--target-port', type=int, default=env_int('INGRESS_TARGET_PORT', 19443))
|
||
|
|
parser.add_argument('--allow-from', default=os.environ.get('INGRESS_ALLOW_FROM', ''),
|
||
|
|
help='comma-separated CIDRs allowed to connect (empty = any)')
|
||
|
|
asyncio.run(serve(parser.parse_args()))
|