fix(chalice): fixed Spot new refresh token

refactor(chalice): customizable-Spot-auth
tahayk · Aug 8, 2024 · 6f88b5d · 6f88b5d
1 parent a7cd363
commit 6f88b5d
Show file tree

Hide file tree

Showing 10 changed files with 46 additions and 32 deletions.
diff --git a/api/auth/auth_jwt.py b/api/auth/auth_jwt.py
@@ -36,7 +36,7 @@ async def __call__(self, request: Request) -> Optional[schemas.CurrentContext]:
             return await self.__process_refresh_call(request)
 
         elif request.url.path in ["/spot/refresh", "/spot/api/refresh"]:
-            return await self.__process_refresh_call(request)
+            return await self.__process_spot_refresh_call(request)
 
         else:
             credentials: HTTPAuthorizationCredentials = await super(JWTAuth, self).__call__(request)

diff --git a/api/chalicelib/core/authorizers.py b/api/chalicelib/core/authorizers.py
@@ -2,10 +2,9 @@
 
 import jwt
 from decouple import config
-from typing import Optional
+
 from chalicelib.core import tenants
 from chalicelib.core import users, spot
-
 from chalicelib.utils.TimeUTC import TimeUTC
 
 logger = logging.getLogger(__name__)
@@ -15,17 +14,25 @@ def get_supported_audience():
     return [users.AUDIENCE, spot.AUDIENCE]
 
 
-def jwt_authorizer(scheme: str, token: str, leeway=0) -> Optional[dict]:
+def is_spot_token(token: str) -> bool:
+    try:
+        decoded_token = jwt.decode(token, options={"verify_signature": False, "verify_exp": False})
+        audience = decoded_token.get("aud")
+        return audience == spot.AUDIENCE
+    except jwt.InvalidTokenError:
+        logger.error(f"Invalid token: {token}")
+        raise
+
+
+def jwt_authorizer(scheme: str, token: str, leeway=0) -> dict | None:
     if scheme.lower() != "bearer":
         return None
     try:
-        payload = jwt.decode(
-            token,
-            config("jwt_secret"),
-            algorithms=config("jwt_algorithm"),
-            audience=get_supported_audience(),
-            leeway=leeway
-        )
+        payload = jwt.decode(jwt=token,
+                             key=config("jwt_secret") if not is_spot_token(token) else config("JWT_SPOT_SECRET"),
+                             algorithms=config("jwt_algorithm"),
+                             audience=get_supported_audience(),
+                             leeway=leeway)
     except jwt.ExpiredSignatureError:
         logger.debug("! JWT Expired signature")
         return None
@@ -40,12 +47,11 @@ def jwt_refresh_authorizer(scheme: str, token: str):
     if scheme.lower() != "bearer":
         return None
     try:
-        payload = jwt.decode(
-            token,
-            config("JWT_REFRESH_SECRET"),
-            algorithms=config("jwt_algorithm"),
-            audience=get_supported_audience()
-        )
+        payload = jwt.decode(jwt=token,
+                             key=config("JWT_REFRESH_SECRET") if not is_spot_token(token) \
+                                 else config("JWT_SPOT_SECRET"),
+                             algorithms=config("jwt_algorithm"),
+                             audience=get_supported_audience())
     except jwt.ExpiredSignatureError:
         logger.debug("! JWT-refresh Expired signature")
         return None
@@ -56,34 +62,36 @@ def jwt_refresh_authorizer(scheme: str, token: str):
     return payload
 
 
-def generate_jwt(user_id, tenant_id, iat, aud):
+def generate_jwt(user_id, tenant_id, iat, aud, for_spot=False):
     token = jwt.encode(
         payload={
             "userId": user_id,
             "tenantId": tenant_id,
-            "exp": iat + config("JWT_EXPIRATION", cast=int),
+            "exp": iat + (config("JWT_EXPIRATION", cast=int) if not for_spot
+                          else config("JWT_SPOT_EXPIRATION", cast=int)),
             "iss": config("JWT_ISSUER"),
             "iat": iat,
             "aud": aud
         },
-        key=config("jwt_secret"),
+        key=config("jwt_secret") if not for_spot else config("JWT_SPOT_SECRET"),
         algorithm=config("jwt_algorithm")
     )
     return token
 
 
-def generate_jwt_refresh(user_id, tenant_id, iat, aud, jwt_jti):
+def generate_jwt_refresh(user_id, tenant_id, iat, aud, jwt_jti, for_spot=False):
     token = jwt.encode(
         payload={
             "userId": user_id,
             "tenantId": tenant_id,
-            "exp": iat + config("JWT_REFRESH_EXPIRATION", cast=int),
+            "exp": iat + (config("JWT_REFRESH_EXPIRATION", cast=int) if not for_spot
+                          else config("JWT_SPOT_REFRESH_EXPIRATION", cast=int)),
             "iss": config("JWT_ISSUER"),
             "iat": iat,
             "aud": aud,
             "jti": jwt_jti
         },
-        key=config("JWT_REFRESH_SECRET"),
+        key=config("JWT_REFRESH_SECRET") if not for_spot else config("JWT_SPOT_REFRESH_SECRET"),
         algorithm=config("jwt_algorithm")
     )
     return token

diff --git a/api/env.default b/api/env.default
@@ -32,6 +32,8 @@ JWT_REFRESH_EXPIRATION=604800
 JWT_REFRESH_SECRET="SET A RANDOM STRING HERE"
 JWT_SPOT_REFRESH_EXPIRATION=604800
 JWT_SPOT_REFRESH_SECRET="SET A RANDOM STRING HERE"
