mirror of
https://github.com/jupyterhub/jupyterhub.git
synced 2025-10-18 15:33:02 +00:00
313 lines
9.2 KiB
Python
313 lines
9.2 KiB
Python
"""Miscellaneous utilities"""
|
|
|
|
# Copyright (c) Jupyter Development Team.
|
|
# Distributed under the terms of the Modified BSD License.
|
|
|
|
from binascii import b2a_hex
|
|
import random
|
|
import errno
|
|
import hashlib
|
|
from hmac import compare_digest
|
|
import os
|
|
import socket
|
|
from threading import Thread
|
|
import uuid
|
|
import warnings
|
|
|
|
from tornado import web, gen, ioloop
|
|
from tornado.httpclient import AsyncHTTPClient, HTTPError
|
|
from tornado.log import app_log
|
|
|
|
|
|
def random_port():
|
|
"""Get a single random port."""
|
|
sock = socket.socket()
|
|
sock.bind(('', 0))
|
|
port = sock.getsockname()[1]
|
|
sock.close()
|
|
return port
|
|
|
|
|
|
# ISO8601 for strptime with/without milliseconds
|
|
ISO8601_ms = '%Y-%m-%dT%H:%M:%S.%fZ'
|
|
ISO8601_s = '%Y-%m-%dT%H:%M:%SZ'
|
|
|
|
|
|
def can_connect(ip, port):
|
|
"""Check if we can connect to an ip:port.
|
|
|
|
Return True if we can connect, False otherwise.
|
|
"""
|
|
if ip in {'', '0.0.0.0'}:
|
|
ip = '127.0.0.1'
|
|
try:
|
|
socket.create_connection((ip, port)).close()
|
|
except socket.error as e:
|
|
if e.errno not in {errno.ECONNREFUSED, errno.ETIMEDOUT}:
|
|
app_log.error("Unexpected error connecting to %s:%i %s", ip, port, e)
|
|
return False
|
|
else:
|
|
return True
|
|
|
|
@gen.coroutine
|
|
def exponential_backoff(
|
|
pass_func,
|
|
fail_message,
|
|
start_wait=0.2,
|
|
scale_factor=2,
|
|
max_wait=5,
|
|
timeout=10,
|
|
timeout_tolerance=0.1,
|
|
*args, **kwargs):
|
|
"""
|
|
Exponentially backoff until `pass_func` is true.
|
|
|
|
The `pass_func` function will wait with **exponential backoff** and
|
|
**random jitter** for as many needed iterations of the Tornado loop,
|
|
until reaching maximum `timeout` or truthiness. If `pass_func` is still
|
|
returning false at `timeout`, a `TimeoutError` will be raised.
|
|
|
|
The first iteration will begin with a wait time of `start_wait` seconds.
|
|
Each subsequent iteration's wait time will scale up by continuously
|
|
multiplying itself by `scale_factor`. This continues for each iteration
|
|
until `pass_func` returns true or an iteration's wait time has reached
|
|
the `max_wait` seconds per iteration.
|
|
|
|
`pass_func` may be a future, although that is not entirely recommended.
|
|
|
|
Parameters
|
|
----------
|
|
pass_func
|
|
function that is to be run
|
|
fail_message : str
|
|
message for a `TimeoutError`
|
|
start_wait : optional
|
|
initial wait time for the first iteration in seconds
|
|
scale_factor : optional
|
|
a multiplier to increase the wait time for each iteration
|
|
max_wait : optional
|
|
maximum wait time per iteration in seconds
|
|
timeout : optional
|
|
maximum time of total wait in seconds
|
|
timeout_tolerance : optional
|
|
a small multiplier used to add jitter to `timeout`'s deadline
|
|
*args, **kwargs
|
|
passed to `pass_func(*args, **kwargs)`
|
|
|
|
Returns
|
|
-------
|
|
value of `pass_func(*args, **kwargs)`
|
|
|
|
Raises
|
|
------
|
|
TimeoutError
|
|
If `pass_func` is still false at the end of the `timeout` period.
|
|
|
|
Notes
|
|
-----
|
|
See https://www.awsarchitectureblog.com/2015/03/backoff.html
|
|
for information about the algorithm and examples. We're using their
|
|
full Jitter implementation equivalent.
|
|
"""
|
|
loop = ioloop.IOLoop.current()
|
|
deadline = loop.time() + timeout
|
|
# add jitter to the deadline itself to prevent re-align of a bunch of
|
|
# timing out calls once the deadline is reached.
|
|
if timeout_tolerance:
|
|
tol = timeout_tolerance * timeout
|
|
deadline = random.uniform(deadline - tol, deadline + tol)
|
|
scale = 1
|
|
while True:
|
|
ret = yield gen.maybe_future(pass_func(*args, **kwargs))
|
|
# Truthy!
|
|
if ret:
|
|
return ret
|
|
remaining = deadline - loop.time()
|
|
if remaining < 0:
|
|
# timeout exceeded
|
|
break
|
|
# add some random jitter to improve performance
|
|
# this prevents overloading any single tornado loop iteration with
|
|
# too many things
|
|
dt = min(max_wait, remaining, random.uniform(0, start_wait * scale))
|
|
scale *= scale_factor
|
|
yield gen.sleep(dt)
|
|
raise TimeoutError(fail_message)
|
|
|
|
|
|
@gen.coroutine
|
|
def wait_for_server(ip, port, timeout=10):
|
|
"""Wait for any server to show up at ip:port."""
|
|
if ip in {'', '0.0.0.0'}:
|
|
ip = '127.0.0.1'
|
|
yield exponential_backoff(
|
|
lambda: can_connect(ip, port),
|
|
"Server at {ip}:{port} didn't respond in {timeout} seconds".format(ip=ip, port=port, timeout=timeout)
|
|
)
|
|
|
|
|
|
@gen.coroutine
|
|
def wait_for_http_server(url, timeout=10):
|
|
"""Wait for an HTTP Server to respond at url.
|
|
|
|
Any non-5XX response code will do, even 404.
|
|
"""
|
|
client = AsyncHTTPClient()
|
|
@gen.coroutine
|
|
def is_reachable():
|
|
try:
|
|
r = yield client.fetch(url, follow_redirects=False)
|
|
return r
|
|
except HTTPError as e:
|
|
if e.code >= 500:
|
|
# failed to respond properly, wait and try again
|
|
if e.code != 599:
|
|
# we expect 599 for no connection,
|
|
# but 502 or other proxy error is conceivable
|
|
app_log.warning(
|
|
"Server at %s responded with error: %s", url, e.code)
|
|
else:
|
|
app_log.debug("Server at %s responded with %s", url, e.code)
|
|
return e.response
|
|
except (OSError, socket.error) as e:
|
|
if e.errno not in {errno.ECONNABORTED, errno.ECONNREFUSED, errno.ECONNRESET}:
|
|
app_log.warning("Failed to connect to %s (%s)", url, e)
|
|
return False
|
|
re = yield exponential_backoff(
|
|
is_reachable,
|
|
"Server at {url} didn't respond in {timeout} seconds".format(url=url, timeout=timeout)
|
|
)
|
|
return re
|
|
|
|
|
|
# Decorators for authenticated Handlers
|
|
def auth_decorator(check_auth):
|
|
"""Make an authentication decorator.
|
|
|
|
I heard you like decorators, so I put a decorator
|
|
in your decorator, so you can decorate while you decorate.
|
|
"""
|
|
def decorator(method):
|
|
def decorated(self, *args, **kwargs):
|
|
check_auth(self)
|
|
return method(self, *args, **kwargs)
|
|
decorated.__name__ = method.__name__
|
|
decorated.__doc__ = method.__doc__
|
|
return decorated
|
|
|
|
decorator.__name__ = check_auth.__name__
|
|
decorator.__doc__ = check_auth.__doc__
|
|
return decorator
|
|
|
|
|
|
@auth_decorator
|
|
def token_authenticated(self):
|
|
"""Decorator for method authenticated only by Authorization token header
|
|
|
|
(no cookies)
|
|
"""
|
|
if self.get_current_user_token() is None:
|
|
raise web.HTTPError(403)
|
|
|
|
|
|
@auth_decorator
|
|
def authenticated_403(self):
|
|
"""Decorator for method to raise 403 error instead of redirect to login
|
|
|
|
Like tornado.web.authenticated, this decorator raises a 403 error
|
|
instead of redirecting to login.
|
|
"""
|
|
if self.get_current_user() is None:
|
|
raise web.HTTPError(403)
|
|
|
|
|
|
@auth_decorator
|
|
def admin_only(self):
|
|
"""Decorator for restricting access to admin users"""
|
|
user = self.get_current_user()
|
|
if user is None or not user.admin:
|
|
raise web.HTTPError(403)
|
|
|
|
|
|
# Token utilities
|
|
|
|
def new_token(*args, **kwargs):
|
|
"""Generator for new random tokens
|
|
|
|
For now, just UUIDs.
|
|
"""
|
|
return uuid.uuid4().hex
|
|
|
|
|
|
def hash_token(token, salt=8, rounds=16384, algorithm='sha512'):
|
|
"""Hash a token, and return it as `algorithm:salt:hash`.
|
|
|
|
If `salt` is an integer, a random salt of that many bytes will be used.
|
|
"""
|
|
h = hashlib.new(algorithm)
|
|
if isinstance(salt, int):
|
|
salt = b2a_hex(os.urandom(salt))
|
|
if isinstance(salt, bytes):
|
|
bsalt = salt
|
|
salt = salt.decode('utf8')
|
|
else:
|
|
bsalt = salt.encode('utf8')
|
|
btoken = token.encode('utf8', 'replace')
|
|
h.update(bsalt)
|
|
for i in range(rounds):
|
|
h.update(btoken)
|
|
digest = h.hexdigest()
|
|
|
|
return "{algorithm}:{rounds}:{salt}:{digest}".format(**locals())
|
|
|
|
|
|
def compare_token(compare, token):
|
|
"""Compare a token with a hashed token.
|
|
|
|
Uses the same algorithm and salt of the hashed token for comparison.
|
|
"""
|
|
algorithm, srounds, salt, _ = compare.split(':')
|
|
hashed = hash_token(token, salt=salt, rounds=int(srounds), algorithm=algorithm).encode('utf8')
|
|
compare = compare.encode('utf8')
|
|
if compare_digest(compare, hashed):
|
|
return True
|
|
return False
|
|
|
|
|
|
def url_path_join(*pieces):
|
|
"""Join components of url into a relative url.
|
|
|
|
Use to prevent double slash when joining subpath. This will leave the
|
|
initial and final / in place.
|
|
|
|
Copied from `notebook.utils.url_path_join`.
|
|
"""
|
|
initial = pieces[0].startswith('/')
|
|
final = pieces[-1].endswith('/')
|
|
stripped = [ s.strip('/') for s in pieces ]
|
|
result = '/'.join(s for s in stripped if s)
|
|
|
|
if initial:
|
|
result = '/' + result
|
|
if final:
|
|
result = result + '/'
|
|
if result == '//':
|
|
result = '/'
|
|
|
|
return result
|
|
|
|
|
|
def default_server_name(user):
|
|
"""Return the default name for a new server for a given user.
|
|
|
|
Will be the first available integer string, e.g. '1' or '2'.
|
|
"""
|
|
existing_names = set(user.spawners)
|
|
# if there are 5 servers, count from 1 to 6
|
|
for n in range(1, len(existing_names) + 2):
|
|
name = str(n)
|
|
if name not in existing_names:
|
|
return name
|
|
raise RuntimeError("It should be impossible to get here")
|
|
|