Implement OAuth endpoints

README.md

@@ -54,6 +54,9 @@ Configure the server in the client:
     }
 ```
 
+4. Ask about your clusters:
+![Example prompt asking about a cluster](images/cluster-prompt-example.png)
+
 ### Providing the Offline Token via Request Header
 
 If you do not set the `OFFLINE_TOKEN` environment variable, you can provide the token as a request header.
@@ -69,9 +72,6 @@ When configuring your MCP client, add the `OCM-Offline-Token` header:
     }
 ```
 
-4. Ask about your clusters:
-![Example prompt asking about a cluster](images/cluster-prompt-example.png)
-
 ## Available Tools
 
 The MCP server provides the following tools for interacting with the OpenShift Assisted Installer:
@@ -135,3 +135,51 @@ The MCP server provides the following tools for interacting with the OpenShift A
 * **Create a cluster**: "Create a new cluster named 'my-cluster' with OpenShift 4.14 and base domain 'example.com'"
 * **Check cluster events**: "What events happened on cluster abc123?"
 * **Install a cluster**: "Start the installation for cluster abc123"
+
+## Authorization Options
+
+The server supports multiple authorization methods for accessing the Assisted Installer API. The
+method used depends on the environment variables and headers you provide. The following methods are
+checked in order of priority; the first one that succeeds will be used, and the rest will be
+ignored:
+
+### 1. Access token in the `Authorization` request header
+
+If the `Authorization` request header contains a bearer token, it will be passed directly to the
+Assisted Installer API. In this case, the OAuth flow will not be triggered, and any values provided
+in the `OFFLINE_TOKEN` environment variable or the `OCM-Offline-Token` request header will be
+ignored.
+
+### 2. OAuth flow
+
+If the `OAUTH_ENABLED` environment variable is set to `true`, the server will use a subset of the
+OAuth protocol that MCP clients (such as the one in VS Code) use for authentication. When you
+attempt to connect, the MCP client will open a browser window where you can enter your credentials.
+The client will then request an access token, which the server will use to authenticate requests to
+the Assisted Installer API.
+
+When using this authentication method, the `OFFLINE_TOKEN` environment variable and the
+`OCM-Offline-Token` header will be ignored.
+
+You can configure the OAuth authorization server and client identifier using the `OAUTH_URL` and
+`OAUTH_CLIENT` environment variables. The default values are:
+
+- `OAUTH_URL`: `https://sso.redhat.com/auth/realms/redhat-external`
+- `OAUTH_CLIENT`: `cloud-services`
+
+The `SELF_URL` environment variable specifies the base URL that the server uses to construct URLs
+referencing itself. For example, when OAuth is enabled, the server will generate the dynamic client
+registration URL by appending `/oauth/register` to this base URL. The default value is
+`http://localhost:8000`, but in production environments, it should be set to the actual URL of the
+server as accessible to clients. For instance, if the server is accessed through a reverse proxy
+using HTTPS and the host `my.host.com`, the value should be set to `https://my.host.com`.
+
+### 3. Offline token via environment variable
+
+If you set the `OFFLINE_TOKEN` environment variable, the server will use this offline token to
+request an access token, which will then be used to call the Assisted Installer API.
+
+### 4. Offline token via request header
+
+If the `OCM-Offline-Token` request header is set, the server will use it to request an access token,
+and will then use that access token to call the Assisted Installer API.

oauth.py

