Relieve backpressure

1 files changed, 70 insertions, 46 deletions

diff --git a/src/scadere/listen.py b/src/scadere/listen.py
index 80929ea..aa322ee 100644
--- a/src/scadere/listen.py
+++ b/src/scadere/listen.py
@@ -19,13 +19,14 @@
 from argparse import ArgumentParser
 from asyncio import run, start_server
 from base64 import urlsafe_b64decode as from_base64
-from datetime import datetime
+from datetime import datetime, timezone
 from functools import partial
+from http import HTTPStatus
 from pathlib import Path
 from urllib.parse import parse_qs, urljoin, urlsplit
 from xml.etree.ElementTree import (Element as xml_element,
                                    SubElement as xml_subelement,
-                                   indent, tostring as str_from_xml)
+                                   indent, tostringlist as strings_from_xml)
 from sys import argv
 
 from . import __version__, GNUHelpFormatter, NetLoc
@@ -43,6 +44,38 @@ def path(hostname, port, issuer, serial):
     return f'{hostname}/{port}/{issuer}/{serial}'
 
 
+async def write_status(writer, status):
+    """Write the given HTTP/1.1 status line."""
+    writer.write(f'HTTP/1.1 {status.value} {status.phrase}\r\n'.encode())
+    await writer.drain()
+
+
+async def write_content_type(writer, content_type):
+    """Write the given HTTP content type."""
+    writer.write(f'Content-Type: {content_type}\r\n'.encode())
+    await writer.drain()
+
+
+async def describe_status(writer, status):
+    """Write a HTTP/1.1 response including status description."""
+    await write_status(writer, status)
+    content = f'{status.description}\n'.encode()
+    await write_content_type(writer, 'text/plain')
+    writer.write(f'Content-Length: {len(content)}\r\n\r\n'.encode())
+    writer.write(content)
+    await writer.drain()
+
+
+async def write_xml(writer, document):
+    content = strings_from_xml(xml(document), 'unicode',
+                               xml_declaration=True, default_namespace=None)
+    length = len(''.join(content))
+    writer.write(f'Content-Length: {length}\r\n\r\n'.encode())
+    for part in content:
+        writer.write(part.encode())
+        await writer.drain()
+
+
 def body(not_before, not_after, hostname, port, serial, issuer):
     """Describe the given certificate in XHTML."""
     return (('h1', 'TLS certificate information'),
@@ -106,56 +139,47 @@ async def handle(certs, base_url, reader, writer):
     domains = tuple(parse_qs(url_parts.query).get('domain', ['']))
 
     if not request.startswith(b'GET '):
-        writer.write(b'HTTP/1.1 405 Method Not Allowed\r\n')
-        await writer.drain()
-        writer.close()
-        await writer.wait_closed()
-        return
+        await describe_status(writer, HTTPStatus.METHOD_NOT_ALLOWED)
     elif url.startswith(b'//'):  # urljoin goes haywire
-        writer.write(b'HTTP/1.1 404 Not Found\r\n')
+        await describe_status(writer, HTTPStatus.NOT_FOUND)
     elif url_parts.path == urlsplit(base_url).path:  # Atom feed
-        writer.write(b'HTTP/1.1 200 OK\r\n')
-        writer.write(b'Content-Type: application/atom+xml\r\n')
-        feed = xml(('feed', {'xmlns': 'http://www.w3.org/2005/Atom'},
-                    ('id', base_url),
-                    ('link', {'rel': 'self', 'href': base_url}),
-                    ('title', certs.name),
-                    ('updated', datetime.now().isoformat()),
-                    ('generator',
-                     {'uri': 'https://trong.loang.net/scadere/about',
-                      'version': __version__},
-                     'Scadere'),
-                    *(entry(base_url, cert)
-                      for cert in lookup.values()
-                      if cert[2].endswith(domains))))
-        content = str_from_xml(feed, 'unicode', xml_declaration=True,
-                               default_namespace=None).encode()
-        writer.write(f'Content-Length: {len(content)}\r\n\r\n'.encode())
-        writer.write(content)
+        await write_status(writer, HTTPStatus.OK)
+        await write_content_type(writer, 'application/atom+xml')
+        feed = ('feed', {'xmlns': 'http://www.w3.org/2005/Atom'},
+                ('id', base_url),
+                ('link', {'rel': 'self', 'href': base_url}),
+                ('title', certs.name),
+                ('updated', datetime.now(tz=timezone.utc).isoformat()),
+                ('generator',
+                 {'uri': 'https://trong.loang.net/scadere/about',
+                  'version': __version__},
+                 'Scadere'),
+                *(entry(base_url, cert) for cert in lookup.values()
+                  if cert[2].endswith(domains)))
+        await write_xml(writer, feed)
     elif url_parts.path in lookup:  # accessible Atom entry's link/ID
-        writer.write(b'HTTP/1.1 200 OK\r\n')
-        writer.write(b'Content-Type: application/xhtml+xml\r\n')
+        await write_status(writer, HTTPStatus.OK)
+        await write_content_type(writer, 'application/xhtml+xml')
         (not_before, not_after,
          hostname, port, serial, issuer) = lookup[url_parts.path]
-        page = xml(('html', {'xmlns': 'http://www.w3.org/1999/xhtml',
-                             'lang': 'en'},
-                    ('head',
-                     ('meta', {'name': 'color-scheme',
-                               'content': 'dark light'}),
-                     ('meta', {'name': 'viewport',
-                               'content': ('width=device-width,'
-                                           'initial-scale=1.0')}),
-                     ('link', {'rel': 'icon', 'href': 'data:,'}),
-                     ('title', f'TLS certificate - {hostname}:{port}')),
-                    ('body', *body(not_before, not_after,
-                                   hostname, port, serial, issuer))))
-        content = str_from_xml(page, 'unicode', xml_declaration=True,
-                               default_namespace=None).encode()
-        writer.write(f'Content-Length: {len(content)}\r\n\r\n'.encode())
-        writer.write(content)
+        page = ('html', {'xmlns': 'http://www.w3.org/1999/xhtml',
+                         'lang': 'en'},
+                ('head',
+                 ('meta', {'name': 'color-scheme',
+                           'content': 'dark light'}),
+                 ('meta', {'name': 'viewport',
+                           'content': ('width=device-width,'
+                                       'initial-scale=1.0')}),
+                 ('link', {'rel': 'icon', 'href': 'data:,'}),
+                 ('title', f'TLS certificate - {hostname}:{port}')),
+                ('body', *body(not_before, not_after,
+                               hostname, port, serial, issuer)))
+        await write_xml(writer, page)
     else:
-        writer.write(b'HTTP/1.1 404 Not Found\r\n')
-    await writer.drain()
+        await describe_status(writer, HTTPStatus.NOT_FOUND)
+
+    assert writer.can_write_eof()
+    writer.write_eof()
     writer.close()
     await writer.wait_closed()