Set env & mappath via CLI, add Args for Unix Sockets

jwindgassen · jwindgassen · commit 5cae144c7ca0 · 2024-12-06T13:45:00.000+01:00
diff --git a/jupyter_server_proxy/standalone/__init__.py b/jupyter_server_proxy/standalone/__init__.py
@@ -4,63 +4,63 @@
 
 from tornado import ioloop
 from tornado.httpserver import HTTPServer
-from tornado.log import app_log as log
+from tornado.log import app_log as log, enable_pretty_logging
 
 from .activity import start_activity_update
-from .proxy import configure_http_client, get_port_from_env, get_ssl_options, make_app
+from .proxy import configure_http_client, get_port_from_env, get_ssl_options, make_proxy_app
 
 
 def run(
-    command,
-    port,
-    destport,
-    ip,
-    debug,
-    logs,
-    overwrite_authentication,
-    timeout,
-    activity_interval,
-    progressive,
-    websocket_max_message_size,
+    command: list[str],
+    port: int,
+    address: str,
+    server_port: int,
+    socket_path: str | None,
+    socket_auto: bool,
+    environment: list[tuple[str, str]] | None,
+    mappath: list[tuple[str, str]] | None,
+    debug: bool,
+    # logs: bool,
+    overwrite_authentication: bool | None,
+    timeout: int,
+    activity_interval: int,
+    # progressive: bool,
+    websocket_max_message_size: int,
 ):
-    if port is None:
-        get_port_from_env()
-
+    # Setup Logging
+    enable_pretty_logging(logger=log)
     if debug:
         log.setLevel(logging.DEBUG)
-    elif logs:
-        log.setLevel(logging.INFO)
 
-    prefix = os.environ.get("JUPYTERHUB_SERVICE_PREFIX", "/")
+    if not port:
+        port = get_port_from_env()
 
-    if len(prefix) > 0 and prefix[-1] == "/":
-        prefix = prefix[:-1]
+    if overwrite_authentication is True:
+        log.info("Enabling Authentication with JupyterHub")
 
-    configure_http_client()
+    prefix = os.environ.get("JUPYTERHUB_SERVICE_PREFIX", "/")
 
