clawhdf5/examples/wasm-viewer/test/serve.py

"""A static HTTP server for the wasm tests, with HTTP Range support and
request counting.

    python serve.py [--root PREFIX=DIR ...]

Serves each DIR under URL PREFIX (the first match wins; PREFIX "" is the
site root), prints the port on its first line of stdout, and runs until
killed. No symlinks or copies are made: files are read where they are.

- `Range: bytes=a-b`, `bytes=a-` and `bytes=-n` get 206 with Content-Range,
  an unsatisfiable range 416; every file answer carries an ETag, and CORS
  headers exposing Content-Range, so a page on another origin can use it.
- Under `/norange/...` the same files are served but Range is ignored
  (200 with the whole file), as by a server without range support.
- `GET /__stats` returns `{"requests": n, "bytes": n, "log": [...]}` for
  file requests since the last `GET /__reset`, which zeroes them.
"""

import argparse
import hashlib
import json
import os
import posixpath
import sys
import threading
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from urllib.parse import unquote, urlsplit

TYPES = {
    ".html": "text/html; charset=utf-8",
    ".js": "text/javascript; charset=utf-8",
    ".mjs": "text/javascript; charset=utf-8",
    ".wasm": "application/wasm",
    ".json": "application/json",
    ".ts": "text/plain; charset=utf-8",
}

lock = threading.Lock()
stats = {"requests": 0, "bytes": 0, "log": []}


def resolve(roots, path):
    """The file for URL `path`, or None. `..` never leaves a root."""
    parts = [p for p in posixpath.normpath(unquote(path)).split("/") if p]
    if any(p in (".", "..") for p in parts):
        return None
    for prefix, root in roots:
        pre = [p for p in prefix.split("/") if p]
        if parts[: len(pre)] == pre:
            rest = parts[len(pre):] or ["index.html"]
            f = os.path.join(root, *rest)
            if os.path.isfile(f):
                return f
    return None


def parse_range(header, size):
    """(start, end exclusive) for a single `bytes=` range, "bad" when
    unsatisfiable, None when absent or unparsable (served whole)."""
    if not header or not header.startswith("bytes=") or "," in header:
        return None
    a, _, b = header[len("bytes="):].strip().partition("-")
    try:
        if a == "":
            n = int(b)
            return (max(0, size - n), size) if n > 0 and size > 0 else "bad"
        start = int(a)
        end = int(b) + 1 if b else size
    except ValueError:
        return None
    if start >= size or end <= start:
        return "bad"
    return start, min(end, size)


def make_handler(roots):
    class Handler(BaseHTTPRequestHandler):
        protocol_version = "HTTP/1.1"

        def log_message(self, *args):
            pass

        def cors(self):
            self.send_header("Access-Control-Allow-Origin", "*")
            self.send_header("Access-Control-Expose-Headers",
                             "Content-Range, Content-Length, ETag, Accept-Ranges")

        def do_OPTIONS(self):
            self.send_response(204)
            self.cors()
            self.send_header("Access-Control-Allow-Headers", "Range")
            self.send_header("Content-Length", "0")
            self.end_headers()

        def do_HEAD(self):
            self.serve(head=True)

        def do_GET(self):
            self.serve(head=False)

        def json(self, obj):
            body = json.dumps(obj).encode()
            self.send_response(200)
            self.cors()
            self.send_header("Content-Type", "application/json")
            self.send_header("Content-Length", str(len(body)))
            self.send_header("Cache-Control", "no-store")
            self.end_headers()
            self.wfile.write(body)

        def serve(self, head):
            path = urlsplit(self.path).path
            if path == "/__stats":
                with lock:
                    return self.json(stats)
            if path == "/__reset":
                with lock:
                    stats.update(requests=0, bytes=0, log=[])
                return self.json({})
            ranges = True
            if path.startswith("/norange/"):
                ranges = False
                path = path[len("/norange"):]
            f = resolve(roots, path)
            if f is None:
                self.send_response(404)
                self.cors()
                self.send_header("Content-Length", "0")
                self.end_headers()
                return
            size = os.path.getsize(f)
            st = os.stat(f)
            etag = '"%s"' % hashlib.sha1(
                f"{f}:{size}:{st.st_mtime_ns}".encode()).hexdigest()[:16]
            r = parse_range(self.headers.get("Range"), size) if ranges else None
            if r == "bad":
                self.send_response(416)
                self.cors()
                self.send_header("Content-Range", f"bytes */{size}")
                self.send_header("Content-Length", "0")
                self.end_headers()
                return
            start, end = r if r else (0, size)
            self.send_response(206 if r else 200)
            self.cors()
            ext = os.path.splitext(f)[1]
            self.send_header("Content-Type", TYPES.get(ext, "application/octet-stream"))
            self.send_header("Content-Length", str(end - start))
            self.send_header("ETag", etag)
            self.send_header("Cache-Control", "no-store")
            if ranges:
                self.send_header("Accept-Ranges", "bytes")
            if r:
                self.send_header("Content-Range", f"bytes {start}-{end - 1}/{size}")
            self.end_headers()
            if not head:
                with lock:
                    stats["requests"] += 1
                    stats["bytes"] += end - start
                    stats["log"].append([path, start, end, 206 if r else 200])
                with open(f, "rb") as fh:
                    fh.seek(start)
                    left = end - start
                    try:
                        while left:
                            buf = fh.read(min(left, 1 << 20))
                            if not buf:
                                break
                            self.wfile.write(buf)
                            left -= len(buf)
                    except (BrokenPipeError, ConnectionResetError):
                        pass

    return Handler


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--root", action="append", default=[],
                    help="PREFIX=DIR: serve DIR under URL PREFIX")
    args = ap.parse_args()
    roots = []
    for spec in args.root:
        prefix, _, d = spec.partition("=")
        roots.append((prefix, os.path.abspath(d)))
    class Server(ThreadingHTTPServer):
        def handle_error(self, request, client_address):
            # A client that drops a connection (a cancelled download) is
            # not an error of the server.
            if not isinstance(sys.exc_info()[1], (ConnectionError, TimeoutError)):
                super().handle_error(request, client_address)

    httpd = Server(("127.0.0.1", 0), make_handler(roots))
    httpd.daemon_threads = True
    print(httpd.server_address[1], flush=True)
    sys.stdout.close()
    httpd.serve_forever()


if __name__ == "__main__":
    main()