-
Notifications
You must be signed in to change notification settings - Fork 29
/
toml.py
277 lines (226 loc) · 9.61 KB
/
toml.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
"""Traefik implementation
Custom proxy implementations can subclass :class:`Proxy`
and register in JupyterHub config:
.. sourcecode:: python
from mymodule import MyProxy
c.JupyterHub.proxy_class = MyProxy
Route Specification:
- A routespec is a URL prefix ([host]/path/), e.g.
'host.tld/path/' for host-based routing or '/path/' for default routing.
- Route paths should be normalized to always start and end with '/'
"""
# Copyright (c) Jupyter Development Team.
# Distributed under the terms of the Modified BSD License.
import json
import os
import asyncio
import string
import escapism
from traitlets import Any, default, Unicode
from . import traefik_utils
from jupyterhub.proxy import Proxy
from jupyterhub_traefik_proxy import TraefikProxy
class TraefikTomlProxy(TraefikProxy):
"""JupyterHub Proxy implementation using traefik and toml config file"""
mutex = Any()
@default("mutex")
def _default_mutex(self):
return asyncio.Lock()
toml_dynamic_config_file = Unicode(
"rules.toml", config=True, help="""traefik's dynamic configuration file"""
)
def __init__(self, **kwargs):
super().__init__(**kwargs)
try:
# Load initial routing table from disk
self.routes_cache = traefik_utils.load_routes(self.toml_dynamic_config_file)
except FileNotFoundError:
self.routes_cache = {}
finally:
if not self.routes_cache:
self.routes_cache = {"backends": {}, "frontends": {}}
async def _setup_traefik_static_config(self):
await super()._setup_traefik_static_config()
self.static_config["file"] = {"filename": "rules.toml", "watch": True}
try:
traefik_utils.persist_static_conf(
self.toml_static_config_file, self.static_config
)
try:
os.stat(self.toml_dynamic_config_file)
except FileNotFoundError:
# Make sure that the dynamic configuration file exists
self.log.info(
f"Creating the toml dynamic configuration file: {self.toml_dynamic_config_file}"
)
open(self.toml_dynamic_config_file, "a").close()
except IOError:
self.log.exception("Couldn't set up traefik's static config.")
raise
except:
self.log.error("Couldn't set up traefik's static config. Unexpected error:")
raise
def _start_traefik(self):
self.log.info("Starting traefik...")
try:
self._launch_traefik(config_type="toml")
except FileNotFoundError as e:
self.log.error(
"Failed to find traefik \n"
"The proxy can be downloaded from https://github.com/containous/traefik/releases/download."
)
raise
def _clean_resources(self):
try:
if self.should_start:
os.remove(self.toml_static_config_file)
os.remove(self.toml_dynamic_config_file)
except:
self.log.error("Failed to remove traefik's configuration files")
raise
def _get_route_unsafe(self, routespec):
safe = string.ascii_letters + string.digits + "_-"
escaped_routespec = escapism.escape(routespec, safe=safe)
result = {"data": "", "target": "", "routespec": routespec}
def get_target_data(d, to_find):
if to_find == "url":
key = "target"
else:
key = to_find
if result[key]:
return
for k, v in d.items():
if k == to_find:
result[key] = v
if isinstance(v, dict):
get_target_data(v, to_find)
for key, value in self.routes_cache["backends"].items():
if escaped_routespec in key:
get_target_data(value, "url")
for key, value in self.routes_cache["frontends"].items():
if escaped_routespec in key:
get_target_data(value, "data")
if not result["data"] and not result["target"]:
self.log.info("No route for {} found!".format(routespec))
result = None
else:
result["data"] = json.loads(result["data"])
return result
async def start(self):
"""Start the proxy.
Will be called during startup if should_start is True.
**Subclasses must define this method**
if the proxy is to be started by the Hub
"""
await super().start()
await self._wait_for_static_config(provider="file")
async def stop(self):
"""Stop the proxy.
Will be called during teardown if should_start is True.
**Subclasses must define this method**
if the proxy is to be started by the Hub
"""
await super().stop()
self._clean_resources()
async def add_route(self, routespec, target, data):
"""Add a route to the proxy.
**Subclasses must define this method**
Args:
routespec (str): A URL prefix ([host]/path/) for which this route will be matched,
e.g. host.name/path/
target (str): A full URL that will be the target of this route.
data (dict): A JSONable dict that will be associated with this route, and will
be returned when retrieving information about this route.
Will raise an appropriate Exception (FIXME: find what?) if the route could
not be added.
The proxy implementation should also have a way to associate the fact that a
route came from JupyterHub.
"""
routespec = self.validate_routespec(routespec)
backend_alias = traefik_utils.generate_alias(routespec, "backend")
frontend_alias = traefik_utils.generate_alias(routespec, "frontend")
data = json.dumps(data)
rule = traefik_utils.generate_rule(routespec)
async with self.mutex:
self.routes_cache["frontends"][frontend_alias] = {
"backend": backend_alias,
"passHostHeader": True,
"routes": {"test": {"rule": rule, "data": data}},
}
self.routes_cache["backends"][backend_alias] = {
"servers": {"server1": {"url": target, "weight": 1}}
}
traefik_utils.persist_routes(
self.toml_dynamic_config_file, self.routes_cache
)
if self.should_start:
try:
# Check if traefik was launched
pid = self.traefik_process.pid
except AttributeError:
self.log.error(
"You cannot add routes if the proxy isn't running! Please start the proxy: proxy.start()"
)
raise
try:
await self._wait_for_route(routespec, provider="file")
except TimeoutError:
self.log.error(
f"Is Traefik configured to watch {self.toml_dynamic_config_file}?"
)
raise
async def delete_route(self, routespec):
"""Delete a route with a given routespec if it exists.
**Subclasses must define this method**
"""
routespec = self.validate_routespec(routespec)
safe = string.ascii_letters + string.digits + "_-"
escaped_routespec = escapism.escape(routespec, safe=safe)
async with self.mutex:
for key, value in self.routes_cache["frontends"].items():
if escaped_routespec in key:
del self.routes_cache["frontends"][key]
break
for key, value in self.routes_cache["backends"].items():
if escaped_routespec in key:
del self.routes_cache["backends"][key]
break
traefik_utils.persist_routes(self.toml_dynamic_config_file, self.routes_cache)
async def get_all_routes(self):
"""Fetch and return all the routes associated by JupyterHub from the
proxy.
**Subclasses must define this method**
Should return a dictionary of routes, where the keys are
routespecs and each value is a dict of the form::
{
'routespec': the route specification ([host]/path/)
'target': the target host URL (proto://host) for this route
'data': the attached data dict for this route (as specified in add_route)
}
"""
all_routes = {}
async with self.mutex:
for key, value in self.routes_cache["frontends"].items():
escaped_routespec = "".join(key.split("_", 1)[1:])
routespec = escapism.unescape(escaped_routespec)
all_routes[routespec] = self._get_route_unsafe(routespec)
return all_routes
async def get_route(self, routespec):
"""Return the route info for a given routespec.
Args:
routespec (str):
A URI that was used to add this route,
e.g. `host.tld/path/`
Returns:
result (dict):
dict with the following keys::
'routespec': The normalized route specification passed in to add_route
([host]/path/)
'target': The target host for this route (proto://host)
'data': The arbitrary data dict that was passed in by JupyterHub when adding this
route.
None: if there are no routes matching the given routespec
"""
routespec = self.validate_routespec(routespec)
async with self.mutex:
return self._get_route_unsafe(routespec)