Merge pull request #238 from moi-si/patch

Safely close StreamWriter and cancel stream task

Author: URenko

accesser/__init__.py

@@ -25,7 +25,6 @@
 import ssl
 import asyncio
 import traceback
-from contextlib import closing
 from urllib import request
 from urllib.parse import urlsplit
 from packaging.version import Version
@@ -46,9 +45,9 @@ async def update_cert(server_name):
     if not is_tld(server_name):
         res = get_tld(server_name, as_object=True, fix_protocol=True)
         if res.subdomain:
-            server_name = res.subdomain.split('.', 1)[-1] + '.' + res.domain + '.' + res.tld
+            server_name = f"{res.subdomain.split('.', 1)[-1]}.{res.domain}.{res.tld}"
         else:
-            server_name = res.domain + '.' + res.tld
+            server_name = f'{res.domain}.{res.tld}'
     async with cert_lock:
         if not server_name in cert_store:
             cm.create_certificate(server_name)
@@ -79,7 +78,7 @@ async def http_redirect(writer: asyncio.StreamWriter, path: str):
         if path.startswith(key):
             path = setting.config['http_redirect'][key] + path[len(key):]
             break
-    logger.debug('Redirect to '+path)
+    logger.debug('Redirect to %s', path)
     writer.write(f'HTTP/1.1 301 Moved Permanently\r\nLocation: https://{path}\r\n\r\n'.encode('iso-8859-1'))
     await writer.drain()
     writer.close()
