Merge pull request #220 from Point72/tkp/apik2

Revamp api key setup in anticipation of other mechanisms, add external api key validation mechanism

Author: timkpaine

csp_gateway/client/client.py

@@ -232,8 +232,8 @@ class GatewayClientConfig(BaseModel):
     host: str = "localhost"
     port: Optional[int] = Field(default=8000, ge=1, le=65535, description="Port number for the gateway server")
     api_route: str = "/api/v1"
-    authenticate: bool = False
     api_key: str = ""
+    bearer_token: Optional[str] = None
     return_type: ReturnType = Field(
         default=ReturnType.Raw,
         description="Determines how REST request responses should be returned. Options: 'raw' (JSON dict), 'pandas' (DataFrame), 'polars' (DataFrame), 'struct' (original type), 'wrapper' (ResponseWrapper object).",
@@ -251,13 +251,13 @@ def _promote_return_type(cls, v):
     def validate_config(self):
         if self.port is not None and self.port < 1:
             raise ValueError("Port must be a positive integer")
-        if self.api_key and not self.authenticate:
-            raise ValueError("API key must be provided if authentication is enabled")
         if self.host.startswith("http"):
             # Switch protocol to host
             protocol, host = self.host.split("://")
             self.__dict__["protocol"] = protocol
             self.__dict__["host"] = host
+        if self.bearer_token and self.api_key:
+            raise ValueError("Cannot provide both bearer_token and api_key. Choose one authentication method.")
         return self
 
     def __hash__(self):
@@ -403,6 +403,15 @@ class BaseGatewayClient(BaseModel):
         default=dict(follow_redirects=True), description="Additional arguments to pass to httpx requests (e.g., headers, auth, etc.)"
     )
 
+    # Additional initialization for bearer_token
+    def __init__(self, config: GatewayClientConfig = None, **kwargs) -> None:
+        # Exists for compatibility with positional argument instantiation
+        if config is None:
+            config = GatewayClientConfig()
+        if kwargs:
+            config = GatewayClientConfig(**{**config.model_dump(exclude_unset=True), **kwargs})
+        super().__init__(config=config)
+
     # openapi configureation
     _initialized: bool = PrivateAttr(default=False)
     _openapi_spec: Dict[Any, Any] = PrivateAttr(default=None)
@@ -424,16 +433,14 @@ class BaseGatewayClient(BaseModel):
     _event_loop: Optional[AbstractEventLoop] = PrivateAttr(default=None)
     _event_loop_thread: Optional[Thread] = PrivateAttr(default=None)
 
-    def __init__(self, config: GatewayClientConfig = None, **kwargs) -> None:
-        # Exists for compatibility with positional argument instantiation
-        if config is None:
-            config = GatewayClientConfig()
-        if kwargs:
-            config = GatewayClientConfig(**{**config.model_dump(exclude_unset=True), **kwargs})
-        super().__init__(config=config)
-
     @model_validator(mode="after")
     def validate_client(self):
+        # Set Authorization header if bearer_token is provided
+        if self.config.bearer_token:
+            headers = self.http_args.get("headers", {}).copy()
+            headers["Authorization"] = f"Bearer {self.config.bearer_token}"
+            self.http_args["headers"] = headers
+
         if self._event_loop is None:
             self._event_loop = _get_or_new_event_loop()
 
@@ -461,10 +468,12 @@ def _initializeStreaming(self) -> None:
     def _initialize(self) -> None:
         if not self._initialized:
             # grab openapi spec
+            openapi_url = f"{_host(self.config)}/openapi.json"
+            openapi_params = {"token": self.config.api_key} if self.config.api_key else None
             self._openapi_spec: Dict[Any, Any] = replace_refs(
                 cast(
                     Dict[Any, Any],
-                    GET(f"{_host(self.config)}/openapi.json", **self.http_args),
+                    GET(openapi_url, params=openapi_params, **self.http_args),
                 ).json(),
             )
 
@@ -504,9 +513,11 @@ def _buildpath(self, route: str) -> str:
 
     def _buildroute(self, route: str) -> str:
         url = f"{_host(self.config)}{self._buildpath(route)}"
-        if self.config.authenticate:
-            return url, {"token": self.config.api_key}
-        return url, {}
+        # If using api_key (not bearer_token), add as query param
+        extra_params = {}
+        if self.config.api_key:
+            extra_params["token"] = self.config.api_key
+        return url, extra_params
 
     def _api_path_and_route(self, route: str) -> str:
         return self.config.api_route + "/" + route