@@ -0,0 +1,140 @@
+"""
+This module contains the Starlette middleware that implemnts OAuth authorization.
+"""
+
+import httpx
+import starlette.middleware
+import starlette.middleware.base
+import starlette.requests
+import starlette.responses
+import starlette.types
+
+
+class Middleware(starlette.middleware.base.BaseHTTPMiddleware):
+    """
+    This middleware implements the OAuth metadata and registration endpoints that MCP clients
+    will try to use when the server responds with a 401 status code.
+    """
+
+    def __init__(
+        self,
+        app: starlette.types.ASGIApp,
+        self_url: str,
+        oauth_url: str,
+        oauth_client: str,
+    ):
+        """
+        Creates a new OAuth middleware.
+
+        Args:
+            app (starlette.types.ASGIApp): The starlette application.
+            self_url (str): Base URL of the service, as seen by clients.
+            oauth_url (str): Base URL of the authorization server.
+            oauth_client (str): The client identifier.
+        """
+        super().__init__(app=app)
+        self._self_url = self_url
+        self._oauth_url = oauth_url
+        self._oauth_client = oauth_client
+
+    async def dispatch(
+        self,
+        request: starlette.requests.Request,
+        call_next: starlette.middleware.base.RequestResponseEndpoint,
+    ) -> starlette.responses.Response:
+        """
+        Dispatches the request, calling the OAuth handlers or else the protected application.
+        """
+        # The OAuth endpoints don't require authentication:
+        method = request.method
+        path = request.url.path
+        if method == "GET" and path == "/.well-known/oauth-protected-resource":
+            return await self._resource(request)
+        if method == "GET" and path == "/.well-known/oauth-authorization-server":
+            return await self._metadata(request)
+        if method == "POST" and path == "/oauth/register":
+            return await self._register(request)
+
+        # The rest of the endpoints do require authentication. Note that we are not validating the
+        # bearer token, just requiring the authorization header, so that the client will receive
+        # the 401 response code and trigger the OAuth flow.
+        auth = request.headers.get("authorization")
+        if auth is None:
+            resource_url = f"{self._self_url}/.well-known/oauth-protected-resource"
+            return starlette.responses.Response(
+                status_code=401,
+                headers={
+                    "WWW-Authenticate": f"Bearer resource_metadata=\"{resource_url}\"",
+                },
+            )
+
+        return await call_next(request)
+
+    async def _resource(self, request: starlette.requests.Request) -> starlette.responses.Response:
+        """
+        This method implements the OAuth protected resource endpoint.
+        """
+        return starlette.responses.JSONResponse(
+            content={
+                "resource": self._self_url,
+                "authorization_servers": [
+                    self._self_url,
+                ],
+                "bearer_methods_supported": [
+                    "header",
+                ],
+                "scopes_supported": [
+                    "openid",
+                    "api.ocm",
+                ],
+            }
+        )
+
+    async def _metadata(self, request: starlette.requests.Request) -> starlette.responses.Response:
+        """
+        This method implements the OAuth metadata endpoint. It gets the metadata from our real authorization
+        server, and replaces a few things that are needed to satisfy MCP clients.
+        """
+        # Get the metadata from the real authorization service:
+        try:
+            async with httpx.AsyncClient() as client:
+                response = await client.get(
+                    url=f"{self._oauth_url}/.well-known/oauth-authorization-server",
+                    timeout=10,
+                )
+                response.raise_for_status()
+                body = response.json()
+        except (httpx.RequestError, httpx.HTTPStatusError):
+            return starlette.responses.Response(status_code=503)
+
+        # The MCP clients will want to dynamically register the client, but we don't want that because our
+        # authorization server doesn't allow us to do it. So we replace the registration endpoint with our
+        # own, where we can return a fake response to make the MCP clients happy.
+        body["registration_endpoint"] = f"{self._self_url}/oauth/register"
+
+        # The MCP clients also try to request all the scopes listed in the metadata, but our authorization
+        # server returns a lot of scopes, and most of them will be rejected for our client. So we replace
+        # that large list with a much smaller list containing only the scopes that we need.
+        body["scopes_supported"] = [
+            "openid",
+            "api.ocm",
+        ]
+
+        # Return the modified metadata:
+        return starlette.responses.JSONResponse(
+            content=body,
+        )
+
+    async def _register(self, request: starlette.requests.Request) -> starlette.responses.Response:
+        """
+        This method implements the OAuth dynamic client registration endpoint. It responds to all requests
+        with a fixed client identifier.
+        """
+        body = await request.json()
+        redirect_uris = body.get("redirect_uris", [])
+        return starlette.responses.JSONResponse(
+            content={
+                "client_id": self._oauth_client,
+                "redirect_uris": redirect_uris,
+            },
+        )

server.py

@@ -1,11 +1,16 @@
-from mcp.server.fastmcp import FastMCP
 import json
 import os
+
+import fastmcp
+import fastmcp.server.dependencies
 import requests
+import uvicorn
+
+import oauth
 
 from service_client import InventoryClient
 
-mcp = FastMCP("AssistedService", host="0.0.0.0")
+mcp = fastmcp.FastMCP("AssistedService")
 
 def get_offline_token() -> str:
     """Retrieve the offline token from environment variables or request headers.
@@ -26,7 +31,8 @@ def get_offline_token() -> str:
     if token:
         return token
 
-    token = mcp.get_context().request_context.request.headers.get("OCM-Offline-Token")
+    headers = fastmcp.server.dependencies.get_http_headers()
+    token = headers.get("ocm-offline-token")
     if token:
         return token
 
@@ -46,13 +52,12 @@ def get_access_token() -> str:
         RuntimeError: If it isn't possible to obtain or generate the access token.
     """
     # First try to get the token from the authorization header:
-    request = mcp.get_context().request_context.request
-    if request is not None:
-        header = request.headers.get("Authorization")
-        if header is not None:
-            parts = header.split()
-            if len(parts) == 2 and parts[0].lower() == "bearer":
-                return parts[1]
+    headers = fastmcp.server.dependencies.get_http_headers()
+    header = headers.get("authorization")
+    if header is not None:
+        parts = header.split()
+        if len(parts) == 2 and parts[0].lower() == "bearer":
+            return parts[1]
 
     # Now try to get the offline token, and generate a new access token from it:
     params = {
@@ -300,4 +305,30 @@ def set_host_role(host_id: str, infraenv_id: str, role: str) -> str:
     return InventoryClient(get_access_token()).update_host(host_id, infraenv_id, host_role=role).to_str()
 
 if __name__ == "__main__":
-    mcp.run(transport="sse")
+    # We create a Starlette application so that we can add middleware:
+    app = mcp.http_app(transport="sse")
+
+    # Add the OAuth middleware if enabled:
+    oauth_enabled = os.getenv("OAUTH_ENABLED", "false").lower() == "true"
+    if oauth_enabled:
+        self_url = os.getenv(
+            "SELF_URL",
+            "http://localhost:8000",
+        )
+        oauth_url = os.getenv(
+            "OAUTH_URL",
+            "https://sso.redhat.com/auth/realms/redhat-external",
+        )
+        oauth_client = os.getenv(
+            "OAUTH_CLIENT",
+            "cloud-services",
+        )
+        app.add_middleware(
+            oauth.Middleware,
+            self_url=self_url,
+            oauth_url=oauth_url,
+            oauth_client=oauth_client,
+        )
+
+    # Start the application
+    uvicorn.run(app, host="0.0.0.0", port=8000)