Merge pull request #57 from maia-iyer/agent_jwks_validation

✨ Implement JWKS validation for the Slack Research Agent

Author: mrsabath

a2a/slack_researcher/README.md

@@ -10,15 +10,6 @@ The Slack Researcher Agent is designed to perform research tasks across Slack ch
 
 ### Configuration Variables
 
-## Running in Kagenti
-When deploying in the Kagenti UI - You will need to attach 3 environments to the agent deployment:
-
-![alt text](docs/environments.png)
-
-1. `ollama` or `openai` (Slack Researcher is looking for the values of `LLM_API_BASE` and `LLM_API_KEY` - Either ollama or any other OpenAI-compatible API will work)
-2. `mcp-slack` - This provides the value of the slack MCP server (`MCP_URL`)
-3. `slack-researcher` - This provides the remainder of the necessary configuration settings
-
 | Variable Name | Description | Required? | Default Value | 
 |---------------|-------------|-----------|---------------| 
 | LLM_API_BASE | The URL for OpenAI compatible API | Yes | `http://localhost:11434/v1` | 
@@ -30,3 +21,22 @@ When deploying in the Kagenti UI - You will need to attach 3 environments to the
 | MCP_URL | Endpoint where the Slack MCP server can be found | No |  "" | 
 | SERVICE_PORT | Port on which the service will run | Yes | `8000` |
 | LOG_LEVEL | Application log level | No | DEBUG |
+| JWKS_URL | Endpoint to obtain JWKS for token validation. Enables token validation | No | - |
+| ISSUER | Expected `iss` value of incoming bearer tokens | No | - |
+| AUDIENCE | Expected `aud` value of incoming bearer tokens | No | - |
+
+> **Note on Authorization configuration**
+> By default, no token validation is performed. To enable token validation, set `JWKS_URL`. 
+> If `ISSUER` is additionally set, the `iss` claim will be checked to equal this value. 
+> If `AUDIENCE` is additionally set, the `aud` claim will be checked to equal this value. 
+
+## Running in Kagenti
+When deploying in the Kagenti UI - You will need to attach 3 environments to the agent deployment:
+
+![alt text](docs/environments.png)
+
+1. `ollama` or `openai` (Slack Researcher is looking for the values of `LLM_API_BASE` and `LLM_API_KEY` - Either ollama or any other OpenAI-compatible API will work)
+2. `mcp-slack` - This provides the value of the slack MCP server (`MCP_URL`)
+3. `slack-researcher` - This provides the remainder of the necessary configuration settings
+
+NOTE if you are connecting to a tool that is OAuth-secured, you MUST set the `JWKS_URL` variable or the token will not be passed to the tool. 

a2a/slack_researcher/a2a_agent.py

@@ -20,35 +20,16 @@
 from a2a.types import AgentCapabilities, AgentCard, AgentSkill, TaskState, TextPart, SecurityScheme, HTTPAuthSecurityScheme
 from a2a.utils import new_agent_text_message, new_task
 
-from starlette.authentication import AuthCredentials, SimpleUser, AuthenticationBackend
 from starlette.middleware.authentication import AuthenticationMiddleware
 
 from slack_researcher.config import settings, Settings
 from slack_researcher.event import Event
 from slack_researcher.main import SlackAgent
+from slack_researcher.auth import on_auth_error, BearerAuthBackend
 
 logger = logging.getLogger(__name__)
 logging.basicConfig(level=logging.DEBUG, stream=sys.stdout, format='%(levelname)s: %(message)s')
 
-class BearerAuthBackend(AuthenticationBackend):
-    """ Very temporary demo to grab auth token and print it"""
-    async def authenticate(self, conn):
-        try:
-            auth = conn.headers.get("authorization")
-            if not auth or not auth.lower().startswith("bearer "):
-                print("No bearer token provided")
-                return
-            token = auth.split(" ", 1)[1]
-            print(f"TOKEN: {token}")
-
-            # Storing the token as the username - not a real life scenario - just demo-ing the passing of creds
-            user = SimpleUser(token)
-            return AuthCredentials(["authenticated"]), user
-        except Exception as e:
-            logger.error("Exception when attempting to obtain user token")
-            logger.error(e)
-    
-
 def get_agent_card(host: str, port: int):
     """Returns the Agent Card for the AG2 Agent."""
     capabilities = AgentCapabilities(streaming=True)
@@ -237,6 +218,10 @@ def run():
     )
 
     app = server.build()  # this returns a Starlette app
-    app.add_middleware(AuthenticationMiddleware, backend=BearerAuthBackend())
+    # if one of the auth variables is set, create middleware
+    # if none of them are set, ignore all authorization headers. No token validation will be performed
+    if not settings.JWKS_URI is None:
+        logging.info("JWKS_URI is set - using JWT Validation middleware")
+        app.add_middleware(AuthenticationMiddleware, backend=BearerAuthBackend(), on_error=on_auth_error)
 
     uvicorn.run(app, host="0.0.0.0", port=settings.SERVICE_PORT)

a2a/slack_researcher/pyproject.toml

@@ -9,6 +9,7 @@ dependencies = [
     "python-dotenv>=1.1.0",
     "a2a-sdk>=0.2.16",
     "python-keycloak>=5.5.1",
+    "authlib>=1.6.0",
 ]
 
 [tool.ruff]

