Coverage for src / quber / playground / auth.py: 39%

1"""WorkOS sign-in for the hosted playground. 

2 

3Every request the app serves passes the login gate. A request carrying a live 

4session cookie for the configured organization goes through; otherwise the 

5gate redirects a browser to ``/login`` or refuses an API call with 401. A few 

6paths are open without a session: the sign-in routes themselves, the health 

7route the load balancer and the startup page poll, and the static assets, 

8which are the app's code and carry no document data. 

9 

10Four routes come with the gate. ``/login`` sends the browser to WorkOS, 

11``/callback`` finishes the sign-in WorkOS sends back and sets the cookie, 

12``/logout`` clears the cookie and ends the WorkOS session, and ``/api/me`` 

13tells the client who is signed in so it can offer sign-out. The sign-in 

14itself, the organization check and the cookie all come from 

15``quber.playground.session``, which the wake function shares: a person who 

16signs in while no task is running gets the same cookie from the function, 

17and the gate here accepts it once the task is up. 

18 

19The gate is a pure ASGI middleware rather than a Starlette 

20``BaseHTTPMiddleware`` so the answer and batch streams pass through untouched. 

21It checks the organization on every request, not only at sign-in, so a 

22cookie issued under another organization, or before the setting changed, is 

23refused like an expired one. 

24 

25``install`` decides whether the sign-in is on. All four settings unset means 

26the developer host, where the app runs open. All four set means the hosted 

27playground. A partial set is a misconfiguration and is refused at startup, 

28because the alternative is a hosted app that silently runs open. 

29 

30The gate is also where signed-in use is counted. Every request that passes 

31with a valid session is one unit of activity, and ``ActivityMeter`` reports 

32the count to CloudWatch once a minute when a namespace is configured. The 

33hosted service's idle alarm watches that metric rather than the load 

34balancer's request count, because a public hostname is never quiet: internet 

35scanners reach the sign-in page every few minutes, and only a session tells a 

36person apart from them. 

37""" 

38 

39from __future__ import annotations 

40 

41import threading 

42import time 

43from typing import Optional 

44from urllib.parse import urlencode 

45 

46from fastapi import FastAPI, Request 

47from loguru import logger 

48from starlette.responses import HTMLResponse, JSONResponse, RedirectResponse, Response 

49from starlette.types import ASGIApp, Receive, Scope, Send 

50 

51from quber.playground import session 

52from quber.settings import PlaygroundSettings 

53 

54OPEN_PATHS = frozenset({"/login", "/callback", "/logout", "/healthz"}) 

55OPEN_PREFIXES = ("/static/",) 

56ACTIVITY_METRIC = "AuthenticatedRequests" 

57ACTIVITY_INTERVAL_SECONDS = 60 

58# The app serves the mark from its static files; the wake function inlines it. 

59MARK_SRC = "/static/qubera-mark-160.png" 

60 

61 

62class ActivityMeter: 

63 """Counts requests that carried a valid session and reports the count to 

64 CloudWatch once a minute. A minute with nothing to report is reported as 

65 zero, so the metric has a datapoint whenever a task is up; the alarm's 

66 missing-data rule covers the time nothing is up at all.""" 

67 

68 def __init__(self, namespace: str, service: str) -> None: 

69 self.namespace = namespace 

70 self.service = service 

71 self.count = 0 

72 self.lock = threading.Lock() 

73 threading.Thread(target=self.run, name="activity-meter", daemon=True).start() 

74 

75 def record(self) -> None: 

76 with self.lock: 

77 self.count += 1 

78 

79 def run(self) -> None: 

80 import boto3 

81 

82 client = boto3.client("cloudwatch") 

83 while True: 

84 time.sleep(ACTIVITY_INTERVAL_SECONDS) 

85 with self.lock: 

86 count, self.count = self.count, 0 

87 try: 

88 client.put_metric_data( 

89 Namespace=self.namespace, 

90 MetricData=[ 

91 { 

92 "MetricName": ACTIVITY_METRIC, 

93 "Dimensions": [{"Name": "Service", "Value": self.service}], 

94 "Value": count, 

95 "Unit": "Count", 

96 } 

97 ], 

98 ) 

99 except Exception as exc: 

100 logger.warning("activity metric not sent: {}", exc) 

101 

102 

103class LoginGate: 

104 def __init__( 

105 self, 

106 app: ASGIApp, 

107 *, 

108 secret: str, 

109 organization_id: str, 

110 meter: Optional[ActivityMeter] = None, 

111 ) -> None: 

112 self.app = app 

113 self.secret = secret 

114 self.organization_id = organization_id 

115 self.meter = meter 

116 

117 async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: 

118 if scope["type"] != "http": 

119 await self.app(scope, receive, send) 

120 return 

121 path = scope["path"] 

122 if path in OPEN_PATHS or path.startswith(OPEN_PREFIXES): 

