Skip to main content

RELAY_SCRIPT

Constant RELAY_SCRIPT 

Source
pub const RELAY_SCRIPT: &str = r#"import os, signal, socket, sys, threading

MAX_CONNS = 8
MAX_BYTES = 2 * 1024 * 1024
IDLE_SECS = 30


def log(msg):
    sys.stderr.write("[ingest-relay] %s\n" % msg)


def addr(s):
    host, _, port = s.rpartition(":")
    return host.strip("[]"), int(port)


def parse(argv):
    opts, i = {}, 0
    while i < len(argv) and argv[i] != "--":
        opts[argv[i]] = argv[i + 1]
        i += 2
    cmd = argv[i + 1:]
    if not cmd:
        raise SystemExit("usage: --listen H:P --upstream H:P [--uid N --gid N] -- cmd ...")
    return (addr(opts["--listen"]), addr(opts["--upstream"]),
            int(opts.get("--uid", "65534")), int(opts.get("--gid", "65534")), cmd)


def drop(uid, gid):
    if os.geteuid() == 0:
        os.setgroups([])
        os.setgid(gid)
        os.setuid(uid)
    no_new_privs()


def no_new_privs():
    """prctl(PR_SET_NO_NEW_PRIVS): setuid/setgid bits and file capabilities
    are ignored from here on (Linux only; a no-op elsewhere)."""
    if not sys.platform.startswith("linux"):
        return
    import ctypes
    libc = ctypes.CDLL(None, use_errno=True)
    if libc.prctl(38, 1, 0, 0, 0) != 0:
        raise OSError(ctypes.get_errno(), "prctl(PR_SET_NO_NEW_PRIVS)")


def pump(src, dst):
    total = 0
    try:
        while True:
            data = src.recv(65536)
            if not data:
                break
            total += len(data)
            if total > MAX_BYTES:
                log("connection over %d bytes, closed" % MAX_BYTES)
                break
            dst.sendall(data)
    except OSError:
        pass
    for s in (src, dst):
        try:
            s.shutdown(socket.SHUT_RDWR)
        except OSError:
            pass


def serve(conn, upstream, slots):
    try:
        try:
            up = socket.create_connection(upstream, timeout=10)
        except OSError as e:
            log("upstream %s:%d: %s" % (upstream[0], upstream[1], e))
            return
        conn.settimeout(IDLE_SECS)
        up.settimeout(IDLE_SECS)
        back = threading.Thread(target=pump, args=(up, conn), daemon=True)
        back.start()
        pump(conn, up)
        back.join(IDLE_SECS)
        up.close()
    finally:
        conn.close()
        slots.release()


def accept_loop(srv, upstream):
    slots = threading.BoundedSemaphore(MAX_CONNS)
    while True:
        conn, _ = srv.accept()
        if not slots.acquire(blocking=False):
            conn.close()
            continue
        threading.Thread(target=serve, args=(conn, upstream, slots), daemon=True).start()


def main():
    listen, upstream, uid, gid, cmd = parse(sys.argv[1:])
    srv = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    srv.bind(listen)
    srv.listen(MAX_CONNS)
    pid = os.fork()
    if pid == 0:
        try:
            srv.close()
            drop(uid, gid)
            os.execv(cmd[0], cmd)
        except BaseException as e:
            log("cog exec failed: %s" % e)
            os._exit(126)
    drop(uid, gid)
    for sig in (signal.SIGTERM, signal.SIGINT):
        signal.signal(sig, lambda s, _f: os.kill(pid, s))
    threading.Thread(target=accept_loop, args=(srv, upstream), daemon=True).start()
    while True:
        try:
            _, status = os.waitpid(pid, 0)
            break
        except InterruptedError:
            continue
    code = os.waitstatus_to_exitcode(status)
    sys.exit(code if code >= 0 else 128 - code)


main()
"#;
Expand description

The relay script (see the module docs).