product-recognition-demo/recognization/code/transport.py

250 lines
7.8 KiB
Python

"""
Request bounds and the LAN token check.
A client can declare a tiny Content-Length and then keep writing, so the
declared length is a hint and the read itself stops at the cap. The token
never travels in a query string: proxies and access logs keep those.
"""
from __future__ import annotations
import hmac
import ipaddress
import os
import shutil
import signal
import socket
import subprocess
from typing import Any, Dict, Optional
LOOPBACK_HOSTS = frozenset({"127.0.0.1", "::1", "localhost"})
def is_loopback_host(host: Optional[str]) -> bool:
text = (host or "").strip().lower()
if text.startswith("[") and text.endswith("]"):
text = text[1:-1]
if "%" in text: # fe80::1%lo0
text = text.split("%", 1)[0]
if text in LOOPBACK_HOSTS:
return True
try:
return ipaddress.ip_address(text).is_loopback
except ValueError:
return False
def require_token_for_bind(host: Optional[str], token: Optional[str]) -> None:
"""A non-loopback bind with no token is reachable by anyone on that network."""
if token or is_loopback_host(host):
return
raise SystemExit(
f"refusing to bind {host} without DORAYAKI_API_TOKEN "
"(loopback is the only bind that may run with auth off)"
)
def token_matches(presented: Optional[str], token: Optional[str]) -> bool:
if not presented or not token:
return False
# compare_digest raises when the lengths differ. Fold the length in so a
# short guess neither throws nor returns early.
return hmac.compare_digest(
len(presented).to_bytes(8, "big") + presented.encode(),
len(token).to_bytes(8, "big") + token.encode(),
)
def bearer_matches(authorization: Optional[str], token: str) -> bool:
if not authorization or not token:
return False
scheme, _, presented = authorization.partition(" ")
if scheme.lower() != "bearer" or not presented or " " in presented:
return False
return token_matches(presented, token)
def declared_length(header: Optional[str]) -> Optional[int]:
if header is None or header == "":
return None
try:
length = int(header)
except (TypeError, ValueError):
return None
return length if length >= 0 else None
async def read_bounded(source, max_bytes: int) -> bytes:
"""Stop once the cap is passed. The caller rejects the oversize body.
A request stream is async-iterable. An UploadFile is not: it has read(n).
"""
out = bytearray()
read = getattr(source, "read", None)
if read is not None and not hasattr(source, "__aiter__"):
while True:
chunk = await read(65536)
if not chunk:
break
out += chunk
if len(out) > max_bytes:
break
return bytes(out)
async for chunk in source:
if not chunk:
continue
out += chunk
if len(out) > max_bytes:
break
return bytes(out)
class BodyLimit:
"""Stop the ASGI receive once the body passes the cap.
A handler that checks Content-Length, or reads request.body(), runs after
the server has already buffered the stream. This counts bytes at receive
and returns 413 without handing the rest to the app. WebSocket frames are
not HTTP bodies and are left alone.
"""
def __init__(self, app: Any, max_bytes: int) -> None:
self.app = app
self.max_bytes = int(max_bytes)
async def __call__(self, scope, receive, send) -> None:
if scope.get("type") != "http":
await self.app(scope, receive, send)
return
class TooLarge(Exception):
pass
seen = 0
async def limited() -> Dict[str, Any]:
nonlocal seen
message = await receive()
if message.get("type") != "http.request":
return message
body = message.get("body") or b""
seen += len(body)
if seen > self.max_bytes:
raise TooLarge
return message
try:
await self.app(scope, limited, send)
except TooLarge:
await send(
{
"type": "http.response.start",
"status": 413,
"headers": [(b"content-type", b"application/json")],
}
)
await send(
{
"type": "http.response.body",
"body": b'{"error":{"code":"PAYLOAD_TOO_LARGE","message":"request body too large"}}',
}
)
SERVICE_TYPE = "_dorayaki._tcp"
def _lan_address() -> Optional[str]:
"""Address a phone on this network can open. Loopback is not one.
The demo launcher uses `ipconfig getifaddr en0` then `en1`. A UDP connect
is the fallback when those interfaces are absent.
"""
if shutil.which("ipconfig"):
for iface in ("en0", "en1"):
try:
out = subprocess.check_output(
["ipconfig", "getifaddr", iface], stderr=subprocess.DEVNULL, text=True
).strip()
except (subprocess.CalledProcessError, OSError):
continue
if out and not is_loopback_host(out):
return out
try:
probe = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
try:
probe.connect(("192.0.2.1", 9))
found = probe.getsockname()[0]
finally:
probe.close()
except OSError:
return None
if not found or is_loopback_host(found):
return None
return found
def advertise_host(bind_host: Optional[str]) -> Optional[str]:
"""The address clients should browse. A loopback bind is not advertised.
A wildcard bind (0.0.0.0, ::) has no address of its own, so the record
uses the LAN address the demo launcher would hand a phone.
"""
text = (bind_host or "").strip()
if text.startswith("[") and text.endswith("]"):
text = text[1:-1]
if is_loopback_host(text):
return None
if text in ("", "0.0.0.0", "::"):
return _lan_address()
try:
parsed = ipaddress.ip_address(text)
except ValueError:
return None
if parsed.is_loopback or parsed.is_unspecified:
return None if parsed.is_loopback else _lan_address()
return str(parsed)
class ServiceAdvertiser:
"""Register `_dorayaki._tcp` with the system mDNS daemon and drop it on close.
ponytail: macOS `dns-sd` only. A Linux host needs `avahi-publish` wired the
same way; do not add a Python mDNS package for this one record.
"""
def __init__(self, host: str, port: int, name: str = "Dorayaki") -> None:
self.host = host
self.port = int(port)
self.name = name
self._process: Optional[subprocess.Popen] = None
def start(self) -> bool:
if self._process is not None:
return True
binary = shutil.which("dns-sd")
if not binary:
return False
self._process = subprocess.Popen(
[binary, "-P", self.name, SERVICE_TYPE, "local", str(self.port), self.host, self.host],
stdout=subprocess.DEVNULL,
stderr=subprocess.DEVNULL,
start_new_session=True,
)
return True
def close(self) -> None:
process = self._process
self._process = None
if process is None or process.poll() is not None:
return
try:
os.killpg(process.pid, signal.SIGTERM)
except ProcessLookupError:
return
try:
process.wait(timeout=2)
except subprocess.TimeoutExpired:
try:
os.killpg(process.pid, signal.SIGKILL)
except ProcessLookupError:
return
process.wait(timeout=2)