+JWT_SPOT_SECRET=SECRET
+JWT_SPOT_EXPIRATION=6000
 jwt_secret="SET A RANDOM STRING HERE"
 pg_dbname=postgres
 pg_host=postgresql.db.svc.cluster.local

diff --git a/api/env.dev b/api/env.dev
@@ -33,6 +33,8 @@ JWT_REFRESH_EXPIRATION=604800
 JWT_REFRESH_SECRET=SECRET2
 JWT_SPOT_REFRESH_EXPIRATION=604800
 JWT_SPOT_REFRESH_SECRET=SECRET3
+JWT_SPOT_SECRET=SECRET
+JWT_SPOT_EXPIRATION=6000
 jwt_secret=SECRET
 LOCAL_DEV=true
 LOGLEVEL=INFO

diff --git a/api/routers/core_dynamic.py b/api/routers/core_dynamic.py
@@ -75,14 +75,12 @@ def login_user(response: JSONResponse, spot: Optional[bool] = False, data: schem
             "user": r
         }
     }
+    response.set_cookie(key="refreshToken", value=refresh_token, path=COOKIE_PATH,
+                        max_age=refresh_token_max_age, secure=True, httponly=True)
     if spot:
         content["spotJwt"] = r.pop("spotJwt")
         spot_refresh_token = r.pop("spotRefreshToken")
         spot_refresh_token_max_age = r.pop("spotRefreshTokenMaxAge")
-
-    response.set_cookie(key="refreshToken", value=refresh_token, path=COOKIE_PATH,
-                        max_age=refresh_token_max_age, secure=True, httponly=True)
-    if spot:
         response.set_cookie(key="spotRefreshToken", value=spot_refresh_token, path="/api/spot/refresh",
                             max_age=spot_refresh_token_max_age, secure=True, httponly=True)
     return content

diff --git a/ee/api/env.default b/ee/api/env.default
@@ -50,6 +50,8 @@ JWT_REFRESH_EXPIRATION=604800
 JWT_REFRESH_SECRET="SET A RANDOM STRING HERE"
 JWT_SPOT_REFRESH_EXPIRATION=604800
 JWT_SPOT_REFRESH_SECRET="SET A RANDOM STRING HERE"
+JWT_SPOT_SECRET=SECRET
+JWT_SPOT_EXPIRATION=6000
 jwt_secret="SET A RANDOM STRING HERE"
 KAFKA_SERVERS=kafka.db.svc.cluster.local:9092
 KAFKA_USE_SSL=false

diff --git a/ee/api/env.dev b/ee/api/env.dev
@@ -49,6 +49,8 @@ JWT_REFRESH_EXPIRATION=604800
 JWT_REFRESH_SECRET=SECRET2
 JWT_SPOT_REFRESH_EXPIRATION=604800
 JWT_SPOT_REFRESH_SECRET=SECRET3
+JWT_SPOT_SECRET=SECRET
+JWT_SPOT_EXPIRATION=6000
 jwt_secret=SECRET
 KAFKA_SERVERS=127.0.0.1:9092
 KAFKA_USE_SSL=false

diff --git a/ee/api/routers/core_dynamic.py b/ee/api/routers/core_dynamic.py
@@ -81,14 +81,12 @@ def login_user(response: JSONResponse, spot: Optional[bool] = False, data: schem
             "user": r
         }
     }
+    response.set_cookie(key="refreshToken", value=refresh_token, path=COOKIE_PATH,
+                        max_age=refresh_token_max_age, secure=True, httponly=True)
     if spot:
         content["spotJwt"] = r.pop("spotJwt")
         spot_refresh_token = r.pop("spotRefreshToken")
         spot_refresh_token_max_age = r.pop("spotRefreshTokenMaxAge")
-
-    response.set_cookie(key="refreshToken", value=refresh_token, path=COOKIE_PATH,
-                        max_age=refresh_token_max_age, secure=True, httponly=True)
-    if spot:
         response.set_cookie(key="spotRefreshToken", value=spot_refresh_token, path="/api/spot/refresh",
                             max_age=spot_refresh_token_max_age, secure=True, httponly=True)
     return content

diff --git a/scripts/docker-compose/chalice.env b/scripts/docker-compose/chalice.env
@@ -26,4 +26,5 @@ version_number="${COMMON_VERSION}"
 CLUSTER_URL=""
 POD_NAMESPACE=""
 JWT_REFRESH_SECRET=${COMMON_JWT_REFRESH_SECRET}
-JWT_SPOT_REFRESH_SECRET=${COMMON_JWT_REFRESH_SECRET}
+JWT_SPOT_REFRESH_SECRET=${COMMON_JWT_REFRESH_SECRET}
+JWT_SPOT_SECRET=${COMMON_JWT_SPOT_SECRET}
diff --git a/scripts/docker-compose/common.env b/scripts/docker-compose/common.env
@@ -1,6 +1,7 @@
 COMMON_PROTOCOL="https"
 COMMON_DOMAIN_NAME="change_me_domain"
 COMMON_JWT_SECRET="change_me_jwt"
+COMMON_JWT_SPOT_SECRET="change_me_jwt"
 COMMON_JWT_REFRESH_SECRET="change_me_jwt_refresh"
 COMMON_S3_KEY="change_me_s3_key"
 COMMON_S3_SECRET="change_me_s3_secret"