123 await self.app(scope, receive, send) 

124 return 

125 request = Request(scope, receive) 

126 identity = session.verify(self.secret, self.organization_id, request.cookies.get(session.COOKIE)) 

127 if identity is not None: 

128 if self.meter is not None: 

129 self.meter.record() 

130 scope.setdefault("state", {})["identity"] = identity 

131 await self.app(scope, receive, send) 

132 return 

133 if path.startswith("/api/"): 

134 response = JSONResponse({"detail": "login required"}, status_code=401) 

135 else: 

136 query = scope.get("query_string", b"").decode() 

137 target = path + ("?" + query if query else "") 

138 response = RedirectResponse("/login?" + urlencode({"next": target}), status_code=302) 

139 await response(scope, receive, send) 

140 

141 

142def request_base(request: Request) -> str: 

143 """The site the browser used. The load balancer terminates TLS, so the 

144 browser's scheme arrives in the forwarded header.""" 

145 scheme = request.headers.get("x-forwarded-proto") or request.url.scheme 

146 return session.base_url(scheme, request.headers.get("host") or request.url.netloc) 

147 

148 

149def secure(request: Request) -> bool: 

150 """Whether to mark the cookie Secure: it is then sent back only over HTTPS, 

151 which is every hosted request.""" 

152 return request.headers.get("x-forwarded-proto") == "https" 

153 

154 

155def install(app: FastAPI, settings: PlaygroundSettings) -> bool: 

156 """Put the sign-in in front of ``app`` when settings ask for it. 

157 

158 Returns whether the sign-in is on. Raises when the four settings are only 

159 partly set, so a hosted app never starts open by accident. ``/api/me`` is 

160 added either way, answering an empty object on the developer host. 

161 """ 

162 values = ( 

163 settings.workos_api_key, 

164 settings.workos_client_id, 

165 settings.workos_organization_id, 

166 settings.session_secret, 

167 ) 

168 if not any(values): 

169 

170 def me_open() -> dict[str, str]: 

171 return {} 

172 

173 app.add_api_route("/api/me", me_open, methods=["GET"], include_in_schema=False) 

174 return False 

175 if not all(values): 

176 raise RuntimeError( 

177 "WORKOS_API_KEY, WORKOS_CLIENT_ID, WORKOS_ORGANIZATION_ID and PLAYGROUND_SESSION_SECRET " 

178 "must be set together; the sign-in is on only when all four are present." 

179 ) 

180 api_key, client_id, organization_id, secret = (str(v) for v in values) 

181 config = session.WorkOS( 

182 api_key=api_key, client_id=client_id, organization_id=organization_id, session_secret=secret 

183 ) 

184 meter = ( 

185 ActivityMeter(settings.activity_namespace, "quber-playground") 

186 if settings.activity_namespace 

187 else None 

188 ) 

189 

190 def login(request: Request, next: Optional[str] = None) -> HTMLResponse: 

191 return HTMLResponse(session.login_page(config, request_base(request), next, MARK_SRC)) 

192 

193 # Plain functions, so the call to WorkOS runs on the thread pool and never 

194 # blocks the event loop. 

195 def callback(request: Request, code: Optional[str] = None, state: Optional[str] = None) -> Response: 

196 outcome = session.complete_sign_in(config, request_base(request), code, state, MARK_SRC) 

197 if outcome.token is None or outcome.location is None: 

198 logger.info("sign-in not accepted: status {}", outcome.status) 

199 return HTMLResponse(outcome.page or "", status_code=outcome.status) 

200 response = RedirectResponse(outcome.location, status_code=303) 

201 response.headers.append("set-cookie", session.cookie_header(outcome.token, secure=secure(request))) 

202 return response 

203 

204 def logout(request: Request) -> RedirectResponse: 

205 identity = session.signed_identity(secret, request.cookies.get(session.COOKIE)) 

206 base = request_base(request) 

207 target = session.logout_url(identity.sid, base) if identity and identity.sid else "/" 

208 response = RedirectResponse(target, status_code=302) 

209 response.headers.append("set-cookie", session.clear_cookie_header(secure=secure(request))) 

210 return response 

211 

212 def me(request: Request) -> dict[str, str]: 

213 identity: Optional[session.Identity] = getattr(request.state, "identity", None) 

214 return {"email": identity.email} if identity is not None else {} 

215 

216 app.add_api_route("/login", login, methods=["GET"], response_class=HTMLResponse, include_in_schema=False) 

217 app.add_api_route("/callback", callback, methods=["GET"], include_in_schema=False) 

218 app.add_api_route("/logout", logout, methods=["GET"], include_in_schema=False) 

219 app.add_api_route("/api/me", me, methods=["GET"], include_in_schema=False) 

220 app.add_middleware(LoginGate, secret=secret, organization_id=organization_id, meter=meter) 

221 return True