T-HUB/infra/peertube-workers/bandwidth-probe.py

#!/usr/bin/python3
"""Bounded disposable transfer endpoint, bound ONLY to the PeerTube Tailscale IP."""
import argparse
import http.server
import ipaddress
import json
import os
import socket
import threading
import time
import urllib.parse

MAX_BYTES = 16 * 1024**2
CHUNK = bytes(range(256)) * 4096
SLOTS = threading.BoundedSemaphore(2)
LOCK = threading.Lock()
RECENT = {}


class Probe(http.server.BaseHTTPRequestHandler):
    protocol_version = 'HTTP/1.1'

    def setup(self):
        super().setup()
        self.connection.settimeout(30)

    def log_message(self, *_):
        pass

    def reply(self, code, data):
        blob = json.dumps(data).encode()
        self.send_response(code)
        self.send_header('Content-Type', 'application/json')
        self.send_header('Content-Length', str(len(blob)))
        self.send_header('Cache-Control', 'no-store')
        self.end_headers()
        self.wfile.write(blob)

    def admission(self, direction):
        address = ipaddress.ip_address(self.client_address[0])
        if address not in ipaddress.ip_network('100.64.0.0/10') and not address.is_loopback:
            self.reply(403, {'error': 'Tailscale clients only'})
            return False
        with LOCK:
            now = time.monotonic()
            for key in list(RECENT):
                if now - RECENT[key] > 60:
                    RECENT.pop(key)
            key = (str(address), direction)
            if now - RECENT.get(key, -60) < 10:
                self.reply(429, {'error': 'probe rate limit'})
                return False
            RECENT[key] = now
        if not SLOTS.acquire(blocking=False):
            self.reply(503, {'error': 'probe busy'})
            return False
        return True

    def do_GET(self):
        parsed = urllib.parse.urlsplit(self.path)
        if parsed.path == '/health':
            return self.reply(200, {'ok': True, 'max_bytes': MAX_BYTES, 'storage': 'discarded'})
        if parsed.path != '/download':
            return self.reply(404, {'error': 'unknown endpoint'})
        try:
            count = int(urllib.parse.parse_qs(parsed.query).get('bytes', [8 * 1024**2])[0])
            if not 65536 <= count <= MAX_BYTES:
                raise ValueError()
        except ValueError:
            return self.reply(400, {'error': 'invalid size'})
        if not self.admission('download'):
            return
        try:
            self.send_response(200)
            self.send_header('Content-Length', str(count))
            self.send_header('Content-Type', 'application/octet-stream')
            self.send_header('Cache-Control', 'no-store')
            self.end_headers()
            while count:
                block = CHUNK[:min(count, len(CHUNK))]
                self.wfile.write(block)
                count -= len(block)
        except (OSError, socket.timeout):
            self.close_connection = True
        finally:
            SLOTS.release()

    def do_PUT(self):
        if self.path != '/upload':
            self.close_connection = True
            return self.reply(404, {'error': 'unknown endpoint'})
        try:
            count = int(self.headers.get('Content-Length', '0'))
            if not 65536 <= count <= MAX_BYTES:
                raise ValueError()
        except ValueError:
            self.close_connection = True
            return self.reply(413, {'error': 'invalid size'})
        if not self.admission('upload'):
            self.close_connection = True
            return
        received = 0
        try:
            while received < count:
                data = self.rfile.read(min(1024**2, count - received))
                if not data:
                    raise OSError('short upload')
                received += len(data)
            self.reply(200, {'bytes_received': received})
        except (OSError, socket.timeout):
            self.close_connection = True
        finally:
            SLOTS.release()


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--bind', required=True)
    parser.add_argument('--port', type=int, default=9044)
    args = parser.parse_args()
    address = ipaddress.ip_address(args.bind)
    if address not in ipaddress.ip_network('100.64.0.0/10') and not address.is_loopback:
        raise SystemExit('Refusing a public/LAN/wildcard bind: require a Tailscale IP')
    server = http.server.ThreadingHTTPServer((args.bind, args.port), Probe)
    notify = os.environ.get('NOTIFY_SOCKET')
    if notify:
        address = '\0' + notify[1:] if notify.startswith('@') else notify
        with socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM) as channel:
            channel.connect(address)
            channel.sendall(b'READY=1')
    server.serve_forever()


if __name__ == '__main__':
    main()