JWT server cache fix and client response update (#8)

* Prevent caching incomplete jwt tokens in `validate_auth` method
Add token expiration to response data so client can manage refreshing

* Fix default `client_id` default to use a factory method
Add rate limit to `/auth/login` endpoint by origin IP address

* Add default values for test backwards-compat.

---------

Co-authored-by: Daniel McKnight <[email protected]>

neon_hana/app/routers/auth.py

@@ -24,7 +24,7 @@
 # NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
 # SOFTWARE,  EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
-from fastapi import APIRouter
+from fastapi import APIRouter, Request
 
 from neon_hana.app.dependencies import client_manager
 from neon_hana.schema.auth_requests import *
@@ -33,8 +33,10 @@
 
 
 @auth_route.post("/login")
-async def check_login(request: AuthenticationRequest) -> AuthenticationResponse:
-    return client_manager.check_auth_request(**dict(request))
+async def check_login(auth_request: AuthenticationRequest,
+                      request: Request) -> AuthenticationResponse:
+    return client_manager.check_auth_request(**dict(auth_request),
+                                             origin_ip=request.client.host)
 
 
 @auth_route.post("/refresh")

neon_hana/auth/client_manager.py

@@ -50,6 +50,7 @@ def __init__(self, config: dict):
         self._jwt_algo = "HS256"
 
     def _create_tokens(self, encode_data: dict) -> dict:
+        token_expiration = encode_data['expire']
         token = jwt.encode(encode_data, self._access_secret, self._jwt_algo)
         encode_data['expire'] = time() + self._refresh_token_lifetime
         encode_data['access_token'] = token
@@ -58,12 +59,25 @@ def _create_tokens(self, encode_data: dict) -> dict:
         return {"username": encode_data['username'],
                 "client_id": encode_data['client_id'],
                 "access_token": token,
-                "refresh_token": refresh}
+                "refresh_token": refresh,
+                "expiration": token_expiration}
 
     def check_auth_request(self, client_id: str, username: str,
-                           password: Optional[str] = None):
+                           password: Optional[str] = None,
+                           origin_ip: str = "127.0.0.1"):
         if client_id in self.authorized_clients:
+            print(f"Using cached client: {self.authorized_clients[client_id]}")
             return self.authorized_clients[client_id]
+
+        if not self.rate_limiter.get_all_buckets(f"auth{origin_ip}"):
+            self.rate_limiter.add_bucket(f"auth{origin_ip}",
+                                         TokenBucket(replenish_time=30,
+                                                     max_tokens=3))
+        if not self.rate_limiter.consume(f"auth{origin_ip}"):
+            raise HTTPException(status_code=429,
+                                detail=f"Too many auth requests from: "
+                                       f"{origin_ip}. Wait 30 seconds.")
+
         if username != "guest":
             # TODO: Validate password here
             pass
@@ -122,9 +136,6 @@ def validate_auth(self, token: str, origin_ip: str) -> bool:
             if auth['expire'] < time():
                 self.authorized_clients.pop(auth['client_id'], None)
                 return False
-            # Keep track of authorized client connections
-            self.authorized_clients[auth['client_id']] = auth
-            # TODO: Consider consuming an extra request for guest sessions
             return True
         except DecodeError:
             # Invalid token supplied

neon_hana/schema/auth_requests.py

@@ -27,13 +27,13 @@
 from typing import Optional
 from uuid import uuid4
 
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 
 
 class AuthenticationRequest(BaseModel):
     username: str = "guest"
     password: Optional[str] = None
-    client_id: str = str(uuid4())
+    client_id: str = Field(default_factory=lambda: str(uuid4()))
 
     model_config = {
         "json_schema_extra": {
@@ -48,14 +48,16 @@ class AuthenticationResponse(BaseModel):
     client_id: str
     access_token: str
     refresh_token: str
+    expiration: float
 
     model_config = {
         "json_schema_extra": {
             "examples": [{
                 "username": "guest",
                 "client_id": "be84ae66-f61c-4aac-a9af-b0da364b82b6",
                 "access_token": "<redacted>",
-                "refresh_token": "<redacted>"
+                "refresh_token": "<redacted>",
+                "expiration": 1706045776.4168212
             }]}}