Skip to content

Commit 6c7ccfd

Browse files
fix: Fix Session.refresh() KeyError by sealing session client-side (#673)
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
1 parent 0bde40a commit 6c7ccfd

2 files changed

Lines changed: 529 additions & 43 deletions

File tree

src/workos/session.py

Lines changed: 98 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -336,18 +336,11 @@ def refresh(
336336
)
337337

338338
try:
339-
# Use raw dict request because the generated AuthenticateResponse
340-
# doesn't include sealed_session, and the request body needs the
341-
# session parameter which isn't in the generated request models.
342339
body: Dict[str, Any] = {
343340
"grant_type": "refresh_token",
344341
"client_id": self._client.client_id,
345342
"client_secret": self._client._api_key,
346343
"refresh_token": session["refresh_token"],
347-
"session": {
348-
"seal_session": True,
349-
"cookie_password": effective_cookie_password,
350-
},
351344
}
352345
if organization_id is not None:
353346
body["organization_id"] = organization_id
@@ -358,31 +351,64 @@ def refresh(
358351
body=body,
359352
)
360353

361-
self.session_data = str(auth_response["sealed_session"])
362-
self.cookie_password = effective_cookie_password
363-
364-
signing_key = self.jwks.get_signing_key_from_jwt(
365-
auth_response["access_token"]
366-
)
367-
decoded = jwt.decode(
368-
auth_response["access_token"],
369-
signing_key.key,
370-
algorithms=self._JWK_ALGORITHMS,
371-
options={"verify_aud": False},
372-
leeway=self._client._jwt_leeway,
354+
access_token = auth_response.get("access_token")
355+
refresh_token = auth_response.get("refresh_token")
356+
if not access_token or not refresh_token:
357+
return RefreshWithSessionCookieErrorResponse(
358+
authenticated=False,
359+
reason=AuthenticateWithSessionCookieFailureReason.REFRESH_DENIED,
360+
)
361+
362+
user = auth_response.get("user") or {}
363+
impersonator = auth_response.get("impersonator")
364+
365+
try:
366+
signing_key = self.jwks.get_signing_key_from_jwt(access_token)
367+
decoded = jwt.decode(
368+
access_token,
369+
signing_key.key,
370+
algorithms=self._JWK_ALGORITHMS,
371+
options={"verify_aud": False},
372+
leeway=self._client._jwt_leeway,
373+
)
374+
except (
375+
jwt.exceptions.InvalidTokenError,
376+
jwt.exceptions.PyJWKClientError,
377+
):
378+
return RefreshWithSessionCookieErrorResponse(
379+
authenticated=False,
380+
reason=AuthenticateWithSessionCookieFailureReason.INVALID_JWT,
381+
)
382+
383+
session_id = decoded.get("sid")
384+
if not session_id:
385+
return RefreshWithSessionCookieErrorResponse(
386+
authenticated=False,
387+
reason=AuthenticateWithSessionCookieFailureReason.INVALID_JWT,
388+
)
389+
390+
new_sealed = seal_session_from_auth_response(
391+
access_token=access_token,
392+
refresh_token=refresh_token,
393+
user=user,
394+
impersonator=impersonator,
395+
cookie_password=effective_cookie_password,
373396
)
374397

398+
self.session_data = new_sealed
399+
self.cookie_password = effective_cookie_password
400+
375401
return RefreshWithSessionCookieSuccessResponse(
376402
authenticated=True,
377-
sealed_session=str(auth_response["sealed_session"]),
378-
session_id=decoded["sid"],
403+
sealed_session=new_sealed,
404+
session_id=session_id,
379405
organization_id=decoded.get("org_id"),
380406
role=decoded.get("role"),
381407
roles=decoded.get("roles"),
382408
permissions=decoded.get("permissions"),
383409
entitlements=decoded.get("entitlements"),
384-
user=auth_response.get("user"),
385-
impersonator=auth_response.get("impersonator"),
410+
user=user,
411+
impersonator=impersonator,
386412
feature_flags=decoded.get("feature_flags"),
387413
)
388414
except Exception as e:
@@ -523,10 +549,6 @@ async def refresh(
523549
"client_id": self._client.client_id,
524550
"client_secret": self._client._api_key,
525551
"refresh_token": session["refresh_token"],
526-
"session": {
527-
"seal_session": True,
528-
"cookie_password": effective_cookie_password,
529-
},
530552
}
531553
if organization_id is not None:
532554
body["organization_id"] = organization_id
@@ -537,31 +559,64 @@ async def refresh(
537559
body=body,
538560
)
539561

540-
self.session_data = str(auth_response["sealed_session"])
541-
self.cookie_password = effective_cookie_password
542-
543-
signing_key = self.jwks.get_signing_key_from_jwt(
544-
auth_response["access_token"]
545-
)
546-
decoded = jwt.decode(
547-
auth_response["access_token"],
548-
signing_key.key,
549-
algorithms=self._JWK_ALGORITHMS,
550-
options={"verify_aud": False},
551-
leeway=self._client._jwt_leeway,
562+
access_token = auth_response.get("access_token")
563+
refresh_token = auth_response.get("refresh_token")
564+
if not access_token or not refresh_token:
565+
return RefreshWithSessionCookieErrorResponse(
566+
authenticated=False,
567+
reason=AuthenticateWithSessionCookieFailureReason.REFRESH_DENIED,
568+
)
569+
570+
user = auth_response.get("user") or {}
571+
impersonator = auth_response.get("impersonator")
572+
573+
try:
574+
signing_key = self.jwks.get_signing_key_from_jwt(access_token)
575+
decoded = jwt.decode(
576+
access_token,
577+
signing_key.key,
578+
algorithms=self._JWK_ALGORITHMS,
579+
options={"verify_aud": False},
580+
leeway=self._client._jwt_leeway,
581+
)
582+
except (
583+
jwt.exceptions.InvalidTokenError,
584+
jwt.exceptions.PyJWKClientError,
585+
):
586+
return RefreshWithSessionCookieErrorResponse(
587+
authenticated=False,
588+
reason=AuthenticateWithSessionCookieFailureReason.INVALID_JWT,
589+
)
590+
591+
session_id = decoded.get("sid")
592+
if not session_id:
593+
return RefreshWithSessionCookieErrorResponse(
594+
authenticated=False,
595+
reason=AuthenticateWithSessionCookieFailureReason.INVALID_JWT,
596+
)
597+
598+
new_sealed = seal_session_from_auth_response(
599+
access_token=access_token,
600+
refresh_token=refresh_token,
601+
user=user,
602+
impersonator=impersonator,
603+
cookie_password=effective_cookie_password,
552604
)
553605

606+
self.session_data = new_sealed
607+
self.cookie_password = effective_cookie_password
608+
554609
return RefreshWithSessionCookieSuccessResponse(
555610
authenticated=True,
556-
sealed_session=str(auth_response["sealed_session"]),
557-
session_id=decoded["sid"],
611+
sealed_session=new_sealed,
612+
session_id=session_id,
558613
organization_id=decoded.get("org_id"),
559614
role=decoded.get("role"),
560615
roles=decoded.get("roles"),
561616
permissions=decoded.get("permissions"),
562617
entitlements=decoded.get("entitlements"),
563-
user=auth_response.get("user"),
564-
impersonator=auth_response.get("impersonator"),
618+
user=user,
619+
impersonator=impersonator,
565620
feature_flags=decoded.get("feature_flags"),
566621
)
567622
except Exception as e:

0 commit comments

Comments
 (0)