Skip to content

Commit

Permalink
Support named servers
Browse files Browse the repository at this point in the history
  • Loading branch information
vsavinkov committed Jan 25, 2024
1 parent a054eda commit 91a290f
Showing 1 changed file with 44 additions and 14 deletions.
58 changes: 44 additions & 14 deletions jupyterhub_ssh/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -34,29 +34,41 @@ def connection_made(self, conn):
def password_auth_supported(self):
return True

async def get_user_server_url(self, session, username):
async def get_user_server_url(self, session, username, force_default=False):
"""
Return user's server url if it is running.
Else return None
"""
server_name = ""
if "-" in username and not force_default:
username, server_name = username.split("-", 1)
print("Using username for custom instance", username, server_name)
else:
print("Using username for default instance", username)
async with session.get(self.app.hub_url / "hub/api/users" / username) as resp:
if resp.status != 200:
return None
user = await resp.json()
print(user)
# URLs will have preceding slash, but yarl forbids those
server = user.get("servers", {}).get("", {})
if server.get("ready", False):
return self.app.hub_url / user["servers"][""]["url"][1:]
return self.app.hub_url / user["servers"][server_name]["url"][1:]
else:
return None

async def start_user_server(self, session, username):
async def start_user_server(self, session, username, force_default=False):
""" """
# REST API reference: https://jupyterhub.readthedocs.io/en/stable/_static/rest-api/index.html#operation--users--name--server-post
# REST API implementation: https://github.com/jupyterhub/jupyterhub/blob/187fe911edce06eb067f736eaf4cc9ea52e69e08/jupyterhub/apihandlers/users.py#L451-L497
create_url = self.app.hub_url / "hub/api/users" / username / "server"
server_name = ""
if "-" in username and not force_default:
username, server_name = username.split("-", 1)

if server_name:
create_url = self.app.hub_url / "hub/api/users" / username / "servers" / server_name

print("create_url: ", create_url, flush=True)

async with session.post(create_url) as resp:
if resp.status == 201 or resp.status == 400:
Expand All @@ -68,6 +80,10 @@ async def start_user_server(self, session, username):
# We manually generate this, even though it's *bad*
# Mostly because when the server is already running, JupyterHub
# doesn't respond with the whole model!
if server_name:
print("Custom-named server started", username, server_name)
return self.app.hub_url / "user" / username / server_name
print("Server started", username, server_name)
return self.app.hub_url / "user" / username
elif resp.status == 202:
# Server start has been requested, now and potentially earlier,
Expand All @@ -78,35 +94,49 @@ async def start_user_server(self, session, username):
async with timeout(self.app.start_timeout):
notebook_url = None
self._conn.send_auth_banner("Starting your server...")
print("Server start requested", username, server_name)
while notebook_url is None:
# FIXME: Exponential backoff + make this configurable
await asyncio.sleep(0.5)
notebook_url = await self.get_user_server_url(
session, username
)
if server_name:
notebook_url = await self.get_user_server_url(
session, username + "/" + server_name, force_default
)
else:
notebook_url = await self.get_user_server_url(
session, username, force_default
)
self._conn.send_auth_banner(".")
self._conn.send_auth_banner("done!\n")
print("Server started", username, server_name)
return notebook_url
except asyncio.TimeoutError:
# Server didn't start on time!
self._conn.send_auth_banner("failed to start server on time!\n")
print("Server start timed out")
return None
elif resp.status == 403:
# Token is wrong!
print(f"The token seems to be wrong for {username}")
return None
elif resp.status == 404 and server_name:
print(f"Custom instance not found, probably {username} contains dash")
return 404
else:
# FIXME: Handle other cases that pop up
print("Unhandled response", resp.status)
resp.raise_for_status()

async def validate_password(self, username, token):
async def validate_password(self, username, token, force_default=False):
self.username = username
self.token = token

headers = {"Authorization": f"token {token}"}
async with ClientSession(headers=headers) as session:
notebook_url = await self.start_user_server(session, username)
notebook_url = await self.start_user_server(session, username, force_default)
if notebook_url is None:
return False
elif notebook_url == 404:
return self.validate_password(username, token, force_default=True)
else:
self.notebook_url = notebook_url
return True
Expand Down Expand Up @@ -225,13 +255,13 @@ class JupyterHubSSH(Application):
debug = Bool(
True,
help="""
Turn on debugg logging
Turn on debug logging
""",
config=True,
)

hub_url = Any(
"",
URL("http://proxy-public.jhub.svc.cluster.local"),
help="""
URL of JupyterHub's proxy to connect to.
Expand Down Expand Up @@ -269,7 +299,7 @@ def _hub_url_cast_string_to_yarl_url(self, proposal):
raise ValueError("hub_url must either be a string or a yarl.URL")

host_key_path = Unicode(
"",
"/etc/jupyterhub-ssh/config/hostKey",
help="""
Path to host's private SSH Key.
Expand Down

0 comments on commit 91a290f

Please sign in to comment.