a2a/slack_researcher/slack_researcher/auth.py

@@ -0,0 +1,104 @@
+import httpx
+import logging
+import sys
+
+from starlette.responses import JSONResponse
+from starlette.requests import Request
+from starlette.authentication import AuthCredentials, SimpleUser, AuthenticationBackend
+from starlette.authentication import AuthenticationError as StarletteAuthenticationError
+
+from authlib.jose import jwt
+from authlib.common.errors import AuthlibBaseError
+
+from slack_researcher.config import settings
+
+logger = logging.getLogger(__name__)
+logging.basicConfig(level=logging.DEBUG, stream=sys.stdout, format='%(levelname)s: %(message)s')
+
+def on_auth_error(request: Request, e: Exception):
+    status_code = 401
+    message = "Authentication failed"
+    headers = {"WWW-Authenticate": "Bearer"}
+
+    if isinstance(e, AuthenticationError):
+        status_code = e.status_code
+    message = str(e)
+
+    return JSONResponse(
+        {"error": message},
+        status_code = status_code,
+        headers=headers if status_code == 401 else None
+    )
+
+class AuthenticationError(StarletteAuthenticationError):
+    def __init__(self, message: str, status_code: int = 401):
+        self.status_code = status_code
+        super().__init__(message)
+
+class BearerAuthBackend(AuthenticationBackend):
+    def __init__(self):
+        if settings.JWKS_URI is None: # TODO implement oidc discovery in this case
+            raise Exception("JWKS_URI env var not set. ")
+        self.jwks_url = settings.JWKS_URI
+
+        self.claims_options = {}
+        if settings.AUDIENCE is None:
+            logger.debug(f"AUDIENCE env var not set. No audience check will be performed. ")
+        else:
+            self.claims_options["aud"] = {"essential": True, "value": settings.AUDIENCE}
+        if settings.ISSUER is None:
+            logger.debug(f"ISSUER env var no set. No issuer check will be performed")
+        else:
+            self.claims_options["iss"] = {"essential": True, "value": settings.ISSUER}
+
+    async def get_jwks(self):
+        logger.debug(f"Fetching JWKS from {self.jwks_url}")
+        jwks = None
+        try: 
+            async with httpx.AsyncClient() as client:
+                response = await client.get(self.jwks_url)
+                response.raise_for_status()
+                jwks = response.json() # TODO should probably use caching
+                return jwks
+        except httpx.HTTPStatusError as e:
+            logger.debug(f"Could not retrieve JWKS from {self.jwks_url}: {e}")
+            return None
+
+    async def get_token(self, conn):
+        logger.debug("Obtaining bearer token...")
+        headers = conn.headers
+        logger.debug(f"Headers obtained: {headers}")
+        auth = headers.get("authorization")
+        if not auth or not auth.lower().startswith("bearer "):
+            logger.error("Expected `Authorization: Bearer` access token; None provided. ")
+            return None
+        token = auth.split(" ", 1)[1]
+        return token
+
+    """This is run upon every received request"""
+    async def authenticate(self, conn):
+        # bypass authentication for agent card
+        if conn.scope.get("path") == "/.well-known/agent.json":
+            logger.debug("Bypassing authentication for public agent card path")
+            return None
+        
+        # extract token
+        token = await self.get_token(conn)
+        if token is None:
+            raise AuthenticationError(message = "Bearer token not found in Authorization header.")
+
+        # fetch jwks
+        jwks = await self.get_jwks()
+
+        try: 
+            # decode and validate claims
+            claims = jwt.decode(s=token, key=jwks, claims_options=self.claims_options)
+            claims.validate()
+            logger.debug("Token successfully validated.")
+
+            # return user
+            user = SimpleUser(token)
+            return AuthCredentials(["authenticated"]), user
+        except AuthlibBaseError as e:
+            logger.error(f"Token validation failed: {e}")
+            raise AuthenticationError(f"Invalid token: {e}, status_code=401")

a2a/slack_researcher/slack_researcher/config.py

@@ -5,7 +5,7 @@
 from pydantic_settings import BaseSettings
 from pydantic import model_validator
 from pydantic import Field
-from typing import Literal
+from typing import Literal, Optional
 
 
 class Settings(BaseSettings):
@@ -36,6 +36,20 @@ class Settings(BaseSettings):
     MCP_URL: str = Field(os.getenv("MCP_URL", "http://slack-tool:8000"), description="Endpoint for an option MCP server")
     SERVICE_PORT: int = Field(os.getenv("SERVICE_URL", 8000), description="Port on which the service will run.")
 
+    # auth variables
+    ISSUER: Optional[str] = Field(
+        os.getenv("ISSUER", None), 
+        description="The issuer for incoming JWT tokens"
+    )
+    JWKS_URI: Optional[str] = Field(
+        os.getenv("JWKS_URI", None),
+        description="Endpoint to obtain JWKS from auth server"
+    )
+    AUDIENCE: Optional[str] = Field(
+        os.getenv("AUDIENCE", None),
+        description="Expected audience value during resource validation"
+    )
+
     class Config:
         env_file = ".env"
         env_file_encoding = "utf-8"