@@ -517,10 +528,11 @@ def _buildroutews(self, route: str) -> str:
             host = host.replace("http://", "ws://")
         elif host.startswith("https://"):
             host = host.replace("https://", "wss://")
-        if self.config.authenticate:
-            auth = f"?token={self.config.api_key}"
-        else:
-            auth = ""
+        # If using api_key (not bearer_token), add as query param
+        auth = ""
+        if self.config.api_key:
+            sep = "&" if "?" in host else "?"
+            auth = f"{sep}token={self.config.api_key}"
         return f"{host}{self.config.api_route}/{route}{auth}"
 
     def _handle_response(self, resp: Response, route: str) -> ResponseType:

csp_gateway/server/config/gateway/omnibus.yaml

@@ -2,7 +2,6 @@
 defaults:
   - _self_
 
-authenticate: ???
 port: ???
 
 modules:
@@ -32,16 +31,12 @@ modules:
     force_mount_all: True
   mount_websocket_routes:
     _target_: csp_gateway.MountWebSocketRoutes
-  mount_api_key_middleware:
-    _target_: csp_gateway.MountAPIKeyMiddleware
 
 gateway:
   _target_: csp_gateway.Gateway
   settings:
     PORT: ${port}
-    AUTHENTICATE: ${authenticate}
     UI: True
-    API_KEY: "12345"
   modules:
     - /modules/logfire
     - /modules/example_module

csp_gateway/server/demo/config/logfire.yaml

@@ -10,7 +10,7 @@ logfire:
   instrument_fastapi: True
   capture_logging: True
   log_level: DEBUG
-  
+
 publish_logfire:
   _target_: csp_gateway.server.modules.logging.PublishLogfire
   selection:
@@ -25,9 +25,7 @@ gateway:
   _target_: csp_gateway.Gateway
   settings:
     PORT: ${port}
-    AUTHENTICATE: ${authenticate}
     UI: True
-    API_KEY: "12345"
   modules:
     - /modules/example_module
     - /modules/mount_outputs
@@ -40,5 +38,4 @@ gateway:
 
 # csp-gateway-start --config-dir=csp_gateway/server/demo +config=logfire
 
-authenticate: False
 port: 8000

csp_gateway/server/demo/config/omnibus.yaml

@@ -5,5 +5,4 @@ defaults:
 
 # csp-gateway-start --config-dir=csp_gateway/server/demo +config=omnibus
 
-authenticate: False
 port: 8000

csp_gateway/server/gateway/csp/module.py

@@ -12,7 +12,7 @@
 from .channels import ChannelsType
 
 if typing.TYPE_CHECKING:
-    from csp_gateway.server import GatewayWebApp
+    from csp_gateway.server import GatewaySettings, GatewayWebApp
 
 
 class Module(BaseModel, Generic[ChannelsType], abc.ABC):
@@ -33,6 +33,8 @@ def connect(self, Channels: ChannelsType) -> None: ...
 
     def rest(self, app: "GatewayWebApp") -> None: ...
 
+    def info(self, settings: "GatewaySettings") -> Optional[str]: ...
+
     @abc.abstractmethod
     def shutdown(self) -> None: ...
 

csp_gateway/server/gateway/gateway.py

@@ -281,11 +281,11 @@ def start(
                 log.info("Launching web server on:")
                 url = f"http://{gethostname()}:{self.settings.PORT}"
 
-                if ui:
-                    if self.settings.AUTHENTICATE:
-                        log.info(f"\tUI: {url}?token={self.settings.API_KEY}")
-                    else:
-                        log.info(f"\tUI: {url}")
+                # Allow module sto log information at statup
+                for module in self.modules:
+                    _info = module.info(self.settings)
+                    if _info:
+                        log.info(_info)
 
                 log.info(f"\tDocs: {url}/docs")
                 log.info(f"\tDocs: {url}/redoc")

csp_gateway/server/middleware/__init__.py

@@ -1 +1,9 @@
 from .api_key import MountAPIKeyMiddleware
+from .api_key_external import MountExternalAPIKeyMiddleware
+from .base import AuthenticationMiddleware
+
+__all__ = (
+    "AuthenticationMiddleware",
+    "MountAPIKeyMiddleware",
+    "MountExternalAPIKeyMiddleware",
+)