@@ -95,18 +94,19 @@ async def forward_stream(reader: asyncio.StreamReader, writer: asyncio.StreamWri
 
 async def handle(reader, writer):
     global context
-    with closing(writer):
+    try:
+        remote_writer = None
         raw_request = await reader.readuntil(b'\r\n\r\n')
         requestline = raw_request.decode('iso-8859-1').splitlines()[0]
         i_addr, i_port, *_ = writer.get_extra_info('peername')
-        logger.debug(f"{i_addr}:{i_port} say: {requestline}")
+        logger.debug("%s:%d say: %s", i_addr, i_port, requestline)
         words = requestline.split()
         command, path = words[:2]
         match command:
             case 'CONNECT':
                 host, port = path.split(':')
                 remote_ip = await DNSquery(host)
-                logger.debug(f'[{i_port:5}] DNS: {host} -> {remote_ip}')
+                logger.debug('[%5d] DNS: %s -> %s', i_port, host, remote_ip)
             case 'GET':
                 if path.startswith('/pac/'):
                     return await send_pac(writer)
@@ -127,13 +127,13 @@ async def handle(reader, writer):
             writer._transport = await writer._loop.start_tls(writer.transport, writer._protocol, context, server_side=True)
         server_hostname_key = next(filter(lambda h:fnmatch.fnmatchcase(host, h), setting.config['alter_hostname']), None)
         server_hostname = '' if server_hostname_key is None else setting.config['alter_hostname'][server_hostname_key]
-        logger.debug(f'[{i_port:5}] {server_hostname=}')
+        logger.debug("[%5d] server_hostname: %s", i_port, server_hostname)
         remote_context = ssl.create_default_context()
         remote_context.check_hostname = False
         remote_reader, remote_writer = await asyncio.open_connection(remote_ip, port, ssl=remote_context, server_hostname=server_hostname)
         cert = remote_writer.get_extra_info('peercert')
         cert_message = f"subjectAltName: {cert.get('subjectAltName', ())}, subject: {cert.get('subject', ())}"
-        logger.debug(f"[{i_port:5}] {cert_message}.")
+        logger.debug("[%5d] %s", i_port, cert_message)
         cert_verify_key = next(filter(lambda h:fnmatch.fnmatchcase(host, h), setting.config.get('cert_verify', ())), None)
         if cert_verify_key is not None:
             cert_verify_list = setting.config['cert_verify'][cert_verify_key]
@@ -144,23 +144,29 @@ async def handle(reader, writer):
         else:
             cert_verify_list = [host]
             cert_policy = setting.config['check_hostname']
-        if  cert_policy is not False and not any(match_hostname(cert, h, cert_policy) for h in cert_verify_list):
-            logger.warning(f"[{i_port:5}] {cert_verify_list} don't match either of {cert_message}.")
+        if cert_policy is not False and not any(match_hostname(cert, h, cert_policy) for h in cert_verify_list):
+            logger.warning("[%5d] %s don't match either of %s.", i_port, cert_verify_list, cert_message)
             return
-        await asyncio.gather(
-            forward_stream(reader, remote_writer),
-            forward_stream(remote_reader, writer)
+        tasks = (
+            asyncio.create_task(forward_stream(reader, remote_writer)),
+            asyncio.create_task(forward_stream(remote_reader, writer))
         )
-    writer.close()
-    remote_writer.close()
-    await remote_writer.wait_closed()
-    await writer.wait_closed()
+        done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
+        for task in pending:
+            task.cancel()
+        await asyncio.gather(*pending)
+    finally:
+        writer.close()
+        if remote_writer:
+            remote_writer.close()
+            await remote_writer.wait_closed()
+        await writer.wait_closed()
 
 async def proxy():
     server = await asyncio.start_server(handle, setting.config['server']['address'], setting.config['server']['port'])
 
     print(f"Serving on {', '.join(str(sock.getsockname()) for sock in server.sockets)}")
-    sysproxy.set_pac('http://localhost:'+str(setting.config['server']['port'])+'/pac/?t='+str(random.randrange(2**16)))
+    sysproxy.set_pac(f"http://localhost:{setting.config['server']['port']}/pac/?t={random.randrange(2**16)}")
 
     try:
         async with server:
@@ -171,7 +177,7 @@ async def proxy():
 async def DNSquery(domain, hosts_only=False):
     global DNSresolver
     try:
-        return next(v for k,v in setting.config['hosts'].items() if k==domain or (k.startswith('.') and domain.endswith(k)))
+        return next(v for k, v in setting.config['hosts'].items() if k == domain or (k.startswith('.') and domain.endswith(k)))
     except StopIteration:
         if hosts_only:
             return
@@ -185,7 +191,7 @@ async def DNSquery(domain, hosts_only=False):
     return ret[0].to_text()
 
 def update_checker():
-    for pypi_url in ['https://pypi.org/pypi/accesser/json', 'https://mirrors.cloud.tencent.com/pypi/json/accesser']:
+    for pypi_url in ('https://pypi.org/pypi/accesser/json', 'https://mirrors.cloud.tencent.com/pypi/json/accesser'):
         try:
             with request.urlopen(pypi_url) as f:
                 v2 = Version(json.load(f)["info"]["version"])
@@ -197,7 +203,7 @@ def update_checker():
             v2 = Version(f.geturl().rsplit('/', maxsplit=1)[-1])
     v1 = Version(__version__)
     if v2 > v1:
-        logger.warning("There is a new version, you can update with 'python3 -m pip install -U accesser' or download from GitHub")
+        logger.warning("There is a new version. You can update with `python3 -m pip install -U accesser` or download from GitHub.")
 
 async def main():
     global context, cert_store, cert_lock, DNSresolver
@@ -207,7 +213,7 @@ async def main():
     if setting.rules_update_case in ('old', 'missing'):
         logger.warning("Updated rules.toml because it is %s.", setting.rules_update_case)
     elif setting.rules_update_case == 'modified':
-        logger.warning("You've already modified rules.toml, so it won't be updated automatically!")
+        logger.warning("You've already modified rules.toml so it won't be updated automatically!")
     else:
         logger.debug("rules.toml status: %s", setting.rules_update_case)
 
@@ -220,14 +226,15 @@ async def main():
         if (_url := urlsplit(nameserver)).netloc == '':
             _url = urlsplit('//' + nameserver)
         address = await DNSquery(_url.hostname, hosts_only=True)
-        if _url.scheme == '':
-            DNSresolver.nameservers.append(dns.nameserver.Do53Nameserver(_url.hostname if address is None else address, 53 if _url.port is None else _url.port))
-        elif _url.scheme == 'https':
-            DNSresolver.nameservers.append(dns.nameserver.DoHNameserver(nameserver, bootstrap_address=address))
-        elif _url.scheme == 'tls':
-            DNSresolver.nameservers.append(dns.nameserver.DoTNameserver(_url.hostname if address is None else address, 853 if _url.port is None else _url.port, hostname=_url.hostname))
-        elif _url.scheme == 'quic':
-            DNSresolver.nameservers.append(dns.nameserver.DoQNameserver(_url.hostname if address is None else address, 853 if _url.port is None else _url.port, server_hostname=_url.hostname))
+        match _url.scheme:
+            case '':
+                DNSresolver.nameservers.append(dns.nameserver.Do53Nameserver(_url.hostname if address is None else address, 53 if _url.port is None else _url.port))
+            case 'https':
+                DNSresolver.nameservers.append(dns.nameserver.DoHNameserver(nameserver, bootstrap_address=address))
+            case 'tls':
+                DNSresolver.nameservers.append(dns.nameserver.DoTNameserver(_url.hostname if address is None else address, 853 if _url.port is None else _url.port, hostname=_url.hostname))
+            case 'quic':
+                DNSresolver.nameservers.append(dns.nameserver.DoQNameserver(_url.hostname if address is None else address, 853 if _url.port is None else _url.port, server_hostname=_url.hostname))
 
     importca.import_ca()