-    app = make_app(
-        destport,
-        prefix,
-        list(command),
-        overwrite_authentication,
+    app = make_proxy_app(
+        command,
+        prefix.removesuffix("/"),
+        server_port,
+        socket_path or socket_auto,
+        dict(environment),
+        dict(mappath),
         timeout,
+        overwrite_authentication is True,
         debug,
-        logs,
-        progressive,
         websocket_max_message_size,
     )
 
     ssl_options = get_ssl_options()
 
     http_server = HTTPServer(app, ssl_options=ssl_options, xheaders=True)
+    http_server.listen(port, address)
 
-    http_server.listen(port or get_port_from_env(), ip)
-
-    log.info(
-        f"Starting standaloneproxy on {ip}:{port}, server is started on Port {destport}"
-    )
-    log.info(f"URL Prefix: {prefix}")
-    log.info(f"Command: {command}")
+    log.info(f"Starting standaloneproxy on '{address}:{port}'")
+    log.info(f"URL Prefix: {prefix!r}")
+    log.info(f"Command: {' '.join(command)!r}")
 
     if activity_interval > 0:
         start_activity_update(activity_interval)
@@ -70,36 +70,81 @@ def run(
 
 def main():
     parser = argparse.ArgumentParser(
-        "jupyter-native-proxy",
-        description="Wrap an arbitrary WebApp so it can be used in place of 'singleuser' in a JupyterHub setting",
+        "jupyter-standalone-proxy",
+        description="Wrap an arbitrary WebApp so it can be used in place of 'jupyterhub-singleuser' in a JupyterHub setting.",
         formatter_class=argparse.ArgumentDefaultsHelpFormatter,
     )
 
     parser.add_argument(
+        "-p",
         "--port",
-        default=None,
+        default=0,
         type=int,
-        help="Port for the proxy server to listen on. Defaults to JupyterHub default.",
+        dest="port",
+        help="Port for the proxy server to listen on (0 for JupyterHub default).",
+    )
+    parser.add_argument(
+        "-a",
+        "--address",
+        default="localhost",
+        type=str,
+        dest="address",
+        help="Address for the proxy server to listen on.",
     )
     parser.add_argument(
-        "--destport",
+        "-s",
+        "--server-port",
         default=0,
         type=int,
-        help="Port for the WebApp should end up running on. Leave at 0 for a random open port.",
+        dest="server_port",
+        help="Port for the WebApp should end up running on (0 for random open port).",
+    )
+    parser.add_argument(
+        "--socket-path",
+        type=str,
+        default=None,
+        help="Path to the Unix Socket to use for proxying. Takes precedence over '-s/--server_port' and '--socket-auto'.",
+    )
+    parser.add_argument(
+        "--socket-auto",
+        action="store_true",
+        help="Use Unix Socket for proxying, but let Jupyter Server Proxy automatically create one.",
+    )
+    parser.add_argument(
+        "--env",
+        "--environment",
+        type=lambda v: tuple(v.split(":")[:2]),
+        default=[],
+        action="append",
+        dest="environment",
+        help="Add an environment variable to the server process. Must be of the form <Name>:<Value>, e.g. --env=MY_VAR:42",
     )
-    parser.add_argument("--ip", default="localhost", help="Address to listen on.")
     parser.add_argument(
-        "--debug", action="store_true", default=False, help="Display debug level logs."
+        "--mappath",
+        type=lambda v: tuple(v.split(":")[:2]),
+        default=[],
+        action="append",
+        help="Add an path mapping to the proxy. Any requests received under <Source> will be redirected to <Target>. "
+        "Must be of the form <Source>:<Target>, e.g. --mappath=/:/index.html",
     )
     parser.add_argument(
-        "--logs",
+        "-d",
+        "--debug",
         action="store_true",
-        default=True,
-        help="Display logs generated by the subprocess.",
+        default=False,
+        dest="debug",
+        help="Display debug level logs.",
     )
+    # parser.add_argument(
+    #     "--logs",
+    #     action="store_true",
+    #     default=True,
+    #     help="Display logs generated by the subprocess.",
+    # )
     parser.add_argument(
         "--overwrite-authentication",
         default=None,
+        type=lambda v: None if v is None else bool(v),
         help="Forcefully enable/disable authentication with JupyterHub.",
     )
     parser.add_argument(
@@ -114,12 +159,12 @@ def main():
         type=int,
         help="Frequency to notify Hub that the WebApp is still running (In seconds, 0 for never).",
     )
-    parser.add_argument(
-        "--progressive",
-        action="store_true",
-        default=False,
-        help="Progressively flush responses as they arrive (good for Voila).",
-    )
+    # parser.add_argument(
+    #     "--progressive",
+    #     action="store_true",
+    #     default=False,
+    #     help="Progressively flush responses as they arrive (good for Voila).",
+    # )
     parser.add_argument(
         "--websocket-max-message-size",
         default=0,
diff --git a/jupyter_server_proxy/standalone/proxy.py b/jupyter_server_proxy/standalone/proxy.py
@@ -6,6 +6,7 @@
 from jupyterhub.services.auth import HubOAuthCallbackHandler, HubOAuthenticated
 from jupyterhub.utils import make_ssl_context
 from tornado import httpclient, web
+from tornado.log import app_log
 from tornado.web import Application, RedirectHandler
 from tornado.websocket import WebSocketHandler
 
@@ -78,57 +79,41 @@ def set_default_headers(self):
         self.set_header("X-JupyterHub-Version", __jh_version__)
 
 
-def _make_native_proxy_handler(
-    command, port, mappath, overwrite_authentication, environment, timeout
+def make_proxy_app(
+    command: list[str],
+    prefix: str,
+    port: int,
+    unix_socket: bool | str,
+    environment: dict[str, str],
+    mappath: dict[str, str],
+    timeout: int,
+    skip_authentication: bool,
+    debug: bool,
+    # progressive: bool,
+    websocket_max_message_size: int,
 ):
     """
     Create a StandaloneHubProxyHandler subclass with given parameters
     """
+    app_log.debug(f"Process will use {port = }")
+    app_log.debug(f"Process will use {unix_socket = }")
+    app_log.debug(f"Process environment: {environment}")
+    app_log.debug(f"Proxy mappath: {mappath}")
 
-    # Try to determine if we are launched by a JupyterHub via environment variables.
-    # See jupyterhub/spawner.py#L1035
-    if overwrite_authentication is not None:
-        base = (
-            StandaloneHubProxyHandler
-            if overwrite_authentication is True
-            else StandaloneProxyHandler
-        )
-    elif "JUPYTERHUB_API_TOKEN" in os.environ and "JUPYTERHUB_API_URL" in os.environ:
-        base = StandaloneHubProxyHandler
-    else:
-        base = StandaloneProxyHandler
-
-    class _Proxy(base):
+    base = StandaloneProxyHandler if skip_authentication else StandaloneHubProxyHandler
+    class Proxy(base):
         def __init__(self, *args, **kwargs):
             super().__init__(*args, **kwargs)
-            self.name = command[0]
+            self.name = f"{command[0]!r} Process"
             self.proxy_base = command[0]
             self.requested_port = port
+            self.requested_unix_socket = unix_socket
             self.mappath = mappath
             self.command = command
             self.environment = environment
             self.timeout = timeout
 
-    return _Proxy
-
-
-def make_app(
-    destport,
-    prefix,
-    command,
-    authtype,
-    timeout,
-    debug,
-    logs,
-    progressive,
-    websocket_max_message_size,
-):
-    # ToDo: Environment
-    proxy_handler = _make_native_proxy_handler(
-        command, destport, {}, authtype, {}, timeout
-    )
-
-    options = dict(
+    settings = dict(
         debug=debug,
         # Required for JupyterHub Authentication
         hub_user=os.environ.get("JUPYTERHUB_USER", ""),
@@ -137,7 +122,8 @@ def make_app(
     )
 
     if websocket_max_message_size:
-        options["websocket_max_message_size"] = websocket_max_message_size
+        app_log.debug(f"Restricting WebSocket Messages to {websocket_max_message_size}")
+        settings["websocket_max_message_size"] = websocket_max_message_size
 
     return Application(
         [
@@ -147,7 +133,7 @@ def make_app(
             ),
             (
                 r"^" + re.escape(prefix) + r"/(.*)",
-                proxy_handler,
+                Proxy,
                 dict(
                     state={},
                     # ToDo: progressive=progressive
@@ -159,7 +145,7 @@ def make_app(
                 dict(url=prefix + "/{0}"),
             ),
         ],
-        **options,
+        **settings,
     )
 
 
