ntunnel/ntunnel/quic.py

# Copyright 2022 John-Mark Gurney.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
# 1. Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS ``AS IS'' AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
# OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
# HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
# OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
# SUCH DAMAGE.
#

import asyncio
import unittest

from . import parsesockstr, connectsockstr, listensockstr

from aioquic.asyncio import QuicConnectionProtocol, serve
from aioquic.asyncio.client import connect
from aioquic.quic.configuration import QuicConfiguration

# copied from wsfwd
async def fwd_data(reader, writer):
	while True:
		data = await reader.read(16384)
		if data == b'':
			#_debprint('fwd_data eof', repr(reader), repr(writer))
			# XXX - aioquic doesn't implement close
			#writer.close()
			#await writer.wait_closed()
			#_debprint('fwd_data done', repr(reader), repr(writer))
			return

		#_debprint('fwd_data data', repr(reader), repr(writer), len(data))
		writer.write(data)
		# XXX - aioquic doesn't implement is_closing
		#await writer.drain()

async def run_connect(dst, rdr, wrr):
	connrdr, connwrr = await connectsockstr(dst)

	await asyncio.gather(fwd_data(connrdr, wrr), fwd_data(rdr, connwrr))

def cmd_quic_serv(args):
	privkey = args.servkey
	cert = args.cert

	quic_logger = None

	quic_conf = QuicConfiguration(
	    alpn_protocols=["ntunnel-01"],
	    is_client=False,
	    quic_logger=quic_logger,
	)

	quic_conf.load_cert_chain(cert, privkey)

	proto, slargs = parsesockstr(args.servlisten)

	if proto != 'udp':
		raise ValueError('protocol for servlisten must be udp')

	def sh(rdr, wrr):
		task = run_connect(args.servtarget, rdr, wrr)
		asyncio.create_task(task)

		# XXX - await task

	print('foo', repr(slargs))

	loop = asyncio.get_event_loop()
	loop.run_until_complete(serve(slargs['host'], slargs['port'],
	    configuration=quic_conf, retry=True, stream_handler=sh))

	try:
		loop.run_forever()
	except KeyboardInterrupt:
		pass

async def client_run(conf, liststr, deststr):
	proto, slargs = parsesockstr(deststr)

	if proto != 'udp':
		raise ValueError('protocol for destination must be udp')

	# XXX - loop when server restarts?
	async with connect(slargs['host'], slargs['port'], configuration=conf) as \
	    client:
		async def connmaker(rdr, wrr):
			connrdr, connwrr = await client.create_stream()

			await asyncio.gather(fwd_data(connrdr, wrr),
			    fwd_data(rdr, connwrr))

		ssock = await listensockstr(liststr, connmaker)

		# XXX - how to break out when new connection needed?
		await ssock.serve_forever()

def cmd_quic_client(args):
	quic_logger = None

	quic_conf = QuicConfiguration(
	    alpn_protocols=["ntunnel-01"],
	    is_client=True,
	    quic_logger=quic_logger,
	)

	if args.ca_certs:
		quic_conf.load_verify_locations(args.ca_certs)

	cr = client_run(quic_conf, args.clientlisten, args.clienttarget)

	loop = asyncio.get_event_loop()

	loop.run_until_complete(cr)


def quic_parsers(subparsers):
	parser_quic_serv = subparsers.add_parser('quic_serv', help='run a QUIC server')
	parser_quic_serv.add_argument("-k", "--servkey", type=str,
	    help="load the TLS private key from the specified file")
	parser_quic_serv.add_argument("-c", "--cert", type=str,
	    required=True,
	    help="load the TLS certificate from the specified file")
	parser_quic_serv.add_argument('servlisten', type=str, help='Connection that the server listens on')
	parser_quic_serv.add_argument('servtarget', type=str, help='Connection that the server connects to')
	parser_quic_serv.set_defaults(func=cmd_quic_serv)

	parser_quic_client = subparsers.add_parser('quic_client', help='run a QUIC client')
	parser_quic_client.add_argument("--ca-certs", type=str,
	    help="load CA certificates from the specified file")
	parser_quic_client.add_argument('clientlisten', type=str, help='Connection that the client listens on')
	parser_quic_client.add_argument('clienttarget', type=str, help='Connection that the client connects to')
	parser_quic_client.set_defaults(func=cmd_quic_client)

class Tests(unittest.IsolatedAsyncioTestCase):
	pass