| #!/usr/bin/env python3 |
| import argparse |
| import asyncio |
| import io |
| import logging |
| import os |
| import re |
| import struct |
| import urllib.parse |
| from typing import Dict, Optional |
| |
| from aioquic.asyncio import QuicConnectionProtocol, serve |
| from aioquic.quic.configuration import QuicConfiguration |
| from aioquic.quic.connection import END_STATES |
| from aioquic.quic.events import StreamDataReceived, QuicEvent |
| from aioquic.tls import SessionTicket |
| |
| SERVER_NAME = 'aioquic-transport' |
| |
| handlers_path = None |
| |
| |
| class EventHandler: |
| def __init__(self, connection: QuicConnectionProtocol, global_dict: Dict): |
| self.connection = connection |
| self.global_dict = global_dict |
| |
| def handle_client_indication( |
| self, |
| origin: str, |
| query: Dict[str, str]) -> None: |
| name = 'handle_client_indication' |
| if name in self.global_dict: |
| self.global_dict[name](self.connection, origin, query) |
| |
| def handle_event(self, event: QuicEvent) -> None: |
| name = 'handle_event' |
| if name in self.global_dict: |
| self.global_dict[name](self.connection, event) |
| |
| |
| class QuicTransportProtocol(QuicConnectionProtocol): |
| def __init__(self, *args, **kwargs) -> None: |
| super().__init__(*args, **kwargs) |
| self.streams = dict() |
| self.pending_events = [] |
| self.client_indication_finished = False |
| self.client_indication_data = b'' |
| self.handler = None |
| |
| def quic_event_received(self, event: QuicEvent) -> None: |
| prefix = '!!' |
| logging.log(logging.INFO, 'QUIC event: %s' % type(event)) |
| try: |
| if (not self.client_indication_finished and |
| isinstance(event, StreamDataReceived) and |
| event.stream_id == 2): |
| # client indication process |
| prefix = 'Client indication error: ' |
| self.client_indication_data += event.data |
| if len(self.client_indication_data) > 65535: |
| raise Exception('too large data for client indication') |
| if event.end_stream: |
| self.process_client_indication() |
| if self.is_closing_or_closed(): |
| return |
| prefix = 'Event handling Error: ' |
| for e in self.pending_events: |
| self.handler.handle_event(e) |
| self.pending_events.clear() |
| elif not self.client_indication_finished: |
| self.pending_events.append(event) |
| elif self.handler is not None: |
| prefix = 'Event handling Error: ' |
| self.handler.handle_event(event) |
| except Exception as e: |
| self.handler = None |
| logging.log(logging.WARN, prefix + str(e)) |
| self.close() |
| |
| def parse_client_indication(self, bs): |
| while True: |
| key_b = bs.read(2) |
| if len(key_b) == 0: |
| return |
| length_b = bs.read(2) |
| if len(key_b) != 2: |
| raise Exception('failed to get "Key" field') |
| if len(length_b) != 2: |
| raise Exception('failed to get "Length" field') |
| key = struct.unpack('!H', key_b)[0] |
| length = struct.unpack('!H', length_b)[0] |
| value = bs.read(length) |
| if len(value) != length: |
| raise Exception('truncated "Value" field') |
| yield (key, value) |
| |
| def process_client_indication(self) -> None: |
| origin = None |
| origin_string = None |
| path = None |
| path_string = None |
| KEY_ORIGIN = 0 |
| KEY_PATH = 1 |
| for (key, value) in self.parse_client_indication( |
| io.BytesIO(self.client_indication_data)): |
| if key == KEY_ORIGIN: |
| origin_string = value.decode() |
| origin = urllib.parse.urlparse(origin_string) |
| elif key == KEY_PATH: |
| path_string = value.decode() |
| path = urllib.parse.urlparse(path_string) |
| else: |
| # We must ignore unrecognized fields. |
| pass |
| logging.log(logging.INFO, |
| 'origin = %s, path = %s' % (origin_string, path_string)) |
| if origin is None: |
| raise Exception('No origin is given') |
| if path is None: |
| raise Exception('No path is given') |
| if origin.scheme != 'https' and origin.scheme != 'http': |
| raise Exception('Invalid origin: %s' % origin_string) |
| if origin.netloc == '': |
| raise Exception('Invalid origin: %s' % origin_string) |
| |
| # To make the situation simple we accept only simple path strings. |
| m = re.compile('^/([a-zA-Z0-9\._\-]+)$').match(path.path) |
| if m is None: |
| raise Exception('Invalid path: %s' % path_string) |
| |
| handler_name = m.group(1) |
| query = dict(urllib.parse.parse_qsl(path.query)) |
| self.handler = self.create_event_handler(handler_name) |
| self.handler.handle_client_indication(origin_string, query) |
| if self.is_closing_or_closed(): |
| return |
| self.client_indication_finished = True |
| logging.log(logging.INFO, 'Client indication finished') |
| |
| def create_event_handler(self, handler_name: str) -> None: |
| global_dict = {} |
| with open(handlers_path + '/' + handler_name) as f: |
| exec(f.read(), global_dict) |
| return EventHandler(self, global_dict) |
| |
| def is_closing_or_closed(self) -> bool: |
| if self._quic._close_pending: |
| return True |
| if self._quic._state in END_STATES: |
| return True |
| return False |
| |
| |
| class SessionTicketStore: |
| ''' |
| Simple in-memory store for session tickets. |
| ''' |
| |
| def __init__(self) -> None: |
| self.tickets: Dict[bytes, SessionTicket] = {} |
| |
| def add(self, ticket: SessionTicket) -> None: |
| self.tickets[ticket.ticket] = ticket |
| |
| def pop(self, label: bytes) -> Optional[SessionTicket]: |
| return self.tickets.pop(label, None) |
| |
| |
| if __name__ == '__main__': |
| parser = argparse.ArgumentParser(description='QUIC server') |
| parser.add_argument( |
| '-c', |
| '--certificate', |
| type=str, |
| required=True, |
| help='load the TLS certificate from the specified file', |
| ) |
| parser.add_argument( |
| '--host', |
| type=str, |
| default='::', |
| help='listen on the specified address (defaults to ::)', |
| ) |
| parser.add_argument( |
| '--port', |
| type=int, |
| default=4433, |
| help='listen on the specified port (defaults to 4433)', |
| ) |
| parser.add_argument( |
| '-k', |
| '--private-key', |
| type=str, |
| required=True, |
| help='load the TLS private key from the specified file', |
| ) |
| parser.add_argument( |
| '--handlers-path', |
| type=str, |
| required=True, |
| help='the directory path of QuicTransport event handlers', |
| ) |
| parser.add_argument( |
| '-v', |
| '--verbose', |
| action='store_true', |
| help='increase logging verbosity' |
| ) |
| args = parser.parse_args() |
| |
| logging.basicConfig( |
| format='%(asctime)s %(levelname)s %(name)s %(message)s', |
| level=logging.DEBUG if args.verbose else logging.INFO, |
| ) |
| |
| configuration = QuicConfiguration( |
| alpn_protocols=['wq-vvv-01'] + ['siduck'], |
| is_client=False, |
| max_datagram_frame_size=65536, |
| ) |
| |
| handlers_path = os.path.abspath(os.path.expanduser(args.handlers_path)) |
| logging.log(logging.INFO, 'port = %s' % args.port) |
| logging.log(logging.INFO, 'handlers path = %s' % handlers_path) |
| |
| # load SSL certificate and key |
| configuration.load_cert_chain(args.certificate, args.private_key) |
| |
| ticket_store = SessionTicketStore() |
| |
| loop = asyncio.get_event_loop() |
| loop.run_until_complete( |
| serve( |
| args.host, |
| args.port, |
| configuration=configuration, |
| create_protocol=QuicTransportProtocol, |
| session_ticket_fetcher=ticket_store.pop, |
| session_ticket_handler=ticket_store.add, |
| ) |
| ) |
| try: |
| loop.run_forever() |
| except KeyboardInterrupt: |
| pass |
